TiKV 源码解析系列文章（十六）TiKV Coprocessor Executor 源码解析

select  *  from students where age >  21  limit  2

select * from t where a > 2 limit 2

pub trait Executor: Send {
    fn next(&mut self) -> Result<Option<Row>>;
    // ...
}
当以上的请求被解析之后，我们会在 ExecutorRunner 里边不断的调用最上层算子的 next() 方法， 直到其无法再返回行。
pub fn handle_request(&mut self) -> Result<SelectResponse> {
    loop {
        match self.executor.next()? {
            Some(row) => {
                // Do some aggregation.
            },
            None => {
                // ...
                return result;
            }
        }
    }
}

pub trait BatchExecutor: Send {

    /// 获取表的 `schema`
    fn schema(&self) -> &[FieldType];

    // 向下层算子要回一张表
    fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult;

    // ...
}

pub struct BatchExecuteResult {
    // 本轮循环 `TableScan` 扫上来的数据
    pub physical_columns: LazyBatchColumnVec,

    /// 记录 `physical_columns` 中有效的行的下标
    pub logical_rows: Vec<usize>,

    // ...

    // 表示下层算子是否已经没有数据可以返回
    pub is_drained: Result<bool>,
}

pub struct BatchTableScanExecutor<S: Storage>(ScanExecutor<S, TableScanExecutorImpl>);

pub struct ScanExecutor<S: Storage, I: ScanExecutorImpl> {
    /// 具体的扫表/扫索引实现。
    imp: I,

    /// 给定一个 `KeyRange`，扫上一行或者多行。
    scanner: RangesScanner<S>,

    // 标记是否已经扫完了所有的行。
    is_ended: bool,
}

    fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {

        // 创建一个列数组
        let mut logical_columns = self.imp.build_column_vec(scan_rows);
        
        // 扫上 `scan_rows` 行， 然后按列填充到创建好的列数组中。
        let is_drained = self.fill_column_vec(scan_rows, &mut logical_columns);

        // 创建一个 `logical_rows`, 表示当前表中所有行有效。后边可能根据 `Selection` 的结果修改这个 `logical_rows`。
        let logical_rows = (0..logical_columns.rows_len()).collect();

        // 判断是否扫完传入的 `KeyRange`
        match &is_drained {
            // Note: `self.is_ended` is only used for assertion purpose.
            Err(_) | Ok(true) => self.is_ended = true,
            Ok(false) => {}
        };

        // 返回 `BatchExecuteResult`
        BatchExecuteResult {
            // ...
        }
    }

    fn fill_column_vec(
        &mut self,
        scan_rows: usize,
        columns: &mut LazyBatchColumnVec,
    ) -> Result<bool> {
        assert!(scan_rows > 0);

        for _ in 0..scan_rows {
            let some_row = self.scanner.next()?;
            if let Some((key, value)) = some_row {
                // 将扫上来的一行放入 `columns` 中
                self.imp.process_kv_pair(&key, &value, columns)?;
            } else {
                // 没有 `KeyRange` 可供扫描，已经完成扫表。
                return Ok(true);
            }
        }

        // 表示下层数据还没有扫完。
        Ok(false)
    }

pub struct BatchSelectionExecutor<Src: BatchExecutor> {
    // ...
    
    // 数据源
    src: Src,

    // 条件表达式
    conditions: Vec<RpnExpression>,
}

    #[inline]
    fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {
        // 从下层算子那会一块数据开始过滤
        let mut src_result = self.src.next_batch(scan_rows);

        // 根据表达式的值，过滤出对应的行。
        if let Err(e) = self.handle_src_result(&mut src_result) {
            src_result.is_drained = src_result.is_drained.and(Err(e));
            src_result.logical_rows.clear();
        } else {
            // ... 
        }

        src_result

    fn handle_src_result(&mut self, src_result: &mut BatchExecuteResult) -> Result<()> {
        let mut src_logical_rows_copy = Vec::with_capacity(src_result.logical_rows.len());
        let mut condition_index = 0;
        while condition_index < self.conditions.len() && !src_result.logical_rows.is_empty() {
            // 拷贝一份下层算子的 `logical_rows`，用做计算表达式。
            src_logical_rows_copy.clear();
            src_logical_rows_copy.extend_from_slice(&src_result.logical_rows);

            // 计算表达式的值，然后根据表达式的值去更新下层算子的 `logical_rows`。
            match self.conditions[condition_index].eval(
                &mut self.context,
                self.src.schema(),
                &mut src_result.physical_columns,
                &src_logical_rows_copy,
                // 表达式产生的结果如果是一列的话, 这里表示表达式应该输出的行数
                src_logical_rows_copy.len(),
            )? {
                RpnStackNode::Scalar { value, .. } => {
                    // 如果表达式是一个标量，根据转换成 `bool` 的值确定是否保留该列。
                    update_logical_rows_by_scalar_value(
                        &mut src_result.logical_rows,
                        &mut self.context,
                        value,
                    )?;
                }
                RpnStackNode::Vector { value, .. } => {
                    // 根据每行的结果，确定是否保留那行。
                    update_logical_rows_by_vector_value(
                    &mut src_result.logical_rows,
                    &mut self.context,
                    eval_result,
                    eval_result_logical_rows,
                    )?;
                }
            }

            condition_index += 1;
        }

        Ok(())
    }
}

select count(*) from t group by a

pub trait AggrFunction: std::fmt::Debug + Send + 'static {
    /// The display name of the function.
    fn name(&self) -> &'static str;

    /// Creates a new state instance. Different states aggregate independently.
    fn create_state(&self) -> Box<dyn AggrFunctionState>;
}

// NOTE: AggrFunctionState 是 AggrFunctionStateUpdatePartial 的 super trait
pub trait AggrFunctionState:
    std::fmt::Debug
    + Send
    + 'static
    + AggrFunctionStateUpdatePartial<Int>
    + AggrFunctionStateUpdatePartial<Real>
    + AggrFunctionStateUpdatePartial<Decimal>
    + AggrFunctionStateUpdatePartial<Bytes>
    + AggrFunctionStateUpdatePartial<DateTime>
    + AggrFunctionStateUpdatePartial<Duration>
    + AggrFunctionStateUpdatePartial<Json>
{
    fn push_result(&self, ctx: &mut EvalContext, target: &mut [VectorValue]) -> Result<()>;
}
pub trait AggrFunctionStateUpdatePartial<T: Evaluable> {
    fn update(&mut self, ctx: &mut EvalContext, value: &Option<T>) -> Result<()>;

    fn update_repeat(
        &mut self,
        ctx: &mut EvalContext,
        value: &Option<T>,
        repeat_times: usize,
    ) -> Result<()>;

    fn update_vector(
        &mut self,
        ctx: &mut EvalContext,
        physical_values: &[Option<T>],
        logical_rows: &[usize],
    ) -> Result<()>;
}

/// The COUNT aggregate function.
#[derive(Debug, AggrFunction)]
#[aggr_function(state = AggrFnStateCount::new())]
pub struct AggrFnCount;

/// The state of the COUNT aggregate function.
#[derive(Debug)]
pub struct AggrFnStateCount {
    count: usize,
}

impl AggrFnStateCount {
    pub fn new() -> Self {
        Self { count: 0 }
    }
}

impl AggrFunctionStateUpdatePartial for AggrFnStateCount { /* .. */ }
impl AggrFunctionState for AggrFnStateCount { /* .. */ }

pub struct BatchFastHashAggregationExecutor<Src: BatchExecutor>(
    AggregationExecutor<Src, FastHashAggregationImpl>,
);

pub struct AggregationExecutor<Src: BatchExecutor, I: AggregationExecutorImpl<Src>> {
    imp: I,
    is_ended: bool,
    entities: Entities<Src>,
}

pub trait AggregationExecutorImpl<Src: BatchExecutor>: Send {
    // 根据 `group by` columns 和 聚合函数初始化 `entities` 中的 `schema`
    fn prepare_entities(&mut self, entities: &mut Entities<Src>);

    // 根据下层算子扫上来的数据做聚合和分组
    fn process_batch_input(
        &mut self,
        entities: &mut Entities<Src>,
        input_physical_columns: LazyBatchColumnVec,
        input_logical_rows: &[usize],
    ) -> Result<()>;

    // 将每个聚合函数的状态更新到列数组中，即写入聚合结果
    // 这里返回的是 `group by` column，在分布式场景如果不把 `group by` column 返回，`TiDB` 没有办法根据分组做二次聚合。
    fn iterate_available_groups(
        &mut self,
        entities: &mut Entities<Src>,
        src_is_drained: bool,
        iteratee: impl FnMut(&mut Entities<Src>, &[Box<dyn AggrFunctionState>]) -> Result<()>,
    ) -> Result<Vec<LazyBatchColumn>>;
}

pub struct Entities<Src: BatchExecutor> {
    pub src: Src,
    
    // ...

    // 聚合后产生的 `schmea`， 包含 `group_by` columns
    pub schema: Vec<FieldType>,

    /// 聚合函数的集合
    pub each_aggr_fn: Vec<Box<dyn AggrFunction>>,

    /// 每个聚合函数输出的列大小，`COUNT` 是 1，`AVG` 是 2
    pub each_aggr_cardinality: Vec<usize>,

    /// 聚合函数里边的表达式
    pub each_aggr_exprs: Vec<RpnExpression>,

    // 每个聚合表达式输出的类型的集合
    pub all_result_column_types: Vec<EvalType>,
}

    fn handle_next_batch(&mut self) -> Result<(Option<LazyBatchColumnVec>, bool)> {
        // 从下层算子取回一个 `batch`
        let src_result = self
            .entities
            .src
            .next_batch(crate::batch:🏃:BATCH_MAX_SIZE);

        self.entities.context.warnings = src_result.warnings;

        let src_is_drained = src_result.is_drained?;

        // 如果下层返回的数据不为空，将根据每行的结果分组并聚合
        if !src_result.logical_rows.is_empty() {
            self.imp.process_batch_input(
                &mut self.entities,
                src_result.physical_columns,
                &src_result.logical_rows,
            )?;
        }

        // 在 `FastHashAggr` 中，只有下层算子没有办法再返回数据的时候，才能认为聚合已经完成，
        // 否则我们返回一个空数据给上层算子，等待下一次 `next_batch` 被调用。
        let result = if src_is_drained {
            Some(self.aggregate_partial_results(src_is_drained)?)
        } else {
            None
        };
        Ok((result, src_is_drained))
    }