refine appendAggregation in analyzer

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h

@@ -53,6 +53,7 @@ class DAGExpressionAnalyzer : private boost::noncopyable
         const tipb::TopN & topN);
 
     /// <aggregation_keys, collators, aggregate_descriptions, before_agg>
+    /// May change the source columns.
     std::tuple<Names, TiDB::TiDBCollators, AggregateDescriptions, ExpressionActionsPtr> appendAggregation(
         ExpressionActionsChain & chain,
         const tipb::Aggregation & agg,
@@ -62,6 +63,8 @@ class DAGExpressionAnalyzer : private boost::noncopyable
         ExpressionActionsChain & chain,
         const std::vector<NameAndTypePair> & columns) const;
 
+    ExpressionActionsChain::Step & initAndGetLastStep(ExpressionActionsChain & chain) const;
+
     void appendJoin(
         ExpressionActionsChain & chain,
         SubqueryForSet & join_query,
@@ -119,7 +122,7 @@ class DAGExpressionAnalyzer : private boost::noncopyable
         const ::google::protobuf::RepeatedPtrField<tipb::ByItem> & order_by);
 
     void appendCastAfterAgg(
-        ExpressionActionsChain & chain,
+        const ExpressionActionsPtr & actions,
         const tipb::Aggregation & agg);
 
     String buildTupleFunctionForGroupConcat(
@@ -131,20 +134,43 @@ class DAGExpressionAnalyzer : private boost::noncopyable
 
     void buildGroupConcat(
         const tipb::Expr & expr,
-        ExpressionActionsChain::Step & step,
+        const ExpressionActionsPtr & actions,
         const String & agg_func_name,
         AggregateDescriptions & aggregate_descriptions,
         NamesAndTypes & aggregated_columns,
         bool result_is_nullable);
 
     void buildCommonAggFunc(
         const tipb::Expr & expr,
-        ExpressionActionsChain::Step & step,
+        const ExpressionActionsPtr & actions,
         const String & agg_func_name,
         AggregateDescriptions & aggregate_descriptions,
         NamesAndTypes & aggregated_columns,
         bool empty_input_as_null);
 
+    void buildAggFuncs(
+        const tipb::Aggregation & aggregation,
+        const ExpressionActionsPtr & actions,
+        AggregateDescriptions & aggregate_descriptions,
+        NamesAndTypes & aggregated_columns);
+
+    void buildAggGroupBy(
+        const google::protobuf::RepeatedPtrField<tipb::Expr> & group_by,
+        const ExpressionActionsPtr & actions,
+        AggregateDescriptions & aggregate_descriptions,
+        NamesAndTypes & aggregated_columns,
+        Names & aggregation_keys,
+        std::unordered_set<String> & agg_key_set,
+        bool group_by_collation_sensitive,
+        TiDB::TiDBCollators & collators);
+
+    void fillAggArgumentDetail(
+        const ExpressionActionsPtr & actions,
+        const tipb::Expr & arg,
+        Names & arg_names,
+        DataTypes & arg_types,
+        TiDB::TiDBCollators & arg_collators);
+
     void makeExplicitSet(
         const tipb::Expr & expr,
         const Block & sample_block,

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzerHelper.cpp

@@ -1,10 +1,12 @@
+#include <AggregateFunctions/AggregateFunctionFactory.h>
 #include <Columns/ColumnSet.h>
 #include <Common/FmtUtils.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeSet.h>
 #include <DataTypes/getLeastSupertype.h>
 #include <Flash/Coprocessor/DAGExpressionAnalyzer.h>
 #include <Flash/Coprocessor/DAGExpressionAnalyzerHelper.h>
+#include <Flash/Coprocessor/DAGUtils.h>
 #include <Functions/FunctionFactory.h>
 #include <Functions/FunctionsTiDBConversion.h>
 #include <Storages/Transaction/TypeMapping.h>
@@ -56,25 +58,6 @@ const std::unordered_map<String, String> DateSub::unit_to_func_name_map
         {"SECOND", "subtractSeconds"}};
 } // namespace
 
-String DAGExpressionAnalyzerHelper::genFuncString(
-    const String & func_name,
-    const Names & argument_names,
-    const TiDB::TiDBCollators & collators)
-{
-    assert(!collators.empty());
-    FmtBuffer buf;
-    buf.fmtAppend("{}({})_collator", func_name, fmt::join(argument_names.begin(), argument_names.end(), ", "));
-    for (const auto & collator : collators)
-    {
-        if (collator == nullptr)
-            buf.append("_0");
-        else
-            buf.fmtAppend("_{}", collator->getCollatorId());
-    }
-    buf.append(" ");
-    return buf.toString();
-}
-
 String DAGExpressionAnalyzerHelper::buildMultiIfFunction(
     DAGExpressionAnalyzer * analyzer,
     const tipb::Expr & expr,

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzerHelper.h

@@ -69,11 +69,6 @@ class DAGExpressionAnalyzerHelper
         const tipb::Expr & expr,
         const ExpressionActionsPtr & actions);
 
-    static String genFuncString(
-        const String & func_name,
-        const Names & argument_names,
-        const TiDB::TiDBCollators & collators);
-
     using FunctionBuilder = std::function<String(DAGExpressionAnalyzer *, const tipb::Expr &, const ExpressionActionsPtr &)>;
     using FunctionBuilderMap = std::unordered_map<String, FunctionBuilder>;
 

dbms/src/Flash/Coprocessor/DAGQueryBlockInterpreter.cpp

@@ -956,8 +956,7 @@ void DAGQueryBlockInterpreter::handleProjection(DAGPipeline & pipeline, const ti
         input_columns.emplace_back(p.name, p.type);
     DAGExpressionAnalyzer dag_analyzer(std::move(input_columns), context);
     ExpressionActionsChain chain;
-    dag_analyzer.initChain(chain, dag_analyzer.getCurrentInputColumns());
-    ExpressionActionsChain::Step & last_step = chain.steps.back();
+    auto & last_step = dag_analyzer.initAndGetLastStep(chain);
     std::vector<NameAndTypePair> output_columns;
     NamesWithAliases project_cols;
     UniqueNameGenerator unique_name_generator;

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -1155,6 +1155,25 @@ SortDescription getSortDescription(const std::vector<NameAndTypePair> & order_co
     return order_descr;
 }
 
+String genFuncString(
+    const String & func_name,
+    const Names & argument_names,
+    const TiDB::TiDBCollators & collators)
+{
+    assert(!collators.empty());
+    FmtBuffer buf;
+    buf.fmtAppend("{}({})_collator", func_name, fmt::join(argument_names.begin(), argument_names.end(), ", "));
+    for (const auto & collator : collators)
+    {
+        if (collator)
+            buf.fmtAppend("_{}", collator->getCollatorId());
+        else
+            buf.append("_0");
+    }
+    buf.append(" ");
+    return buf.toString();
+}
+
 TiDB::TiDBCollatorPtr getCollatorFromFieldType(const tipb::FieldType & field_type)
 {
     if (field_type.collate() < 0)

dbms/src/Flash/Coprocessor/DAGUtils.h

@@ -36,6 +36,10 @@ DataTypePtr inferDataType4Literal(const tipb::Expr & expr);
 SortDescription getSortDescription(
     const std::vector<NameAndTypePair> & order_columns,
     const google::protobuf::RepeatedPtrField<tipb::ByItem> & by_items);
+String genFuncString(
+    const String & func_name,
+    const Names & argument_names,
+    const TiDB::TiDBCollators & collators);
 
 extern const Int8 VAR_SIZE;