Pipeline: Support pipeline agg

dbms/src/Flash/Pipeline/Pipeline.cpp

@@ -140,6 +140,7 @@ bool Pipeline::isSupported(const tipb::DAGRequest & dag_request)
             case tipb::ExecType::TypeSelection:
             case tipb::ExecType::TypeLimit:
             case tipb::ExecType::TypeTopN:
+            case tipb::ExecType::TypeAggregation:
             // Only support mock table_scan/exchange_sender/exchange_receiver in test mode now.
             case tipb::ExecType::TypeTableScan:
             case tipb::ExecType::TypeExchangeSender:

dbms/src/Flash/Planner/PlanType.h

@@ -36,8 +36,10 @@ struct PlanType
         TableScan = 11,
         MockTableScan = 12,
         Join = 13,
-        GetResult = 14,
-        Expand = 15,
+        PreAggregation = 14,
+        ConvergentAggregation = 15,
+        Expand = 16,
+        GetResult = 17
     };
     PlanTypeEnum enum_value;
 

dbms/src/Flash/Planner/Plans/PhysicalAggregation.cpp

@@ -25,6 +25,8 @@
 #include <Flash/Planner/FinalizeHelper.h>
 #include <Flash/Planner/PhysicalPlanHelper.h>
 #include <Flash/Planner/Plans/PhysicalAggregation.h>
+#include <Flash/Planner/Plans/PhysicalBuildAggregation.h>
+#include <Flash/Planner/Plans/PhysicalConvergentAggregation.h>
 #include <Interpreters/Context.h>
 
 namespace DB
@@ -155,16 +157,36 @@ void PhysicalAggregation::buildBlockInputStreamImpl(DAGPipeline & pipeline, Cont
 
 void PhysicalAggregation::buildPipeline(PipelineBuilder & builder)
 {
+    auto agg_context = std::make_shared<AggregateContext>(
+        is_final_agg,
+        log->identifier());
+    // TODO support fine grained shuffle.
+    assert(!fine_grained_shuffle.enable());
+    auto pre_agg = std::make_shared<PhysicalBuildAggregation>(
+        executor_id,
+        schema,
+        log->identifier(),
+        child,
+        before_agg_actions,
+        aggregation_keys,
+        aggregation_collators,
+        is_final_agg,
+        aggregate_descriptions,
+        expr_after_agg,
+        agg_context);
     // Break the pipeline for pre-agg.
-    // FIXME: Should be newly created PhysicalPreAgg.
-    auto pre_agg_builder = builder.breakPipeline(shared_from_this());
+    auto pre_agg_builder = builder.breakPipeline(pre_agg);
     // Pre-agg pipeline.
     child->buildPipeline(pre_agg_builder);
     pre_agg_builder.build();
     // Final-agg pipeline.
-    // FIXME: Should be newly created PhysicalFinalAgg.
-    builder.addPlanNode(shared_from_this());
-    throw Exception("Unsupport");
+    auto convergent_agg = std::make_shared<PhysicalConvergentAggregation>(
+        executor_id,
+        schema,
+        log->identifier(),
+        agg_context,
+        expr_after_agg);
+    builder.addPlanNode(convergent_agg);
 }
 
 void PhysicalAggregation::finalize(const Names & parent_require)

dbms/src/Flash/Planner/Plans/PhysicalBuildAggregation.cpp

@@ -0,0 +1,54 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Flash/Coprocessor/AggregationInterpreterHelper.h>
+#include <Flash/Planner/Plans/PhysicalBuildAggregation.h>
+#include <Operators/AggregateSinkOp.h>
+#include <Operators/ExpressionTransformOp.h>
+
+namespace DB
+{
+void PhysicalBuildAggregation::buildPipelineExec(PipelineExecGroupBuilder & group_builder, Context & context, size_t /*concurrency*/)
+{
+    if (!before_agg_actions->getActions().empty())
+    {
+        group_builder.transform([&](auto & builder) {
+            builder.appendTransformOp(std::make_unique<ExpressionTransformOp>(group_builder.exec_status, before_agg_actions, log->identifier()));
+        });
+    }
+
+    size_t build_index = 0;
+    group_builder.transform([&](auto & builder) {
+        builder.setSinkOp(std::make_unique<AggregateSinkOp>(group_builder.exec_status, build_index++, agg_context, log->identifier()));
+    });
+
+    Block before_agg_header = group_builder.getCurrentHeader();
+    size_t concurrency = group_builder.concurrency;
+    AggregationInterpreterHelper::fillArgColumnNumbers(aggregate_descriptions, before_agg_header);
+    SpillConfig spill_config(context.getTemporaryPath(), fmt::format("{}_aggregation", log->identifier()), context.getSettingsRef().max_cached_data_bytes_in_spiller, context.getSettingsRef().max_spilled_rows_per_file, context.getSettingsRef().max_spilled_bytes_per_file, context.getFileProvider());
+
+    auto params = AggregationInterpreterHelper::buildParams(
+        context,
+        before_agg_header,
+        concurrency,
+        concurrency,
+        aggregation_keys,
+        aggregation_collators,
+        aggregate_descriptions,
+        is_final_agg,
+        spill_config);
+
+    agg_context->init(params, concurrency);
+}
+} // namespace DB

dbms/src/Flash/Planner/Plans/PhysicalBuildAggregation.h

@@ -0,0 +1,66 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
+#include <Flash/Planner/Plans/PhysicalUnary.h>
+#include <Flash/Planner/Plans/PipelineBreakerHelper.h>
+#include <Interpreters/AggregateDescription.h>
+#include <Interpreters/Context.h>
+#include <Interpreters/ExpressionActions.h>
+#include <Operators/AggregateContext.h>
+
+namespace DB
+{
+class PhysicalBuildAggregation : public PhysicalUnary
+{
+public:
+    PhysicalBuildAggregation(
+        const String & executor_id_,
+        const NamesAndTypes & schema_,
+        const String & req_id,
+        const PhysicalPlanNodePtr & child_,
+        const ExpressionActionsPtr & before_agg_actions_,
+        const Names & aggregation_keys_,
+        const TiDB::TiDBCollators & aggregation_collators_,
+        bool is_final_agg_,
+        const AggregateDescriptions & aggregate_descriptions_,
+        const ExpressionActionsPtr & expr_after_agg_,
+        const AggregateContextPtr & agg_context_)
+        : PhysicalUnary(executor_id_, PlanType::PreAggregation, schema_, req_id, child_)
+        , before_agg_actions(before_agg_actions_)
+        , aggregation_keys(aggregation_keys_)
+        , aggregation_collators(aggregation_collators_)
+        , is_final_agg(is_final_agg_)
+        , aggregate_descriptions(aggregate_descriptions_)
+        , expr_after_agg(expr_after_agg_)
+        , agg_context(agg_context_)
+    {}
+
+    void buildPipelineExec(PipelineExecGroupBuilder & group_builder, Context & context, size_t /*concurrency*/) override;
+
+private:
+    DISABLE_USELESS_FUNCTION_FOR_BREAKER
+
+private:
+    ExpressionActionsPtr before_agg_actions;
+    Names aggregation_keys;
+    TiDB::TiDBCollators aggregation_collators;
+    bool is_final_agg;
+    AggregateDescriptions aggregate_descriptions;
+    ExpressionActionsPtr expr_after_agg;
+    AggregateContextPtr agg_context;
+};
+} // namespace DB

dbms/src/Flash/Planner/Plans/PhysicalConvergentAggregation.cpp

@@ -0,0 +1,55 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include <Flash/Planner/Plans/PhysicalConvergentAggregation.h>
+#include <Operators/AggregateConvergentSourceOp.h>
+#include <Operators/ExpressionTransformOp.h>
+#include <Operators/NullSourceOp.h>
+
+namespace DB
+{
+
+void PhysicalConvergentAggregation::buildPipelineExec(PipelineExecGroupBuilder & group_builder, Context & /*context*/, size_t /*concurrency*/)
+{
+    aggregate_context->initConvergent();
+    if (aggregate_context->useNullSource())
+    {
+        group_builder.init(1);
+        group_builder.transform([&](auto & builder) {
+            builder.setSourceOp(std::make_unique<NullSourceOp>(
+                group_builder.exec_status,
+                aggregate_context->getHeader(),
+                log->identifier()));
+        });
+    }
+    else
+    {
+        group_builder.init(aggregate_context->getConcurrency());
+        size_t index = 0;
+        group_builder.transform([&](auto & builder) {
+            builder.setSourceOp(std::make_unique<AggregateConvergentSourceOp>(
+                group_builder.exec_status,
+                aggregate_context,
+                index++,
+                log->identifier()));
+        });
+    }
+
+    if (!expr_after_agg->getActions().empty())
+    {
+        group_builder.transform([&](auto & builder) {
+            builder.appendTransformOp(std::make_unique<ExpressionTransformOp>(group_builder.exec_status, expr_after_agg, log->identifier()));
+        });
+    }
+}
+} // namespace DB

dbms/src/Flash/Planner/Plans/PhysicalConvergentAggregation.h

@@ -0,0 +1,48 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
+#include <Flash/Planner/Plans/PhysicalLeaf.h>
+#include <Flash/Planner/Plans/PipelineBreakerHelper.h>
+#include <Interpreters/ExpressionActions.h>
+#include <Operators/AggregateContext.h>
+
+namespace DB
+{
+class PhysicalConvergentAggregation : public PhysicalLeaf
+{
+public:
+    PhysicalConvergentAggregation(
+        const String & executor_id_,
+        const NamesAndTypes & schema_,
+        const String & req_id,
+        const AggregateContextPtr & aggregate_context_,
+        const ExpressionActionsPtr & expr_after_agg_)
+        : PhysicalLeaf(executor_id_, PlanType::ConvergentAggregation, schema_, req_id)
+        , expr_after_agg(expr_after_agg_)
+        , aggregate_context(aggregate_context_)
+    {}
+
+    void buildPipelineExec(PipelineExecGroupBuilder & group_builder, Context & /*context*/, size_t concurrency) override;
+
+private:
+    DISABLE_USELESS_FUNCTION_FOR_BREAKER
+
+private:
+    ExpressionActionsPtr expr_after_agg;
+    AggregateContextPtr aggregate_context;
+};
+} // namespace DB

dbms/src/Flash/Planner/Plans/PipelineBreakerHelper.h

@@ -0,0 +1,36 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+namespace DB
+{
+#define DISABLE_USELESS_FUNCTION_FOR_BREAKER                                  \
+    void buildPipeline(PipelineBuilder &) override                            \
+    {                                                                         \
+        throw Exception("Unsupport");                                         \
+    }                                                                         \
+    void finalize(const Names &) override                                     \
+    {                                                                         \
+        throw Exception("Unsupport");                                         \
+    }                                                                         \
+    const Block & getSampleBlock() const override                             \
+    {                                                                         \
+        throw Exception("Unsupport");                                         \
+    }                                                                         \
+    void buildBlockInputStreamImpl(DAGPipeline &, Context &, size_t) override \
+    {                                                                         \
+        throw Exception("Unsupport");                                         \
+    }
+} // namespace DB