Sharding-JDBC 源码之结果集归并

Sharding-JDBC 系列

1. 第一篇 Sharding-JDBC 源码之启动流程分析
2. 第二篇 Sharding-JDBC 源码之 SQL 解析
3. 第三篇 Sharding-JDBC 源码之 SQL 路由
4. 第四篇 Sharding-JDBC 源码之 SQL 改写
5. 第五篇 Sharding-JDBC 源码之 SQL 执行
6. 第六篇 Sharding-JDBC 源码之结果集归并（本文）

---

将从各个数据节点获取的多数据结果集，组合成为一个结果集并正确的返回至请求客户端，称为结果归并。由于从数据库中返回的结果集是逐条返回的，并不需要将所有的数据一次性加载至内存中，因此，在进行结果归并时，沿用数据库返回结果集的方式进行归并，能够极大减少内存的消耗，是归并方式的优先选择。

归并类型

流式归并

指每一次从结果集中获取到的数据，都能够通过逐条获取的方式返回正确的单条数据，它与数据库原生的返回结果集的方式最为契合。遍历、排序以及流式分组都属于流式归并的一种。

内存归并

需要将结果集的所有数据都遍历并存储在内存中，再通过统一的分组、排序以及聚合等计算之后，再将其封装成为逐条访问的数据结果集返回。

具体归并类型请参考官网：内核剖析-归并引擎

源码分析

将执行完成的 PreparedStatement 对象交给 resultSetHandler 处理

org.apache.ibatis.executor.statement.PreparedStatementHandler#query

    public <E> List<E> query(Statement statement, ResultHandler resultHandler) throws SQLException {
    PreparedStatement ps = (PreparedStatement) statement;
    // 执行完成
    ps.execute();
    // 处理结果集
    return resultSetHandler.handleResultSets(ps);
  }

执行完成后，将 PreparedStatement 对象 ps 传给结果处理器 resultSetHandler 的 handleResultSets 方法；

执行 DefaultResultSetHandler 中的 handleResultSets 方法

    public List<Object> handleResultSets(Statement stmt) throws SQLException {
    ErrorContext.instance().activity("handling results").object(mappedStatement.getId());

    final List<Object> multipleResults = new ArrayList<>();

    int resultSetCount = 0;
    // 1. 获取第一个结果集
    ResultSetWrapper rsw = getFirstResultSet(stmt);

    // ...... 暂时省略
  }
  
  private ResultSetWrapper getFirstResultSet(Statement stmt) throws SQLException {
    //  2. stmt 是 ShardingPreparedStatement 类型，所以执行该类中的方法
    ResultSet rs = stmt.getResultSet();
    while (rs == null) {
      // move forward to get the first resultset in case the driver
      // doesn't return the resultset as the first result (HSQLDB 2.1)
      if (stmt.getMoreResults()) {
        rs = stmt.getResultSet();
      } else {
        if (stmt.getUpdateCount() == -1) {
          // no more results. Must be no resultset
          break;
        }
      }
    }
    return rs != null ? new ResultSetWrapper(rs, configuration) : null;
  }
  

1 .获取第一个结果集，这里的 stmt 对象是 ShardingPreparedStatement 类型；
2. getFirstResultSet 方法中，stmt 调用 getResultSet 方法获取结果集。

执行 ShardingPreparedStatement 中的 getResultSet 方法

    public ResultSet getResultSet() throws SQLException {
        if (null != currentResultSet) {
            return currentResultSet;
        }
        
        // 1. 如果只查询一张库表，则不用进行结果归并
        if (1 == routedStatements.size() && routeResult.getSqlStatement() instanceof DQLStatement) {
            currentResultSet = routedStatements.iterator().next().getResultSet();
            return currentResultSet;
        }
        List<ResultSet> resultSets = new ArrayList<>(routedStatements.size());
        List<QueryResult> queryResults = new ArrayList<>(routedStatements.size());
        
        // 2.将多个路由执行的结果集封装在 queryResults 列表中
        for (PreparedStatement each : routedStatements) {
            ResultSet resultSet = each.getResultSet();
            resultSets.add(resultSet);
            queryResults.add(new JDBCQueryResult(resultSet));
        }
        if (routeResult.getSqlStatement() instanceof SelectStatement || routeResult.getSqlStatement() instanceof DALStatement) {
            // 3. 根据 SQLStatement 类型创建合并引擎
            MergeEngine mergeEngine = MergeEngineFactory.newInstance(connection.getShardingContext().getShardingRule(), queryResults, routeResult.getSqlStatement());
            // 4. 执行结果归并
            currentResultSet = new ShardingResultSet(resultSets, mergeEngine.merge(), this);
        }
        return currentResultSet;
    }
    

1. routedStatements.size() 等于 1 并且是查询 SQL 表示路由后的 SQL 只有一个，不涉及到多表结果合并，所以直接从返回第一个结果集；
2. 将多个结果集装入 queryResults 列表中，为初始化合并引擎做准备；
3. 根据 SQLStatement 类型创建合并引擎，本例是查询语句，所以创建了 DQLMergeEngine 类型的引擎，并初始化引擎中属性；

执行归并引擎 DQLMergeEngine 中 merge 方法

    public MergedResult merge() throws SQLException {
        // 1. columnLabelIndexMap 存储列名称和对应位置下标的映射
        selectStatement.setIndexForItems(columnLabelIndexMap);
        return decorate(build());
    }
    
    public void setIndexForItems(final Map<String, Integer> columnLabelIndexMap) {
        // 标记聚合函数使用的列在查询字段中的下标
        setIndexForAggregationItem(columnLabelIndexMap);
        // 标记排序字段在查询字段中的下标
        setIndexForOrderItem(columnLabelIndexMap, orderByItems);
        // 标记分组字段在查询字段中的下标
        setIndexForOrderItem(columnLabelIndexMap, groupByItems);
    }
    
    private void setIndexForAggregationItem(final Map<String, Integer> columnLabelIndexMap) {
        for (AggregationSelectItem each : getAggregationSelectItems()) {
            Preconditions.checkState(columnLabelIndexMap.containsKey(each.getColumnLabel()), String.format("Can't find index: %s, please add alias for aggregate selections", each));
            // 2. 设置聚合函数 item 的下标
            each.setIndex(columnLabelIndexMap.get(each.getColumnLabel()));
            for (AggregationSelectItem derived : each.getDerivedAggregationSelectItems()) {
                Preconditions.checkState(columnLabelIndexMap.containsKey(derived.getColumnLabel()), String.format("Can't find index: %s", derived));
                derived.setIndex(columnLabelIndexMap.get(derived.getColumnLabel()));
            }
        }
    }
    
    private void setIndexForOrderItem(final Map<String, Integer> columnLabelIndexMap, final List<OrderItem> orderItems) {
        for (OrderItem each : orderItems) {
            if (-1 != each.getIndex()) {
                continue;
            }
            Preconditions.checkState(columnLabelIndexMap.containsKey(each.getColumnLabel()), String.format("Can't find index: %s", each));
            if (columnLabelIndexMap.containsKey(each.getColumnLabel())) {
                // 3. 设置排序或分组 item 的下标
                each.setIndex(columnLabelIndexMap.get(each.getColumnLabel()));
            }
        }
    }

1. columnLabelIndexMap 存储 select 关键字后的列及下标位置，key 存储列名（别名优先），value 为下标位置（从 1 开始）；
2. 设置聚合函数 item 的下标；
3. 设置排序或分组 item 的下标；
4. 获取聚合函数或分组、排序列所在下标的目的是为了下一步采取流式还是内存归并做准备。

构建 MergedResult

    public MergedResult merge() throws SQLException {
        selectStatement.setIndexForItems(columnLabelIndexMap);
        // build 构建合并结果
        return decorate(build());
    }
    
    private MergedResult build() throws SQLException {
        if (!selectStatement.getGroupByItems().isEmpty() || !selectStatement.getAggregationSelectItems().isEmpty()) {
            // 1. 判断分组和排序内容是否完全一样，相同则流式归并
            if (selectStatement.isSameGroupByAndOrderByItems()) {
                return new GroupByStreamMergedResult(columnLabelIndexMap, queryResults, selectStatement);
            } else {
                // 否则内存归并
                return new GroupByMemoryMergedResult(columnLabelIndexMap, queryResults, selectStatement);
            }
        }
        
        // 2. 只有 order by 排序，则流式排序
        if (!selectStatement.getOrderByItems().isEmpty()) {
            return new OrderByStreamMergedResult(queryResults, selectStatement.getOrderByItems());
        }
        
        // 2. 线性流式合并
        return new IteratorStreamMergedResult(queryResults);
    }

1. 判断分组 item 和 排序 item 内容是否相同，如果相同则进行流式归并，否则进行内存归并，这里的内容是指分组和排序 item 中的属性必须相等，包括：name（列名称）、列别名（alias）、（排序顺序）orderDirection、（默认排序 ASC）nullOrderDirection、（下标位置）index 等；
   如下图
2. 只有排序 item 则进行流式排序归并；
3. 如果不涉及分组、排序、聚合等操作，则只是简单的将结果集按单链表拼接在一起即可，采用流式归并。

在归并结果的构造函数中处理数据

不同的归并对象处理归并结果集的逻辑不同，他们都会在生成对象的同时在构造函数中对归并结果进行一些预处理，
比如：

分组和排序 item 相同时的归并对象 GroupByStreamMergedResult

    private final List<Object> currentRow;
    
    private List<?> currentGroupByValues;
    
    public GroupByStreamMergedResult(
            final Map<String, Integer> labelAndIndexMap, final List<QueryResult> queryResults, final SelectStatement selectStatement) throws SQLException {
        super(queryResults, selectStatement.getOrderByItems());
        this.labelAndIndexMap = labelAndIndexMap;
        this.selectStatement = selectStatement;
        currentRow = new ArrayList<>(labelAndIndexMap.size());
        currentGroupByValues = getOrderByValuesQueue().isEmpty() ? Collections.emptyList() : new GroupByValue(getCurrentQueryResult(), selectStatement.getGroupByItems()).getGroupValues();
    }

初始化成员属性 currentRow （当前行对象），并获取当前分组值；

分组归并对象 GroupByMemoryMergedResult

    private final Iterator<MemoryQueryResultRow> memoryResultSetRows;
    
    public GroupByMemoryMergedResult(
            final Map<String, Integer> labelAndIndexMap, final List<QueryResult> queryResults, final SelectStatement selectStatement) throws SQLException {
        super(labelAndIndexMap);
        this.selectStatement = selectStatement;
        // 初始化迭代器
        memoryResultSetRows = init(queryResults);
    }
    
    private Iterator<MemoryQueryResultRow> init(final List<QueryResult> queryResults) throws SQLException {
        // 1. 获取分组 key 对象的行对象
        Map<GroupByValue, MemoryQueryResultRow> dataMap = new HashMap<>(1024);
        // 2. 分组 key 对应的行对象
        Map<GroupByValue, Map<AggregationSelectItem, AggregationUnit>> aggregationMap = new HashMap<>(1024);
        for (QueryResult each : queryResults) {
            while (each.next()) {
                GroupByValue groupByValue = new GroupByValue(each, selectStatement.getGroupByItems());
                initForFirstGroupByValue(each, groupByValue, dataMap, aggregationMap);
                aggregate(each, groupByValue, aggregationMap);
            }
        }
        setAggregationValueToMemoryRow(dataMap, aggregationMap);
        List<MemoryQueryResultRow> result = getMemoryResultSetRows(dataMap);
        if (!result.isEmpty()) {
            setCurrentResultSetRow(result.get(0));
        }
        return result.iterator();
    }

1. GroupByValue 是执行完SQL后按照某列分组后的具体值，比如按照 goods_id 分组，一共两个 goods_id 分别为：（1001,、1002），则一个GroupByValue 对象就是 1001，另一个是 1002，MemoryQueryResultRow 对象存着对应 goods_id 具体行数据；
2. 聚合函数类似，GroupByValue 存的是具体分组后的 goods_id 值，value 中的 AggregationUnit 为聚合函数的计算结果；
3. 对归并结果进行预处理后，还需要做最后一步处理，便是根据分页属性对数据进行分页或截断处理。

分页装饰

对前一步归并的结果进行处理，如果有 limit 关键字，则要进行相应的分页或截断

    public MergedResult merge() throws SQLException {
        selectStatement.setIndexForItems(columnLabelIndexMap);
        // 装饰归并结果
        return decorate(build());
    }
    
    private MergedResult decorate(final MergedResult mergedResult) throws SQLException {
        // 获取分页信息
        Limit limit = selectStatement.getLimit();
        if (null == limit) {
            return mergedResult;
        }
        if (DatabaseType.MySQL == limit.getDatabaseType() || DatabaseType.PostgreSQL == limit.getDatabaseType() || DatabaseType.H2 == limit.getDatabaseType()) {
            return new LimitDecoratorMergedResult(mergedResult, selectStatement.getLimit());
        }
        if (DatabaseType.Oracle == limit.getDatabaseType()) {
            return new RowNumberDecoratorMergedResult(mergedResult, selectStatement.getLimit());
        }
        if (DatabaseType.SQLServer == limit.getDatabaseType()) {
            return new TopAndRowNumberDecoratorMergedResult(mergedResult, selectStatement.getLimit());
        }
        return mergedResult;
    }

根据数据库类型封装为装饰分页合并结果，并在构造函数中处理分页或截断数据。例如：
在 Mysql 类型的装饰对象 LimitDecoratorMergedResult 中

    private final Limit limit;
    
    private final boolean skipAll;
    public LimitDecoratorMergedResult(final MergedResult mergedResult, final Limit limit) throws SQLException {
        super(mergedResult);
        this.limit = limit;
        // 跳过
        skipAll = skipOffset();
    }
    
    private boolean skipOffset() throws SQLException {
        // 小于偏移量继续取值
        for (int i = 0; i < limit.getOffsetValue(); i++) {
            if (!getMergedResult().next()) {
                return true;
            }
        }
        rowNumber = 0;
        return false;
    }

根据 limit 对象中的 offset 进行取值，超过偏移量或无数据则返回 false。

到这里，数据归并就已经全部完成了，后续就是在 mybatis 中继续获取归并后的结果集并根据 ResultMap 封装数据返回给客户端了，回过头看 handleResultSets 方法后半部分

org.apache.ibatis.executor.resultset.DefaultResultSetHandler#handleResultSets

    public List<Object> handleResultSets(Statement stmt) throws SQLException {
    //...... 省略
    String[] resultSets = mappedStatement.getResultSets();
    if (resultSets != null) {
      while (rsw != null && resultSetCount < resultSets.length) {
        ResultMapping parentMapping = nextResultMaps.get(resultSets[resultSetCount]);
        if (parentMapping != null) {
          String nestedResultMapId = parentMapping.getNestedResultMapId();
          ResultMap resultMap = configuration.getResultMap(nestedResultMapId);
          // 获取归并结果集入口
          handleResultSet(rsw, resultMap, null, parentMapping);
        }
        rsw = getNextResultSet(stmt);
        cleanUpAfterHandlingResultSet();
        resultSetCount++;
      }
    }

    return collapseSingleResultList(multipleResults);
  }
  
  private void handleResultSet(ResultSetWrapper rsw, ResultMap resultMap, List<Object> multipleResults, ResultMapping parentMapping) throws SQLException {
    //......省略
    // 继续调用此方法
    handleRowValues(rsw, resultMap, null, RowBounds.DEFAULT, parentMapping);
  }
  
  private void handleRowValuesForNestedResultMap(ResultSetWrapper rsw, ResultMap resultMap, ResultHandler<?> resultHandler, RowBounds rowBounds, ResultMapping parentMapping) throws SQLException {
    final DefaultResultContext<Object> resultContext = new DefaultResultContext<>();
    // 注意，这里的 ResuleSet 是 ShardingResultSet
    ResultSet resultSet = rsw.getResultSet();
    skipRows(resultSet, rowBounds);
    Object rowValue = previousRowValue;
    // resultSet.next() 为 ShardingResultSet 中的方法
    while (shouldProcessMoreRows(resultContext, rowBounds) && !resultSet.isClosed() && resultSet.next()) {
    //......省略 mybatis中代码  
  }

  // ShardingResultSet 重写后的方法
  public boolean next() throws SQLException {
        return mergeResultSet.next();
    }

1. 因为 ResultSet 是 ShardingResultSet 类型，所以调用重写后的方法，因此获取到的就是归并后的结果；
2. 剩下来的结果集映射就是 Mybatis 该处理的事了。至此，Sharding-JDBC 在一个 SQL 执行中的使命已完成。

总结

1. 简单介绍了流式归并和内存归并的区别，详情可参照官网介绍；
2. 如果路由后的 SQL 只有一个并且是查询语句，不会进行结果归并，直接返回结果；多个路由 SQL ，会将查询结果封装到 queryResults 列表中，并初始化归并引擎 MergeEngine；
3. 归并引擎中的 merge 会计算出分组、排序或聚合函数列所在查询列中的下标位置，为后面是使用流式还是内存归并做准备；
4. 在创建归并结果对象的时候，会在构造函数中对归并结果进行一些预处理，如分组、排序数据合并，LIMIT 分页限制等；
5. Mybatis 在遍历结果集的时候，由于持有的是 ShardingResultSet 对象，所以会去拿到 mergeResultSet 中数据，并按照正常 Mybatis 的流程继续将数据封装完成后返回给调用的客户端。

最最后，贴一张整个 SQL 执行时序图：