YQL-15555 YQL-15435 Keep constraints on expand PartitionsByKeys.

4 files changed, 95 insertions, 88 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index 8b6b194314..e6d73f249d 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -95,83 +95,6 @@ bool CanRewriteToEmptyContainer(const TExprNode& src) {
     return true;
 }
 
-TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, TExprContext& ctx) {
-    if (!sorted) {
-        return node;
-    }
-    const auto& constent = sorted->GetContent();
-    return ctx.Builder(node->Pos())
-        .Callable("AssumeSorted")
-            .Add(0, std::move(node))
-            .List(1)
-                .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                    size_t index = 0;
-                    for (auto c : constent) {
-                        parent.Callable(index++, "Bool")
-                            .Atom(0, ToString(c.second), TNodeFlags::Default)
-                        .Seal();
-                        if (1U < c.first.front().size())
-                            break;
-                    }
-                    return parent;
-                })
-            .Seal()
-            .Lambda(2)
-                .Param("item")
-                .List()
-                    .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                        size_t index = 0;
-                        for (auto c : constent) {
-                            if (c.first.front().empty())
-                                parent.Arg(index++, "item");
-                            else if (1U == c.first.front().size()) {
-                                parent.Callable(index++, "Member")
-                                    .Arg(0, "item")
-                                    .Atom(1, c.first.front().front())
-                                .Seal();
-                            } else {
-                                parent.Callable(index++, "Null").Seal();
-                                break;
-                            }
-                        }
-                        return parent;
-                    })
-                .Seal()
-            .Seal()
-        .Seal()
-        .Build();
-}
-
-template<bool Distinct>
-TExprNode::TPtr KeepUniqueConstraint(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx) {
-    if (const auto uniq = src.GetConstraint<TUniqueConstraintNodeBase<Distinct>>()) {
-        TExprNode::TListType columns;
-        for (const auto& set : uniq->GetAllSets())
-            for (const auto& path : set)
-                if (!path.empty())
-                    columns.emplace_back(ctx.NewAtom(node->Pos(), path.front()));
-        const auto& name = std::conditional_t<Distinct, TCoAssumeDistinct, TCoAssumeUnique>::CallableName();
-        return columns.empty() ?
-            ctx.NewCallable(node->Pos(), name, {std::move(node)}):
-            ctx.Builder(node->Pos())
-                .Callable(name)
-                    .Add(0, std::move(node))
-                    .List(1)
-                        .Add(std::move(columns))
-                    .Seal()
-                .Seal()
-                .Build();
-    }
-    return node;
-}
-
-TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx) {
-    auto res = KeepSortedConstraint(node, src.GetConstraint<TSortedConstraintNode>(), ctx);
-    res = KeepUniqueConstraint<true>(std::move(res), src, ctx);
-    res = KeepUniqueConstraint<false>(std::move(res), src, ctx);
-    return res;
-}
-
 TExprNode::TPtr MakeEmptyCollectionWithConstraintsFrom(const TExprNode& node, TExprContext& ctx, TOptimizeContext& optCtx) {
     auto res = ctx.NewCallable(node.Pos(), GetEmptyCollectionName(node.GetTypeAnn()),
         { ExpandType(node.Pos(), *node.GetTypeAnn(), 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 58fa541ae0..62b6b3c1db 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
@@ -2376,8 +2376,8 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
 
     auto settings = ctx.Builder(node->Pos())
         .List()
-            .Atom(0, "Auto")
-            .Atom(1, "Many")
+            .Atom(0, "Auto", TNodeFlags::Default)
+            .Atom(1, "Many", TNodeFlags::Default)
         .Seal()
         .Build();
 
@@ -2411,22 +2411,22 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
                 .Callable("OrderedFlatMap")
                     .Callable(0, "SqueezeToDict")
                         .Add(0, node->HeadPtr())
-                        .Add(1, keyExtractor)
-                        .Add(2, idLambda)
-                        .Add(3, settings)
+                        .Add(1, std::move(keyExtractor))
+                        .Add(2, std::move(idLambda))
+                        .Add(3, std::move(settings))
                     .Seal()
-                    .Add(1, flatten)
+                    .Add(1, std::move(flatten))
                 .Seal()
                 .Build();
         } else {
             sort = ctx.Builder(node->Pos())
-                .Apply(flatten)
+                .Apply(*flatten)
                     .With(0)
                         .Callable("ToDict")
                             .Add(0, node->HeadPtr())
-                            .Add(1, keyExtractor)
-                            .Add(2, idLambda)
-                            .Add(3, settings)
+                            .Add(1, std::move(keyExtractor))
+                            .Add(2, std::move(idLambda))
+                            .Add(3, std::move(settings))
                         .Seal()
                     .Done()
                 .Seal()
@@ -2434,7 +2434,7 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
         }
     }
 
-    return ctx.ReplaceNode(node->Tail().TailPtr(), node->Tail().Head().Head(), std::move(sort));
+    return KeepConstraints(ctx.ReplaceNode(node->Tail().TailPtr(), node->Tail().Head().Head(), std::move(sort)), *node, ctx);
 }
 
 TExprNode::TPtr ExpandIsKeySwitch(const TExprNode::TPtr& node, TExprContext& ctx) {
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index d724c8e617..764d1432d7 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -17,6 +17,33 @@ namespace NYql {
 
 using namespace NNodes;
 
+namespace {
+
+template<bool Distinct>
+TExprNode::TPtr KeepUniqueConstraint(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx) {
+    if (const auto uniq = src.GetConstraint<TUniqueConstraintNodeBase<Distinct>>()) {
+        TExprNode::TListType columns;
+        for (const auto& set : uniq->GetAllSets())
+            for (const auto& path : set)
+                if (!path.empty())
+                    columns.emplace_back(ctx.NewAtom(node->Pos(), path.front()));
+        const auto& name = std::conditional_t<Distinct, TCoAssumeDistinct, TCoAssumeUnique>::CallableName();
+        return columns.empty() ?
+            ctx.NewCallable(node->Pos(), name, {std::move(node)}):
+            ctx.Builder(node->Pos())
+                .Callable(name)
+                    .Add(0, std::move(node))
+                    .List(1)
+                        .Add(std::move(columns))
+                    .Seal()
+                .Seal()
+                .Build();
+    }
+    return node;
+}
+
+}
+
 TExprNode::TPtr MakeBoolNothing(TPositionHandle position, TExprContext& ctx) {
     return ctx.NewCallable(position, "Nothing", {
         ctx.NewCallable(position, "OptionalType", {
@@ -1692,4 +1719,58 @@ bool HasDependsOn(const TExprNode::TPtr& root, const TExprNode::TPtr& arg) {
     return withDependsOn;
 }
 
+TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, TExprContext& ctx) {
+    if (!sorted) {
+        return node;
+    }
+    const auto& constent = sorted->GetContent();
+    return ctx.Builder(node->Pos())
+        .Callable("AssumeSorted")
+            .Add(0, std::move(node))
+            .List(1)
+                .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
+                    size_t index = 0;
+                    for (auto c : constent) {
+                        parent.Callable(index++, "Bool")
+                            .Atom(0, ToString(c.second), TNodeFlags::Default)
+                        .Seal();
+                        if (1U < c.first.front().size())
+                            break;
+                    }
+                    return parent;
+                })
+            .Seal()
+            .Lambda(2)
+                .Param("item")
+                .List()
+                    .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
+                        size_t index = 0;
+                        for (auto c : constent) {
+                            if (c.first.front().empty())
+                                parent.Arg(index++, "item");
+                            else if (1U == c.first.front().size()) {
+                                parent.Callable(index++, "Member")
+                                    .Arg(0, "item")
+                                    .Atom(1, c.first.front().front())
+                                .Seal();
+                            } else {
+                                parent.Callable(index++, "Null").Seal();
+                                break;
+                            }
+                        }
+                        return parent;
+                    })
+                .Seal()
+            .Seal()
+        .Seal()
+        .Build();
+}
+
+TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx) {
+    auto res = KeepSortedConstraint(node, src.GetConstraint<TSortedConstraintNode>(), ctx);
+    res = KeepUniqueConstraint<true>(std::move(res), src, ctx);
+    res = KeepUniqueConstraint<false>(std::move(res), src, ctx);
+    return res;
+}
+
 }
diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
index caa29d0e49..8f95e989a7 100644
--- a/ydb/library/yql/core/yql_opt_utils.h
+++ b/ydb/library/yql/core/yql_opt_utils.h
@@ -134,4 +134,7 @@ bool IsYieldTransparent(const TExprNode::TPtr& root, const TTypeAnnotationContex
 bool IsStrict(const TExprNode::TPtr& node);
 bool HasDependsOn(const TExprNode::TPtr& node, const TExprNode::TPtr& arg);
 
+TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, TExprContext& ctx);
+TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx);
+
 }