grouping function

6 files changed, 93 insertions, 8 deletions

diff --git a/ydb/library/yql/core/common_opt/CMakeLists.txt b/ydb/library/yql/core/common_opt/CMakeLists.txt
index aa19918c9b3..df2157294c9 100644
--- a/ydb/library/yql/core/common_opt/CMakeLists.txt
+++ b/ydb/library/yql/core/common_opt/CMakeLists.txt
@@ -29,5 +29,6 @@ target_sources(yql-core-common_opt PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
   ${CMAKE_SOURCE_DIR}/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
   ${CMAKE_SOURCE_DIR}/ydb/library/yql/core/common_opt/yql_co_simple2.cpp
+  ${CMAKE_SOURCE_DIR}/ydb/library/yql/core/common_opt/yql_co_simple3.cpp
   ${CMAKE_SOURCE_DIR}/ydb/library/yql/core/common_opt/yql_co_transformer.cpp
 )
diff --git a/ydb/library/yql/core/common_opt/yql_co.h b/ydb/library/yql/core/common_opt/yql_co.h
index 4c82b44241f..167dc25f97d 100644
--- a/ydb/library/yql/core/common_opt/yql_co.h
+++ b/ydb/library/yql/core/common_opt/yql_co.h
@@ -79,6 +79,7 @@ struct TCoCallableRules {
     enum {
         SIMPLE_STEP_1,
         SIMPLE_STEP_2,
+        SIMPLE_STEP_3,
         SIMPLE_STEPS
     };
 
@@ -104,6 +105,7 @@ struct TCoCallableRules {
 
 void RegisterCoSimpleCallables1(TCallableOptimizerMap& map);
 void RegisterCoSimpleCallables2(TCallableOptimizerMap& map);
+void RegisterCoSimpleCallables3(TCallableOptimizerMap& map);
 void RegisterCoFlowCallables1(TCallableOptimizerMap& map);
 void RegisterCoFlowCallables2(TCallableOptimizerMap& map);
 void RegisterCoFinalizers(TFinalizingOptimizerMap& map);
diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
index 3ed7ec37864..ab0a21c8c9b 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -1889,6 +1889,44 @@ TExprNode::TPtr BuildGroup(TPositionHandle pos, TExprNode::TPtr list,
             .Seal()
             .Build();
 
+        if (currentKeys.size() < groupExprs->Tail().ChildrenSize()) {
+            // mark missing columns
+            aggregate = ctx.Builder(pos)
+                .Callable("OrderedMap")
+                    .Add(0, aggregate)
+                    .Lambda(1)
+                        .Param("row")
+                        .Callable("FlattenMembers")
+                            .List(0)
+                                .Atom(0, "")
+                                .Arg(1, "row")
+                            .Seal()
+                            .List(1)
+                                .Atom(0, "")
+                                .Callable(1, "AsStruct")
+                                    .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
+                                        ui32 j = 0;
+                                        for (ui32 i = 0; i < groupExprs->Tail().ChildrenSize(); ++i) {
+                                            if (!currentKeys.contains(i)) {
+                                                parent.List(j++)
+                                                    .Atom(0, "_yql_grouping_" + ToString(i))
+                                                    .Callable(1, "Int32")
+                                                        .Atom(0, "1")
+                                                    .Seal()
+                                                    .Seal();
+                                            }
+                                        }
+
+                                        return parent;
+                                    })
+                                .Seal()
+                            .Seal()
+                        .Seal()
+                    .Seal()
+                .Seal()
+                .Build();
+        }
+
         unionAllItems.push_back(aggregate);
         // shift iterator
         ui32 i = 0;
@@ -3323,15 +3361,49 @@ TExprNode::TPtr ExpandPgGroupRef(const TExprNode::TPtr& node, TExprContext& ctx,
 }
 
 TExprNode::TPtr ExpandPgGrouping(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
-    Y_UNUSED(node);
     Y_UNUSED(optCtx);
-    return ctx.Builder(node->Pos())
-        .Callable("PgConst")
-            .Atom(0, "0")
-            .Callable(1, "PgType")
-                .Atom(0, "int4")
+    TExprNode::TPtr sum;
+    YQL_ENSURE(node->ChildrenSize() >= 1);
+    for (ui32 i = 0; i < node->ChildrenSize(); ++i) {
+        auto child = node->Child(i);
+        YQL_ENSURE(child->IsCallable("PgGroupRef"));
+        auto row = child->HeadPtr();
+        auto index = child->Tail().Content();
+        auto value = ctx.Builder(node->Pos())
+            .Callable("Coalesce")
+                .Callable(0, "TryMember")
+                    .Add(0, row)
+                    .Atom(1, TString("_yql_grouping_") + index)
+                    .Callable(2, "Null")
+                    .Seal()
+                .Seal()
+                .Callable(1, "Int32")
+                    .Atom(0, "0")
+                .Seal()
             .Seal()
-        .Seal()
+            .Build();
+
+        if (!sum) {
+            sum = value;
+        } else {
+            sum = ctx.Builder(node->Pos())
+                .Callable("+")
+                    .Callable(0, "*")
+                        .Add(0, sum)
+                        .Callable(1, "Int32")
+                            .Atom(0, "2")
+                        .Seal()
+                    .Seal()
+                    .Add(1, value)
+                .Seal()
+                .Build();
+        }
+    }
+
+    return ctx.Builder(node->Pos())

+        .Callable("ToPg")

+            .Add(0, sum)

+        .Seal()

         .Build();
 }
 
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple2.cpp b/ydb/library/yql/core/common_opt/yql_co_simple2.cpp
index 5531d687df4..100bc88cd42 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple2.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple2.cpp
@@ -500,7 +500,6 @@ void RegisterCoSimpleCallables2(TCallableOptimizerMap& map) {
         return node;
     };
 
-    map["PgGroupRef"] = ExpandPgGroupRef;
     map["PgGrouping"] = ExpandPgGrouping;
 }
 
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple3.cpp b/ydb/library/yql/core/common_opt/yql_co_simple3.cpp
new file mode 100644
index 00000000000..fbcd8edde8d
--- /dev/null
+++ b/ydb/library/yql/core/common_opt/yql_co_simple3.cpp
@@ -0,0 +1,10 @@
+#include "yql_co.h"
+#include "yql_co_pgselect.h"
+
+namespace NYql {
+
+void RegisterCoSimpleCallables3(TCallableOptimizerMap& map) {
+    map["PgGroupRef"] = ExpandPgGroupRef;
+}
+
+}
diff --git a/ydb/library/yql/core/common_opt/yql_co_transformer.cpp b/ydb/library/yql/core/common_opt/yql_co_transformer.cpp
index 42e2cbb5fb4..6872461ca2f 100644
--- a/ydb/library/yql/core/common_opt/yql_co_transformer.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_transformer.cpp
@@ -183,6 +183,7 @@ const TCoCallableRules& TCoCallableRules::Instance() {
 TCoCallableRules::TCoCallableRules() {
     RegisterCoSimpleCallables1(SimpleCallables[SIMPLE_STEP_1]);
     RegisterCoSimpleCallables2(SimpleCallables[SIMPLE_STEP_2]);
+    RegisterCoSimpleCallables3(SimpleCallables[SIMPLE_STEP_3]);
     RegisterCoFlowCallables1(FlowCallables[FLOW_STEP_1]);
     RegisterCoFlowCallables2(FlowCallables[FLOW_STEP_2]);
     RegisterCoFinalizers(Finalizers);