CountAggregateRewrite: support of AggApply

3 files changed, 93 insertions, 67 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_flow1.cpp b/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
index 4b61c4ebd85..251ccc1efac 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
@@ -1269,81 +1269,113 @@ TExprNode::TPtr CountAggregateRewrite(const TCoAggregate& node, TExprContext& ct
     const bool isDistinct = (aggregatedColumn.Ref().ChildrenSize() == 3);
 
     auto traits = aggregatedColumn.Ref().Child(1);
-    if (!traits->IsCallable("AggregationTraits")) {
-        return node.Ptr();
-    }
-
     auto outputColumn = aggregatedColumn.Ref().HeadPtr();
-    // validation of traits
-    auto inputItemType = traits->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
-    const bool isOptionalColumn = inputItemType->GetKind() == ETypeAnnotationKind::Optional;
 
-    auto init = TCoLambda(traits->Child(1));
+    // validation of traits
+    const TTypeAnnotationNode* inputItemType;
     bool onlyColumn = true;
     bool onlyZero = true;
     TExprNode::TPtr initVal;
-    TExprNode::TPtr updateVal;
-    if (init.Body().Ref().IsCallable("Uint64") &&
-        init.Body().Ref().Head().Content() == "1") {
-        onlyZero = false;
-    } else if (init.Body().Ref().IsCallable("Uint64") &&
-        init.Body().Ref().Head().Content() == "0") {
-        onlyColumn = false;
-    } else if (init.Body().Ref().IsCallable("AggrCountInit")) {
-        initVal = init.Body().Ref().HeadPtr();
-        onlyColumn = onlyColumn && init.Body().Ref().Child(0) == init.Args().Arg(0).Raw();
-        onlyZero = false;
-    } else {
-        return node.Ptr();
-    }
+    if (traits->IsCallable("AggregationTraits")) {
+        inputItemType = traits->Head().GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+
+        auto init = TCoLambda(traits->Child(1));
+        TExprNode::TPtr updateVal;
+        if (init.Body().Ref().IsCallable("Uint64") &&
+            init.Body().Ref().Head().Content() == "1") {
+            onlyZero = false;
+        } else if (init.Body().Ref().IsCallable("Uint64") &&
+            init.Body().Ref().Head().Content() == "0") {
+            onlyColumn = false;
+        } else if (init.Body().Ref().IsCallable("AggrCountInit")) {
+            initVal = init.Body().Ref().HeadPtr();
+            onlyColumn = onlyColumn && init.Body().Ref().Child(0) == init.Args().Arg(0).Raw();
+            onlyZero = false;
+        } else {
+            return node.Ptr();
+        }
 
-    auto update = TCoLambda(traits->Child(2));
-    auto inc = update.Body().Ptr();
-    if (inc->IsCallable("Inc") && inc->Child(0) == update.Args().Arg(1).Raw()) {
-        onlyZero = false;
-    } else if (inc->IsCallable("AggrCountUpdate") && inc->Child(1) == update.Args().Arg(1).Raw()) {
-        updateVal = inc->HeadPtr();
-        onlyColumn = onlyColumn && inc->Child(0) == update.Args().Arg(0).Raw();
-        onlyZero = false;
-    } else if (inc == update.Args().Arg(1).Raw()) {
-        onlyColumn = false;
-    } else {
-        return node.Ptr();
-    }
+        auto update = TCoLambda(traits->Child(2));
+        auto inc = update.Body().Ptr();
+        if (inc->IsCallable("Inc") && inc->Child(0) == update.Args().Arg(1).Raw()) {
+            onlyZero = false;
+        } else if (inc->IsCallable("AggrCountUpdate") && inc->Child(1) == update.Args().Arg(1).Raw()) {
+            updateVal = inc->HeadPtr();
+            onlyColumn = onlyColumn && inc->Child(0) == update.Args().Arg(0).Raw();
+            onlyZero = false;
+        } else if (inc == update.Args().Arg(1).Raw()) {
+            onlyColumn = false;
+        } else {
+            return node.Ptr();
+        }
 
-    auto save = TCoLambda(traits->Child(3));
-    if (save.Body().Raw() != save.Args().Arg(0).Raw()) {
-        return node.Ptr();
-    }
+        auto save = TCoLambda(traits->Child(3));
+        if (save.Body().Raw() != save.Args().Arg(0).Raw()) {
+            return node.Ptr();
+        }
 
-    auto load = TCoLambda(traits->Child(4));
-    if (load.Body().Raw() != load.Args().Arg(0).Raw()) {
-        return node.Ptr();
-    }
+        auto load = TCoLambda(traits->Child(4));
+        if (load.Body().Raw() != load.Args().Arg(0).Raw()) {
+            return node.Ptr();
+        }
 
-    auto merge = TCoLambda(traits->Child(5));
-    {
-        auto& plus = merge.Body().Ref();
-        if (!plus.IsCallable("+")) {
+        auto merge = TCoLambda(traits->Child(5));
+        {
+            auto& plus = merge.Body().Ref();
+            if (!plus.IsCallable("+")) {
+                return node.Ptr();
+            }
+
+            if (!(plus.Child(0) == merge.Args().Arg(0).Raw() &&
+                plus.Child(1) == merge.Args().Arg(1).Raw())) {
+                return node.Ptr();
+            }
+        }
+
+        auto finish = TCoLambda(traits->Child(6));
+        if (finish.Body().Raw() != finish.Args().Arg(0).Raw()) {
             return node.Ptr();
         }
 
-        if (!(plus.Child(0) == merge.Args().Arg(0).Raw() &&
-            plus.Child(1) == merge.Args().Arg(1).Raw())) {
+        auto defVal = traits->Child(7);
+        if (!defVal->IsCallable("Uint64") || defVal->Head().Content() != "0") {
             return node.Ptr();
         }
-    }
 
-    auto finish = TCoLambda(traits->Child(6));
-    if (finish.Body().Raw() != finish.Args().Arg(0).Raw()) {
-        return node.Ptr();
-    }
+        if (!isDistinct) {
+            if (!onlyZero && !onlyColumn) {
+                if (!initVal || !updateVal || initVal != updateVal) {
+                    return node.Ptr();
+                }
+            }
+        }
+    } else if (traits->IsCallable("AggApply")) {
+        if (traits->Head().Content() != "count_all" && traits->Head().Content() != "count") {
+            return node.Ptr();
+        }
 
-    auto defVal = traits->Child(7);
-    if (!defVal->IsCallable("Uint64") || defVal->Head().Content() != "0") {
+        inputItemType = traits->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
+        onlyZero = false;
+        onlyColumn = false;
+        if (&traits->Child(2)->Head().Head() == &traits->Child(2)->Tail()) {
+            onlyColumn = true;
+        }
+
+        if (!isDistinct) {
+            if (IsDepended(traits->Child(2)->Tail(), traits->Child(2)->Head().Head())) {
+                return node.Ptr();
+            }
+
+            if (traits->Head().Content() == "count") {
+                initVal = traits->Child(2)->TailPtr();
+            }
+        }
+    } else {
         return node.Ptr();
     }
 
+    const bool isOptionalColumn = inputItemType->GetKind() == ETypeAnnotationKind::Optional;
+
     if (!isDistinct) {
         auto length = ctx.Builder(node.Pos())
             .Callable("Length")
@@ -1357,11 +1389,7 @@ TExprNode::TPtr CountAggregateRewrite(const TCoAggregate& node, TExprContext& ct
                     .Atom(0, "0", TNodeFlags::Default)
                 .Seal()
                 .Build();
-        } else if (!onlyColumn) {
-            if (!initVal || !updateVal || initVal != updateVal) {
-                return node.Ptr();
-            }
-
+        } else if (!onlyColumn && initVal) {
             length = ctx.Builder(node.Pos())
                 .Callable("If")
                     .Callable(0, "Exists")
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
index e9b27e10c91..ea686d5bab4 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -4874,12 +4874,6 @@ namespace {
         }
 
         auto itemType = input->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
-        if (itemType->GetKind() != ETypeAnnotationKind::Struct) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
-                TStringBuilder() << "Expected struct type, but got: " << *itemType));
-            return IGraphTransformer::TStatus::Error;
-        }
-
         auto& lambda = input->ChildRef(2);
         const auto status = ConvertToLambda(lambda, ctx.Expr, 1);
         if (status.Level != IGraphTransformer::TStatus::Ok) {
diff --git a/ydb/library/yql/core/yql_opt_aggregate.cpp b/ydb/library/yql/core/yql_opt_aggregate.cpp
index edd7b4d26b2..84670cbc146 100644
--- a/ydb/library/yql/core/yql_opt_aggregate.cpp
+++ b/ydb/library/yql/core/yql_opt_aggregate.cpp
@@ -6,6 +6,8 @@
 
 namespace NYql {
 
+namespace {
+
 TExprNode::TPtr ExpandAggApply(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx) {
     auto name = node->Head().Content();
     auto exportsPtr = typesCtx.Modules->GetModule("/lib/yql/aggregate.yql");
@@ -30,6 +32,8 @@ TExprNode::TPtr ExpandAggApply(const TExprNode::TPtr& node, TExprContext& ctx, T
     return traits;
 }
 
+}
+
 TExprNode::TPtr ExpandAggregate(bool allowPickle, const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx, bool forceCompact) {
     auto list = node->HeadPtr();
     auto keyColumns = node->ChildPtr(1);