分段批量插入性能优化

TL;DR

在我们的项目中，我们会针对一个文档做分段，分成多个segment之后，需要把他们都插入到数据库中。在最开始的代码中，我们是这么做的：

@Override
@Transactional
@DistributeLock(scene = "document-split", keyExpression = "#document.docId", waitTime = 0)
public int split(KnowledgeDocument document) {
    // 1. 查询文档


    // 2. 从MinIO下载文件内容


    // 3. 使用 MarkdownHeaderParentTextSplitter 进行切分


    // 4. 转换为 KnowledgeSegment 并保存


    // 5. 批量保存片段
    boolean saveResult = knowledgeSegmentService.saveBatch(knowledgeSegments);
    Assert.isTrue(saveResult, "保存知识片段失败");

    //...省略部分代码
}

也就是说，这里其实是调用的com.baomidou.mybatisplus.extension.service.IService#saveBatch(java.util.Collection)方法。但是这个方法是有性能问题的。他底层调用的是executeBatch方法，实现如下：

public static <E> boolean executeBatch(SqlSessionFactory sqlSessionFactory, Log log, Collection<E> list, int batchSize, BiConsumer<SqlSession, E> consumer) {
    Assert.isFalse(batchSize < 1, "batchSize must not be less than one");
    return !CollectionUtils.isEmpty(list) && executeBatch(sqlSessionFactory, log, sqlSession -> {
        int size = list.size();
        int idxLimit = Math.min(batchSize, size);
        int i = 1;
        for (E element : list) {
            consumer.accept(sqlSession, element);
            if (i == idxLimit) {
                sqlSession.flushStatements();
                idxLimit = Math.min(idxLimit + batchSize, size);
            }
            i++;
        }
    });
}

看到这里面的for循环了么？也就是说，所谓的saveBatch，默认实现看似是批量操作，但其底层逻辑实际上是在一个循环中逐条调用 save() 方法。在这种模式下，每调用一次 save()，框架就会立即向数据库发送一条独立的 INSERT 语句。一旦遇到一个文档的分段有很多的时候，性能极其差。而且，事务是在所有的insert都执行完之后，一次性提交的，也就是说，insert过程中你在数据库中是看不到数据的，并且如果其中有一个sql一旦失败了，则整个事务会回滚。那么解决思路是，在 JDBC URL 中添加 rewriteBatchedStatements=true

spring:
  datasource:
    url: jdbc:mysql://localhost:3306/konw-engine?useUnicode=true&characterEncoding=UTF-8&rewriteBatchedStatements=true

该参数会指示 MySQL 驱动将多条独立的 INSERT INTO table VALUES (...) 语句合并成一条高效的 INSERT INTO table VALUES (...), (...), ... 语句，降低数据库的解析和执行开销。另外，saveBatch 方法允许指定 batchSize 参数（默认为 1000）。根据实际数据量和服务器内存情况，可以调整此值以获得最佳性能。测试用的md文档：