YQL-14629 fixed peephole

ref:4aabe645f25e6dc11bca71d9e7dcfd56c9ecc2f3

7 files changed, 116 insertions, 122 deletions

diff --git a/ydb/core/kqp/provider/yql_kikimr_opt.cpp b/ydb/core/kqp/provider/yql_kikimr_opt.cpp
index cdbb6e600dd..e04c518f2f7 100644
--- a/ydb/core/kqp/provider/yql_kikimr_opt.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_opt.cpp
@@ -120,7 +120,7 @@ TExprNode::TPtr KiRewriteAggregate(TExprBase node, TExprContext& ctx) {
     }
 
     YQL_CLOG(INFO, ProviderKikimr) << "KiRewriteAggregate";
-    return ExpandAggregate(node.Ptr(), ctx);
+    return ExpandAggregate(true, node.Ptr(), ctx);
 }
 
 TExprNode::TPtr KiRedundantSortByPk(TExprBase node, TExprContext& ctx,
diff --git a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
index 128e1883b76..2df9e5af1e5 100644
--- a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -2234,78 +2234,74 @@ TExprNode::TPtr ExpandContainerIf(const TExprNode::TPtr& input, TExprContext& ct
 
 TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext& ctx) {
     YQL_CLOG(DEBUG, CorePeepHole) << "Expand " << node->Content();
+    const bool isStream = node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Flow ||
+        node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Stream;
     TExprNode::TPtr sort;
-    if (node->Child(2)->IsCallable("Void") || node->Child(3)->IsCallable("Void")) {
-        sort = ctx.NewCallable(node->Pos(), "Sort", {node->HeadPtr(), MakeBool<true>(node->Pos(), ctx), node->ChildPtr(1)});
+    auto keyExtractor = node->ChildPtr(1);
+    const bool haveSort = !node->Child(2)->IsCallable("Void");
+    auto idLambda = ctx.Builder(node->Pos())
+        .Lambda()
+            .Param("x")
+            .Arg("x")
+        .Seal()
+        .Build();
+
+    auto sortLambda =  ctx.Builder(node->Pos())
+        .Lambda()
+            .Param("x")
+            .Callable("Sort")
+                .Arg(0, "x")
+                .Add(1, node->ChildPtr(2))
+                .Add(2, node->ChildPtr(3))
+            .Seal()
+        .Seal()
+        .Build();
+
+    auto settings = ctx.Builder(node->Pos())
+        .List()
+            .Atom(0, "Hashed")
+            .Atom(1, "Many")
+            .Atom(2, "Compact")
+        .Seal()
+        .Build();
+
+    auto flatten = ctx.Builder(node->Pos())
+        .Lambda()
+            .Param("dict")
+            .Callable("OrderedFlatMap")
+                .Callable(0, "DictPayloads")
+                    .Arg(0, "dict")
+                .Seal()
+                .Add(1, haveSort ? sortLambda : idLambda)
+            .Seal()
+        .Seal()
+        .Build();
+
+    if (isStream) {
+        sort = ctx.Builder(node->Pos())
+            .Callable("OrderedFlatMap")
+                .Callable(0, "SqueezeToDict")
+                    .Add(0, node->HeadPtr())
+                    .Add(1, keyExtractor)
+                    .Add(2, idLambda)
+                    .Add(3, settings)
+                .Seal()
+                .Add(1, flatten)
+            .Seal()
+            .Build();
     } else {
-        if (node->Child(2)->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple) {
-            if (const auto size = node->Child(2)->GetTypeAnn()->Cast<TTupleExprType>()->GetSize()) {
-                sort = ctx.Builder(node->Pos())
-                    .Callable("Sort")
+        sort = ctx.Builder(node->Pos())
+            .Apply(flatten)
+                .With(0)
+                    .Callable("ToDict")
                         .Add(0, node->HeadPtr())
-                        .List(1)
-                            .Callable(0, "Bool")
-                                .Atom(0, "true", TNodeFlags::Default)
-                            .Seal()
-                            .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                                const bool isList = node->Child(2)->IsList();
-                                for (ui32 i = 0U; i < size; ++i) {
-                                    if (isList) {
-                                        parent.Add(i + 1U, node->Child(2)->ChildPtr(i));
-                                    } else {
-                                        parent.Callable(i + 1U, "Nth")
-                                                .Add(0, node->ChildPtr(2))
-                                                .Atom(1, ToString(i), TNodeFlags::Default)
-                                            .Seal();
-                                    }
-                                }
-                                return parent;
-                            })
-                        .Seal()
-                        .Lambda(2)
-                            .Param("item")
-                            .List()
-                                .Apply(0, *node->Child(1))
-                                    .With(0, "item")
-                                .Seal()
-                                .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                                    for (ui32 pos = 0U; pos < size; ++pos) {
-                                        parent.Callable(pos + 1U, "Nth")
-                                            .Apply(0, *node->Child(3))
-                                                .With(0, "item")
-                                            .Seal()
-                                            .Atom(1, ToString(pos), TNodeFlags::Default)
-                                        .Seal();
-                                    }
-                                    return parent;
-                                })
-                            .Seal()
-                        .Seal()
-                    .Seal().Build();
-            }
-        } else {
-            sort = ctx.Builder(node->Pos())
-                .Callable("Sort")
-                    .Add(0, node->HeadPtr())
-                    .List(1)
-                        .Callable(0, "Bool")
-                            .Atom(0, "true", TNodeFlags::Default)
-                        .Seal()
-                        .Add(1, node->ChildPtr(2))
-                    .Seal()
-                    .Lambda(2)
-                        .Param("item")
-                        .List()
-                            .Apply(0, *node->Child(1))
-                                .With(0, "item")
-                            .Seal()
-                            .Apply(1, *node->Child(3))
-                                .With(0, "item")
-                            .Seal()
-                        .Seal()
+                        .Add(1, keyExtractor)
+                        .Add(2, idLambda)
+                        .Add(3, settings)
                     .Seal()
-                .Seal().Build();
-        }
+                .Done()
+            .Seal()
+            .Build();
     }
 
     return ctx.ReplaceNode(node->Tail().TailPtr(), node->Tail().Head().Head(), std::move(sort));
@@ -5789,7 +5785,7 @@ struct TPeepHoleRules {
         {"OptionalReduce", &ExpandOptionalReduce},
         {"AggrMin", &ExpandAggrMinMax<true>},
         {"AggrMax", &ExpandAggrMinMax<false>},
-        {"Aggregate", &ExpandAggregateCompact},
+        {"Aggregate", &ExpandAggregatePeephole},
         {"And", &OptimizeLogicalDups<true>},
         {"Or", &OptimizeLogicalDups<false>},
         {"CombineByKey", &ExpandCombineByKey},
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 4b099648326..1fe8d5d3eb6 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -3107,9 +3107,12 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        if (!lambdaKeySelector->GetTypeAnn()->IsHashable() || !lambdaKeySelector->GetTypeAnn()->IsEquatable()) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambdaKeySelector->Pos()), TStringBuilder()
-                << "Expected hashable and equatable type for key, but got: " << *lambdaKeySelector->GetTypeAnn()));
+        auto keyType = lambdaKeySelector->GetTypeAnn();
+        if (!EnsureHashableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureEquatableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
@@ -3192,9 +3195,12 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        if (!lambdaKeySelector->GetTypeAnn()->IsHashable() || !lambdaKeySelector->GetTypeAnn()->IsEquatable()) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambdaKeySelector->Pos()), TStringBuilder()
-                << "Expected hashable and equatable type for key, but got: " << *lambdaKeySelector->GetTypeAnn()));
+        auto keyType = lambdaKeySelector->GetTypeAnn();
+        if (!EnsureHashableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureEquatableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
@@ -3882,13 +3888,12 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        if (!EnsureOneOrTupleOfDataOrOptionalOfData(*lambdaKeySelector, ctx.Expr)) {
+        auto keyType = lambdaKeySelector->GetTypeAnn();
+        if (!EnsureHashableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
-        if (lambdaKeySelector->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple &&
-            lambdaKeySelector->GetTypeAnn()->Cast<TTupleExprType>()->GetSize() < 2) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambdaKeySelector->Pos()), "Tuple must contain at least 2 items"));
+        if (!EnsureEquatableKey(lambdaKeySelector->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
@@ -5896,9 +5901,6 @@ namespace {
         if (!keyExtractor->GetTypeAnn()) {
             return IGraphTransformer::TStatus::Repeat;
         }
-        if (!EnsureOneOrTupleOfDataOrOptionalOfData(*keyExtractor, ctx.Expr)) {
-            return IGraphTransformer::TStatus::Error;
-        }
         auto keyType = keyExtractor->GetTypeAnn();
         if (!EnsureHashableKey(keyExtractor->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
@@ -5992,14 +5994,16 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        if (!keyExtractor->GetTypeAnn()->IsHashable() || !keyExtractor->GetTypeAnn()->IsEquatable()) {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(keyExtractor->Pos()), TStringBuilder()
-                << "Expected hashable and equatable type for key, but got: " << *keyExtractor->GetTypeAnn()));
+        auto keyType = keyExtractor->GetTypeAnn();
+        if (!EnsureHashableKey(keyExtractor->Pos(), keyType, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureEquatableKey(keyExtractor->Pos(), keyType, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
 
         // groupSwitch
-        auto keyType = keyExtractor->GetTypeAnn();
         if (!UpdateLambdaAllArgumentsTypes(groupSwitch, {keyType, itemType}, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
         }
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index 08225d058b2..3dfbd353314 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -2849,49 +2849,43 @@ bool IsInstantEqual(const TTypeAnnotationNode& type) {
     return true;
 }
 
-static bool EnsureKeyProperty(TPositionHandle position, const TTypeAnnotationNode* keyType,
-    const std::function<bool(TPositionHandle, EDataSlot, TExprContext&)>& propertyCheck, TExprContext& ctx) {
-
+bool EnsureHashableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
     if (HasError(keyType, ctx)) {
         return false;
     }
 
-    TTypeAnnotationNode::TListType keyTypes;
-    if (keyType->GetKind() != ETypeAnnotationKind::Tuple) {
-        keyTypes.push_back(keyType);
-    }
-    else {
-        auto tuple = keyType->Cast<TTupleExprType>();
-        for (auto& x : tuple->GetItems()) {
-            keyTypes.push_back(x);
-        }
+    if (!keyType->IsHashable()) {
+        ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected hashable type, but got: " << *keyType));
+        return false;
     }
 
-    for (auto& x : keyTypes) {
-        bool isOptional;
-        const TDataExprType* dataType;
-        if (!EnsureDataOrOptionalOfData(position, x, isOptional, dataType, ctx)) {
-            return false;
-        }
+    return true;
+}
 
-        if (!propertyCheck(position, dataType->GetSlot(), ctx)) {
-            return false;
-        }
+bool EnsureComparableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
+    if (HasError(keyType, ctx)) {
+        return false;
+    }
+
+    if (!keyType->IsComparable()) {
+        ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected comparable type, but got: " << *keyType));
+        return false;
     }
 
     return true;
 }
 
-bool EnsureHashableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
-    return EnsureKeyProperty(position, keyType, EnsureHashableDataType, ctx);
-}
+bool EnsureEquatableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
+    if (HasError(keyType, ctx)) {
+        return false;
+    }
 
-bool EnsureComparableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
-    return EnsureKeyProperty(position, keyType, EnsureComparableDataType, ctx);
-}
+    if (!keyType->IsEquatable()) {
+        ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() << "Expected equatable type, but got: " << *keyType));
+        return false;
+    }
 
-bool EnsureEquatableKey(TPositionHandle position, const TTypeAnnotationNode* keyType, TExprContext& ctx) {
-    return EnsureKeyProperty(position, keyType, EnsureEquatableDataType, ctx);
+    return true;
 }
 
 bool UpdateLambdaAllArgumentsTypes(TExprNode::TPtr& lambda, const std::vector<const TTypeAnnotationNode*>& argumentsAnnotations, TExprContext& ctx) {
diff --git a/ydb/library/yql/core/yql_opt_aggregate.cpp b/ydb/library/yql/core/yql_opt_aggregate.cpp
index 7e5d24dc4b2..662402e6315 100644
--- a/ydb/library/yql/core/yql_opt_aggregate.cpp
+++ b/ydb/library/yql/core/yql_opt_aggregate.cpp
@@ -5,7 +5,7 @@
 
 namespace NYql {
 
-TExprNode::TPtr ExpandAggregate(const TExprNode::TPtr& node, TExprContext& ctx, bool forceCompact) {
+TExprNode::TPtr ExpandAggregate(bool allowPickle, const TExprNode::TPtr& node, TExprContext& ctx, bool forceCompact) {
     auto list = node->HeadPtr();
     auto keyColumns = node->ChildPtr(1);
     auto aggregatedColumns = node->Child(2);
@@ -124,7 +124,7 @@ TExprNode::TPtr ExpandAggregate(const TExprNode::TPtr& node, TExprContext& ctx,
         YQL_ENSURE(index, "Unknown column: " << keyColumn->Content());
         auto type = rowType->GetItems()[*index]->GetItemType();
         keyItemTypes.push_back(type);
-        needPickle = needPickle || !IsDataOrOptionalOfData(type);
+        needPickle = needPickle || allowPickle && !IsDataOrOptionalOfData(type);
     }
 
     const TTypeAnnotationNode* pickleType = nullptr;
diff --git a/ydb/library/yql/core/yql_opt_aggregate.h b/ydb/library/yql/core/yql_opt_aggregate.h
index d254bffecf4..9707405544c 100644
--- a/ydb/library/yql/core/yql_opt_aggregate.h
+++ b/ydb/library/yql/core/yql_opt_aggregate.h
@@ -3,9 +3,9 @@
 
 namespace NYql {
 
-TExprNode::TPtr ExpandAggregate(const TExprNode::TPtr& node, TExprContext& ctx, bool forceCompact = false);
-inline TExprNode::TPtr ExpandAggregateCompact(const TExprNode::TPtr& node, TExprContext& ctx) {
-    return ExpandAggregate(node, ctx, true);
+TExprNode::TPtr ExpandAggregate(bool allowPickle, const TExprNode::TPtr& node, TExprContext& ctx, bool forceCompact = false);
+inline TExprNode::TPtr ExpandAggregatePeephole(const TExprNode::TPtr& node, TExprContext& ctx) {
+    return ExpandAggregate(false, node, ctx, true);
 }
 
 }
diff --git a/ydb/library/yql/dq/opt/dq_opt_log.cpp b/ydb/library/yql/dq/opt/dq_opt_log.cpp
index d2aded1075b..333075d777c 100644
--- a/ydb/library/yql/dq/opt/dq_opt_log.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_log.cpp
@@ -18,7 +18,7 @@ TExprBase DqRewriteAggregate(TExprBase node, TExprContext& ctx) {
         return node;
     }
 
-    auto result = ExpandAggregate(node.Ptr(), ctx);
+    auto result = ExpandAggregate(true, node.Ptr(), ctx);
     YQL_ENSURE(result);
 
     return TExprBase(result);