Merge 08e8f83 into f12f073

ydb-platform · Nov 4, 2024 · 9976d44 · 9976d44
2 parents f12f073 + 08e8f83
commit 9976d44
Showing 1 changed file with 38 additions and 0 deletions.
diff --git a/ydb/core/kqp/opt/kqp_query_plan.cpp b/ydb/core/kqp/opt/kqp_query_plan.cpp
@@ -954,6 +954,8 @@ class TxPlanSerializer {
             operatorId = Visit(maybeCondense.Cast(), planNode);
         } else if (auto maybeCombiner = TMaybeNode<TCoCombineCore>(node)) {
             operatorId = Visit(maybeCombiner.Cast(), planNode);
+        } else if (auto maybeBlockCombine = TMaybeNode<TCoBlockCombineHashed>(node)) {
+            operatorId = Visit(maybeBlockCombine.Cast(), planNode);
         } else if (auto maybeSort = TMaybeNode<TCoSort>(node)) {
             operatorId = Visit(maybeSort.Cast(), planNode);
         } else if (auto maybeTop = TMaybeNode<TCoTop>(node)) {
@@ -1026,6 +1028,8 @@ class TxPlanSerializer {
                 } else {
                     operatorId = Visit(olapTable, planNode);
                 }
+            } else if (TMaybeNode<TCoToFlow>(node)) {
+                // do nothing
             } else {
                 for (const auto& child : node->Children()) {
                     if(!child->IsLambda()) {
@@ -1168,6 +1172,40 @@ class TxPlanSerializer {
         return AddOperator(planNode, "Aggregate", std::move(op));
     }
 
+    std::variant<ui32, TArgContext> Visit(const TCoBlockCombineHashed& blockCombine, TQueryPlanNode& planNode) {
+
+        static const THashMap<TString, TString> aggregations = {
+            {"count", "COUNT"},
+            {"count_all", "COUNT"},
+            {"min", "MIN"},
+            {"max", "MAX"},
+            {"avg", "AVG"},
+            {"sum", "SUM"}
+        };
+
+        TOperator op;
+        op.Properties["Name"] = "Aggregate";
+        op.Properties["GroupBy"] = NPlanUtils::PrettyExprStr(blockCombine.Keys());
+
+        if (blockCombine.Aggregations().Ref().IsList()) {
+            TVector<TString> aggrs;
+            for (ui32 index = 0; index < blockCombine.Aggregations().Ref().ChildrenSize(); index++) {
+                auto child = blockCombine.Aggregations().Ref().Child(index);
+                if (child && child->IsList() && child->ChildrenSize() >= 2) {
+                    auto callable = child->Child(0);
+                    if (callable->ChildrenSize() >= 1) {
+                        if (auto aggrName = TMaybeNode<TCoAtom>(callable->Child(0))) {
+                            aggrs.push_back(aggregations.Value(aggrName.Cast().StringValue(), "??"));
+                        }
+                    }
+                }
+            }
+            op.Properties["Aggregation"] = JoinStrings(std::move(aggrs), ",");
+        }
+
+        return AddOperator(planNode, "Aggregate", std::move(op));
+    }
+
     std::variant<ui32, TArgContext> Visit(const TCoSort& sort, TQueryPlanNode& planNode) {
         TOperator op;
         op.Properties["Name"] = "Sort";