presto中count算子的实现

一. 引言

count是presto中支持的聚合算子的一种，本文走读presto中是如何实现select count(*) from xxx来探讨presto中是如何实现数据聚合操作的。了解presto中的聚合操作的实现对于了解presto在读文件方面的优化和了解presto的数据传递全貌方面甚至动态算子代码生成方面有较大的帮助意义。

本文以扫描一个parquent表的行数为例进行代码走读和探讨。

       presto中聚合操作的流程大体如下：

二. 数据源扫描代码走读

在presto中，扫描数据在ConnectorPageSource：：getNextPage中不断实时扫描，每次扫描都给上层返回当前已经扫描到的Page，直到将数据都扫描完成，给上层返回一个Page的条数为-1来实现的。聚合操作也遵循此流程。以Parquet为例， 其实现在ParquetPageSource::getNextPage中，其实现如下：

       

@Override
public Page getNextPage()
{
     // 上层在实时不断地拿着PageSource已经产生的Page，batchId记录着当前给上层返回了多少个Page
     batchId++;
     // batchSize记录着当前Page有多少条记录，batchSize是count计算的关键
     int batchSize = parquetReader.nextBatch();
     ....
        
     // 在count（*）的聚合计算中，列将被裁剪到0列，因此如下列计算不会被执行，因此count的计算非常快
     for (int fieldId = 0; fieldId < blocks.length; fieldId++) {
        .....   
      }
      // batchSize连随着Page返回去，在count的计算中，blocks为空
      return new Page(batchSize, blocks);
      ....
}

parquetReader.nextBatch的实现如下：

public int nextBatch()
{
    // 文件已经读完了，pagesize为-1表示为最后一个Page了
    if (nextRowInGroup >= currentGroupRowCount && !advanceToNextRowGroup()) {
        return -1;
    }
    
    // 下边的计算代表每次Page读取的条数要在上一次的基础上*2，也即batchId=1时读1条，batchId=2时读2条，batchId=3时读4条，batchId=4时读8条...直到最大的1024条或者文件结束
    batchSize = toIntExact(min(nextBatchSize, maxBatchSize));
    nextBatchSize = min(batchSize * BATCH_SIZE_GROWTH_FACTOR, MAX_VECTOR_LENGTH);
    batchSize = toIntExact(min(batchSize, currentGroupRowCount - nextRowInGroup));

    nextRowInGroup += batchSize;
    currentPosition += batchSize;
    
    // 遍历每列，从parquet文件读取batchSize条数的数据，塞到block的数组中。当前count计算的场景下，columnReaders的大小为0
    Arrays.stream(columnReaders)
            .forEach(reader -> reader.prepareNextRead(batchSize));
    return batchSize;
}

      上边就是数据源在tablescan处理中count算子的实现。

      在getNextPage已经返回了每个Page的数据大小，在聚合计算中，将Page的大小做聚合即可。

三. count聚合

      tablescan会实时通过addInput接口将page传递给AggregationOperator中，Aggregator中处理每一个Page的聚合计算代码如下：

       

public void processPage(Page page)
{
    if (step.isInputRaw()) {
        aggregation.addInput(page);
    }
    else {
        aggregation.addIntermediate(page.getBlock(intermediateChannel));
    }
}

       但是在presto中，count算子的计算是动态生成的class，也就是说上边的aggregation.evaluateIntermediate和aggregation.evaluateFinal 是无法在presto找到其静态的代码实现的。其实其实现代码在AccumulatorCompiler中动态生成的，其生成的函数分别是generateAddInput和generateAddIntermediateAsCombine：

private static void generateAddInput(...)
{
    ...
    BytecodeBlock block = generateInputForLoop(
     ....);

    body.append(block);
    body.ret();
}

private static BytecodeBlock generateInputForLoop(...)
{
	.....
    // getPositionCount获取到的就是上边tablescan的pageSize
    BytecodeBlock block = new BytecodeBlock()
            .append(page)
            .invokeVirtual(Page.class, "getPositionCount", int.class)
            .putVariable(rowsVariable)
            .initializeVariable(positionVariable);
      ....
      return block;
 }

      presto会调用每一个Page的getPositionCount获取到Page的条目，然后本地task做每个Page累加，本地完成后再回传给coordaintor做最后的汇总操作完成count计算。