YQL-16719 Fix keep sorted by tuple element.

3 files changed, 31 insertions, 66 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 b0beeffe08..8ed678f2c4 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -2533,11 +2533,10 @@ TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx)
             }
         }
 
-        if (auto inputConstr = node->Head().GetConstraint<TSortedConstraintNode>()) {
-            if (auto topConstr = node->GetConstraint<TSortedConstraintNode>()) {
+        if (const auto inputConstr = node->Head().GetConstraint<TSortedConstraintNode>()) {
+            if (const auto topConstr = node->GetConstraint<TSortedConstraintNode>()) {
                 if (topConstr->IsPrefixOf(*inputConstr)) {
                     YQL_CLOG(DEBUG, Core) << node->Content() << " over sorted input";
-
                     auto res = ctx.Builder(node->Pos())
                         .Callable("Take")
                             .Add(0, node->HeadPtr())
@@ -2545,11 +2544,8 @@ TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx)
                         .Seal()
                         .Build();
 
-                    if (topConstr->Equals(*inputConstr)) {
-                        return res;
-                    }
-
-                    return KeepSortedConstraint(res, topConstr, ctx);
+                    return topConstr->Equals(*inputConstr) ? res :
+                        KeepSortedConstraint(std::move(res), topConstr, GetSeqItemType(node->GetTypeAnn()), ctx);
                 }
             }
             YQL_CLOG(DEBUG, Core) << node->Content() << " over input with " << *inputConstr;
@@ -2585,8 +2581,9 @@ TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx)
         if (const auto inputConstr = node->Head().GetConstraint<TSortedConstraintNode>()) {
             if (const auto sortConstr = node->GetConstraint<TSortedConstraintNode>()) {
                 if (sortConstr->IsPrefixOf(*inputConstr)) {
-                    YQL_CLOG(DEBUG, Core) << node->Content() << " over sorted input";
-                    return KeepSortedConstraint(node->HeadPtr(), sortConstr, ctx);
+                    YQL_CLOG(DEBUG, Core) << node->Content() << " over sorted input.";
+                    return //TODO: sortConstr->Equals(*inputConstr) ? node->HeadPtr() :
+                        KeepSortedConstraint(node->HeadPtr(), sortConstr, GetSeqItemType(node->GetTypeAnn()), ctx);
                 }
             }
             YQL_CLOG(DEBUG, Core) << node->Content() << " over input with " << *inputConstr;
@@ -2684,50 +2681,6 @@ TExprNode::TPtr OptimizeReorder(const TExprNode::TPtr& node, TExprContext& ctx)
     return node;
 }
 
-void FixSortness(const TExprNode& origNode, TExprNode::TPtr& node, TExprContext& ctx) {
-    if (const auto sorted = origNode.GetConstraint<TSortedConstraintNode>()) {
-        if (const auto simple = sorted->FilterFields(ctx, [](const TPartOfConstraintBase::TPathType& path) { return 1U == path.size(); })) {
-            const auto& content = simple->GetContent();
-            node = ctx.Builder(origNode.Pos())
-                .Callable("Sort")
-                    .Add(0, std::move(node))
-                    .List(1)
-                        .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                            size_t index = 0;
-                            for (auto c : content) {
-                                parent.Callable(index++, "Bool")
-                                        .Atom(0, ToString(c.second), TNodeFlags::Default)
-                                        .Seal();
-                            }
-                            return parent;
-                        })
-                    .Seal()
-                    .Lambda(2)
-                        .Param("item")
-                        .List()
-                            .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
-                                size_t index = 0;
-                                for (auto c : content) {
-                                    if (c.first.empty())
-                                        parent.Arg(0, "item");
-                                    else {
-                                        YQL_ENSURE(c.first.front().size() == 1U, "Just column expected.");
-                                        parent.Callable(index++, "Member")
-                                                .Arg(0, "item")
-                                                .Atom(1, c.first.front().front())
-                                                .Seal();
-                                    }
-                                }
-                                return parent;
-                            })
-                        .Seal()
-                    .Seal()
-                .Seal()
-                .Build();
-        }
-    }
-}
-
 TExprNode::TPtr ConvertSqlInPredicatesPrefixToJoins(const TExprNode::TPtr& flatMap, const TPredicateChain& chain,
     const TExprNode::TPtr& sqlInTail, TExprContext& ctx)
 {
@@ -2773,7 +2726,7 @@ TExprNode::TPtr ConvertSqlInPredicatesPrefixToJoins(const TExprNode::TPtr& flatM
     YQL_CLOG(DEBUG, Core) << "FlatMapOverJoinableSqlInChain of size " << chain.size();
 
     auto eq = BuildEquiJoinForSqlInChain(flatMap, chain, ctx);
-    FixSortness(*flatMap, eq, ctx);
+    eq = MakeSortByConstraint(std::move(eq), flatMap->GetConstraint<TSortedConstraintNode>(), GetSeqItemType(flatMap->GetTypeAnn()), ctx);
 
     auto tail = sqlInTail ? sqlInTail : MakeBool<true>(flatMap->Pos(), ctx);
     return RebuildFlatmapOverPartOfPredicate(flatMap, eq, tail, true, ctx);
@@ -3702,7 +3655,7 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
         auto right = node->ChildPtr(1);
         if (left->GetTypeAnn()->GetKind() != ETypeAnnotationKind::Resource && right->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Resource) {
             return ctx.ChangeChild(*node, 0, ctx.NewCallable(node->Pos(), "TypeHandle", {left}));
-        } 
+        }
         if (left->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Resource && right->GetTypeAnn()->GetKind() != ETypeAnnotationKind::Resource) {
             return ctx.ChangeChild(*node, 1, ctx.NewCallable(node->Pos(), "TypeHandle", {right}));
         }
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index f4406b1ed8..094d3f765e 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -71,13 +71,15 @@ TExprNode::TPtr KeepChoppedConstraint(TExprNode::TPtr node, const TExprNode& src
     return node;
 }
 
-TExprNodeBuilder GetterBuilder(TExprNodeBuilder parent, ui32 index, TPartOfConstraintBase::TPathType path) {
+TExprNodeBuilder GetterBuilder(TExprNodeBuilder parent, ui32 index, const TTypeAnnotationNode& type, TPartOfConstraintBase::TPathType path) {
     if (path.empty())
         return parent.Arg(index, "item");
 
     const auto& name = path.back();
     path.pop_back();
-    return GetterBuilder(parent.Callable(index, "Member"), 0U, path).Atom(1, name).Seal();
+    const auto parentType = TPartOfConstraintBase::GetSubTypeByPath(path, type);
+    YQL_ENSURE(parentType, "Wrong path '" << path << "' to component of: " << type);
+    return GetterBuilder(parent.Callable(index, ETypeAnnotationKind::Struct == parentType->GetKind() ? "Member" : "Nth"), 0U, type, path).Atom(1, name).Seal();
 }
 
 }
@@ -225,7 +227,7 @@ bool IsRenameFlatMapWithMapping(const NNodes::TCoFlatMapBase& node, TExprNode::T
     return true;
 }
 
-// Check if the flat map is a simple rename flat map or a flatmap that also computes some 
+// Check if the flat map is a simple rename flat map or a flatmap that also computes some
 // values in 1-1 fashion
 bool IsRenameOrApplyFlatMapWithMapping(const NNodes::TCoFlatMapBase& node, TExprNode::TPtr& structNode,
     THashMap<TString, TString> & mapping, TSet<TString> & apply) {
@@ -1853,13 +1855,14 @@ 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) {
+template<bool Assume>
+TExprNode::TPtr MakeSortConstraintImpl(TExprNode::TPtr node, const TSortedConstraintNode* sorted, const TTypeAnnotationNode* rowType, TExprContext& ctx) {
+    if (!(sorted && rowType))
         return node;
-    }
+
     const auto& constent = sorted->GetContent();
     return ctx.Builder(node->Pos())
-        .Callable("AssumeSorted")
+        .Callable(Assume ? "AssumeSorted" : "Sort")
             .Add(0, std::move(node))
             .List(1)
                 .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
@@ -1878,7 +1881,7 @@ TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstrai
                     .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
                         size_t index = 0;
                         for (const auto& c : constent)
-                            GetterBuilder(parent, index++, c.first.front());
+                            GetterBuilder(parent, index++, *rowType, c.first.front());
                         return parent;
                     })
                 .Seal()
@@ -1887,8 +1890,16 @@ TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstrai
         .Build();
 }
 
+TExprNode::TPtr KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, const TTypeAnnotationNode* rowType, TExprContext& ctx) {
+    return MakeSortConstraintImpl<true>(std::move(node), sorted, rowType, ctx);
+}
+
+TExprNode::TPtr MakeSortByConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, const TTypeAnnotationNode* rowType, TExprContext& ctx) {
+    return MakeSortConstraintImpl<false>(std::move(node), sorted, rowType, ctx);
+}
+
 TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx) {
-    auto res = KeepSortedConstraint(node, src.GetConstraint<TSortedConstraintNode>(), ctx);
+    auto res = KeepSortedConstraint(node, src.GetConstraint<TSortedConstraintNode>(), GetSeqItemType(src.GetTypeAnn()), ctx);
     res = KeepChoppedConstraint(std::move(res), src, ctx);
     res = KeepUniqueConstraint<true>(std::move(res), src, ctx);
     res = KeepUniqueConstraint<false>(std::move(res), src, ctx);
diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
index 2501efbbd0..34529d44a5 100644
--- a/ydb/library/yql/core/yql_opt_utils.h
+++ b/ydb/library/yql/core/yql_opt_utils.h
@@ -141,7 +141,8 @@ 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 KeepSortedConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, const TTypeAnnotationNode* rowType, TExprContext& ctx);
+TExprNode::TPtr MakeSortByConstraint(TExprNode::TPtr node, const TSortedConstraintNode* sorted, const TTypeAnnotationNode* rowType, TExprContext& ctx);
 TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExprContext& ctx);
 
 }