1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use tipb::{Expr, ExprType, FieldType};

use crate::impl_bit_op::*;
use crate::impl_max_min::*;
use crate::impl_variance::*;
use crate::AggrFunction;
use tidb_query_common::Result;
use tidb_query_datatype::expr::EvalContext;
use tidb_query_expr::{RpnExpression, RpnExpressionBuilder};

/// Parse a specific aggregate function definition from protobuf.
///
/// All aggregate function implementations should include an impl for this trait as well as
/// add a match arm in `map_pb_sig_to_aggr_func_parser` so that the aggregate function can be
/// actually utilized.
pub trait AggrDefinitionParser {
    /// Checks whether the inner expression of the aggregate function definition is supported.
    /// It is ensured that `aggr_def.tp` maps the current parser instance.
    fn check_supported(&self, aggr_def: &Expr) -> Result<()>;

    /// Parses and transforms the aggregate function definition.
    ///
    /// The schema of this aggregate function will be appended in `out_schema` and the final
    /// RPN expression (maybe wrapped by some casting according to types) will be appended in
    /// `out_exp`.
    ///
    /// The parser may choose particular aggregate function implementation based on the data
    /// type, so `schema` is also needed in case of data type depending on the column.
    ///
    /// # Panic
    ///
    /// May panic if the aggregate function definition is not supported by this parser.
    fn parse(
        &self,
        mut aggr_def: Expr,
        ctx: &mut EvalContext,
        src_schema: &[FieldType],
        out_schema: &mut Vec<FieldType>,
        out_exp: &mut Vec<RpnExpression>,
    ) -> Result<Box<dyn AggrFunction>> {
        // Rewrite expression to insert CAST() if needed.
        let child = aggr_def.take_children().into_iter().next().unwrap();
        let exp = RpnExpressionBuilder::build_from_expr_tree(child, ctx, src_schema.len())?;

        Self::parse_rpn(&self, aggr_def, exp, ctx, src_schema, out_schema, out_exp)
    }

    #[inline]
    fn parse_rpn(
        &self,
        _root_expr: Expr,
        _exp: RpnExpression,
        _ctx: &mut EvalContext,
        _src_schema: &[FieldType],
        _out_schema: &mut Vec<FieldType>,
        _out_exp: &mut Vec<RpnExpression>,
    ) -> Result<Box<dyn AggrFunction>> {
        unimplemented!(
            "This struct neither implemented parse nor parse_rpn, which is not expected."
        )
    }
}

#[inline]
fn map_pb_sig_to_aggr_func_parser(value: ExprType) -> Result<Box<dyn AggrDefinitionParser>> {
    match value {
        ExprType::Count => Ok(Box::new(super::impl_count::AggrFnDefinitionParserCount)),
        ExprType::Sum => Ok(Box::new(super::impl_sum::AggrFnDefinitionParserSum)),
        ExprType::Avg => Ok(Box::new(super::impl_avg::AggrFnDefinitionParserAvg)),
        ExprType::First => Ok(Box::new(super::impl_first::AggrFnDefinitionParserFirst)),
        ExprType::AggBitAnd => Ok(Box::new(AggrFnDefinitionParserBitOp::<BitAnd>::new())),
        ExprType::AggBitOr => Ok(Box::new(AggrFnDefinitionParserBitOp::<BitOr>::new())),
        ExprType::AggBitXor => Ok(Box::new(AggrFnDefinitionParserBitOp::<BitXor>::new())),
        ExprType::Max => Ok(Box::new(AggrFnDefinitionParserExtremum::<Max>::new())),
        ExprType::Min => Ok(Box::new(AggrFnDefinitionParserExtremum::<Min>::new())),
        ExprType::Variance | ExprType::VarPop => {
            Ok(Box::new(AggrFnDefinitionParserVariance::<Population>::new()))
        }
        ExprType::VarSamp => Ok(Box::new(AggrFnDefinitionParserVariance::<Sample>::new())),
        v => Err(other_err!(
            "Aggregation function meet blacklist aggr function {:?}",
            v
        )),
    }
}

/// Parse all aggregate function definition from protobuf.
pub struct AllAggrDefinitionParser;

impl AggrDefinitionParser for AllAggrDefinitionParser {
    /// Checks whether the aggregate function definition is supported.
    #[inline]
    fn check_supported(&self, aggr_def: &Expr) -> Result<()> {
        let parser = map_pb_sig_to_aggr_func_parser(aggr_def.get_tp())?;
        parser.check_supported(aggr_def).map_err(|e| {
            other_err!(
                "Aggregation function meet blacklist expr type {:?}: {}",
                aggr_def.get_tp(),
                e
            )
        })
    }

    /// Parses and transforms the aggregate function definition to generate corresponding
    /// `AggrFunction` instance.
    ///
    /// # Panic
    ///
    /// May panic if the aggregate function definition is not supported.
    #[inline]
    fn parse(
        &self,
        aggr_def: Expr,
        ctx: &mut EvalContext,
        src_schema: &[FieldType],
        out_schema: &mut Vec<FieldType>,
        out_exp: &mut Vec<RpnExpression>,
    ) -> Result<Box<dyn AggrFunction>> {
        let parser = map_pb_sig_to_aggr_func_parser(aggr_def.get_tp()).unwrap();
        parser.parse(aggr_def, ctx, src_schema, out_schema, out_exp)
    }
}