YQL-16031 Make system column for ORDER BY EXPRESSSION on peephole final stage.

2 files changed, 91 insertions, 93 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 df93de9c43..88ecbf0b88 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -2394,16 +2394,6 @@ TExprNode::TPtr DropReorder(const TExprNode::TPtr& node, TExprContext& ctx) {
     return node;
 }
 
-bool IsExpression(const TExprNode& root, const TExprNode& arg) {
-    if (&root == &arg)
-        return false;
-
-    if (root.IsCallable({"Member", "Nth"}))
-        return IsExpression(root.Head(), arg);
-
-    return true;
-}
-
 template <bool IsTop, bool IsSort>
 TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx) {
     const ui32 ascIndex = node->ChildrenSize() - 2U;
@@ -2596,88 +2586,6 @@ TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx)
         }
     }
 
-    if constexpr (IsTop || IsSort) {
-        if (/*TODO: enable later*/false && ETypeAnnotationKind::Struct == GetSeqItemType(*node->Head().GetTypeAnn()).GetKind()) {
-            std::set<ui32> indexes;
-            if (node->Tail().Tail().IsList())
-                for (auto i = 0U; i < node->Tail().Tail().ChildrenSize(); ++i)
-                    if (IsExpression(*node->Tail().Tail().Child(i), node->Tail().Head().Head()))
-                        indexes.emplace(i);
-
-            if (!indexes.empty() || (!node->Tail().Tail().IsList() && IsExpression(node->Tail().Tail(), node->Tail().Head().Head()))) {
-                YQL_CLOG(DEBUG, Core) << "Make system columns for " << node->Content() << " keys.";
-
-                auto argIn = ctx.NewArgument(node->Tail().Pos(), "row");
-                auto argOut = ctx.NewArgument(node->Tail().Pos(), "row");
-                auto bodyIn = argIn;
-                auto bodyOut = argOut;
-
-                auto selector = node->TailPtr();
-                if (indexes.empty()) {
-                    auto column = ctx.NewAtom(selector->Pos(), "_yql_sys_order_by", TNodeFlags::Default);
-                    bodyIn = ctx.Builder(bodyIn->Pos())
-                        .Callable("AddMember")
-                            .Add(0, std::move(bodyIn))
-                            .Add(1, column)
-                            .Apply(2, node->Tail())
-                                .With(0, argIn)
-                            .Seal()
-                        .Seal().Build();
-                    bodyOut = ctx.Builder(bodyOut->Pos())
-                        .Callable("RemoveMember")
-                            .Add(0, std::move(bodyOut))
-                            .Add(1, column)
-                        .Seal().Build();
-                    selector = ctx.Builder(selector->Pos())
-                        .Lambda()
-                            .Param("row")
-                            .Callable("Member")
-                                .Arg(0, "row")
-                                .Add(1, std::move(column))
-                            .Seal()
-                        .Seal().Build();
-                } else {
-                    auto items = selector->Tail().ChildrenList();
-                    for (const auto index : indexes) {
-                        auto column = ctx.NewAtom(items[index]->Pos(), TString("_yql_sys_order_by_") += ToString(index), TNodeFlags::Default);
-                        bodyIn = ctx.Builder(bodyIn->Pos())
-                            .Callable("AddMember")
-                                .Add(0, std::move(bodyIn))
-                                .Add(1, column)
-                                .ApplyPartial(2, node->Tail().HeadPtr(), std::move(items[index]))
-                                    .With(0, argIn)
-                                .Seal()
-                            .Seal().Build();
-                        bodyOut = ctx.Builder(bodyOut->Pos())
-                            .Callable("RemoveMember")
-                                .Add(0, std::move(bodyOut))
-                                .Add(1, column)
-                            .Seal().Build();
-                        items[index] = ctx.Builder(selector->Pos())
-                            .Callable("Member")
-                                .Add(0, selector->Head().HeadPtr())
-                                .Add(1, std::move(column))
-                            .Seal().Build();
-                    }
-                    selector = ctx.DeepCopyLambda(*selector, ctx.NewList(selector->Tail().Pos(), std::move(items)));
-                }
-
-                auto children = node->ChildrenList();
-                children.back() = std::move(selector);
-                children.front() = ctx.Builder(node->Pos())
-                        .Callable("OrderedMap")
-                            .Add(0, std::move(children.front()))
-                            .Add(1, ctx.NewLambda(node->Tail().Pos(), ctx.NewArguments(node->Tail().Head().Pos(), {std::move(argIn)}), std::move(bodyIn)))
-                        .Seal().Build();
-                return ctx.Builder(node->Pos())
-                        .Callable("OrderedMap")
-                            .Add(0, ctx.ChangeChildren(*node, std::move(children)))
-                            .Add(1, ctx.NewLambda(node->Tail().Pos(), ctx.NewArguments(node->Tail().Head().Pos(), {std::move(argOut)}), std::move(bodyOut)))
-                        .Seal().Build();
-            }
-        }
-    }
-
     return node;
 }
 
@@ -3063,7 +2971,7 @@ TExprNode::TPtr Normalize(const TCoAggregate& node, TExprContext& ctx) {
                     .With(0, arg)
                 .Seal()
                 .Build();
-                
+
             for (auto& [name, idx] : originalIndexes) {
                 nameNodes.emplace_back(ctx.NewAtom(aggTuple.ColumnName().Pos(), name));
                 finishBody.emplace_back(ctx.Builder(traits.FinishHandler().Pos())
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 6d4f26fa67..b0eca4eaf2 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
@@ -4126,8 +4126,98 @@ TExprNode::TPtr OptimizeCombineCore(const TExprNode::TPtr& node, TExprContext& c
     return node;
 }
 
+bool IsExpression(const TExprNode& root, const TExprNode& arg) {
+    if (&root == &arg)
+        return false;
+
+    if (root.IsCallable({"Member", "Nth"}))
+        return IsExpression(root.Head(), arg);
+
+    return true;
+}
+
 template<bool Sort, bool HasCount>
 TExprNode::TPtr OptimizeTopOrSort(const TExprNode::TPtr& node, TExprContext& ctx) {
+    if (ETypeAnnotationKind::Struct == GetSeqItemType(*node->Head().GetTypeAnn()).GetKind()) {
+        std::set<ui32> indexes;
+        if (node->Tail().Tail().IsList())
+            for (auto i = 0U; i < node->Tail().Tail().ChildrenSize(); ++i)
+                if (IsExpression(*node->Tail().Tail().Child(i), node->Tail().Head().Head()))
+                    indexes.emplace(i);
+
+        if (!indexes.empty() || (!node->Tail().Tail().IsList() && IsExpression(node->Tail().Tail(), node->Tail().Head().Head()))) {
+            YQL_CLOG(DEBUG, Core) << "Make system columns for " << node->Content() << " keys.";
+
+            auto argIn = ctx.NewArgument(node->Tail().Pos(), "row");
+            auto argOut = ctx.NewArgument(node->Tail().Pos(), "row");
+            auto bodyIn = argIn;
+            auto bodyOut = argOut;
+
+            auto selector = node->TailPtr();
+            if (indexes.empty()) {
+                auto column = ctx.NewAtom(selector->Pos(), "_yql_sys_order_by", TNodeFlags::Default);
+                bodyIn = ctx.Builder(bodyIn->Pos())
+                    .Callable("AddMember")
+                        .Add(0, std::move(bodyIn))
+                        .Add(1, column)
+                        .Apply(2, node->Tail())
+                            .With(0, argIn)
+                        .Seal()
+                    .Seal().Build();
+                bodyOut = ctx.Builder(bodyOut->Pos())
+                    .Callable("RemoveMember")
+                        .Add(0, std::move(bodyOut))
+                        .Add(1, column)
+                    .Seal().Build();
+                selector = ctx.Builder(selector->Pos())
+                    .Lambda()
+                        .Param("row")
+                        .Callable("Member")
+                            .Arg(0, "row")
+                            .Add(1, std::move(column))
+                        .Seal()
+                    .Seal().Build();
+            } else {
+                auto items = selector->Tail().ChildrenList();
+                for (const auto index : indexes) {
+                    auto column = ctx.NewAtom(items[index]->Pos(), TString("_yql_sys_order_by_") += ToString(index), TNodeFlags::Default);
+                    bodyIn = ctx.Builder(bodyIn->Pos())
+                        .Callable("AddMember")
+                            .Add(0, std::move(bodyIn))
+                            .Add(1, column)
+                            .ApplyPartial(2, node->Tail().HeadPtr(), std::move(items[index]))
+                                .With(0, argIn)
+                            .Seal()
+                        .Seal().Build();
+                    bodyOut = ctx.Builder(bodyOut->Pos())
+                        .Callable("RemoveMember")
+                            .Add(0, std::move(bodyOut))
+                            .Add(1, column)
+                        .Seal().Build();
+                    items[index] = ctx.Builder(selector->Pos())
+                        .Callable("Member")
+                            .Add(0, selector->Head().HeadPtr())
+                            .Add(1, std::move(column))
+                        .Seal().Build();
+                }
+                selector = ctx.DeepCopyLambda(*selector, ctx.NewList(selector->Tail().Pos(), std::move(items)));
+            }
+
+            auto children = node->ChildrenList();
+            children.back() = std::move(selector);
+            children.front() = ctx.Builder(node->Pos())
+                    .Callable("OrderedMap")
+                        .Add(0, std::move(children.front()))
+                        .Add(1, ctx.NewLambda(node->Tail().Pos(), ctx.NewArguments(node->Tail().Head().Pos(), {std::move(argIn)}), std::move(bodyIn)))
+                    .Seal().Build();
+            return ctx.Builder(node->Pos())
+                    .Callable("OrderedMap")
+                        .Add(0, ctx.ChangeChildren(*node, std::move(children)))
+                        .Add(1, ctx.NewLambda(node->Tail().Pos(), ctx.NewArguments(node->Tail().Head().Pos(), {std::move(argOut)}), std::move(bodyOut)))
+                    .Seal().Build();
+        }
+    }
+
     if (const auto& input = node->Head(); input.IsCallable("NarrowMap") && input.Tail().Tail().IsCallable("AsStruct")) {
         TNodeMap<size_t> indexes(input.Tail().Tail().ChildrenSize());
         input.Tail().Tail().ForEachChild([&](const TExprNode& field) {