YQL-16653 Fix sorted over nested OrderedFlatMap with internal Sort.

4 files changed, 43 insertions, 6 deletions

diff --git a/ydb/library/yql/ast/yql_constraint.cpp b/ydb/library/yql/ast/yql_constraint.cpp
index d2a8242610..e4c36cb847 100644
--- a/ydb/library/yql/ast/yql_constraint.cpp
+++ b/ydb/library/yql/ast/yql_constraint.cpp
@@ -1442,6 +1442,14 @@ TPartOfConstraintNode<TOriginalConstraintNode>::CompleteOnly(TExprContext& ctx)
 }
 
 template<class TOriginalConstraintNode>
+const TPartOfConstraintNode<TOriginalConstraintNode>*
+TPartOfConstraintNode<TOriginalConstraintNode>:: RemoveOriginal(TExprContext& ctx, const TMainConstraint* original) const {
+    TMapType mapping(Mapping_);
+    mapping.erase(original);
+    return mapping.empty() ? nullptr : ctx.MakeConstraint<TPartOfConstraintNode>(std::move(mapping));
+}
+
+template<class TOriginalConstraintNode>
 typename TPartOfConstraintNode<TOriginalConstraintNode>::TMapType
 TPartOfConstraintNode<TOriginalConstraintNode>::GetColumnMapping(const std::string_view& asField) const {
     auto mapping = Mapping_;
diff --git a/ydb/library/yql/ast/yql_constraint.h b/ydb/library/yql/ast/yql_constraint.h
index ed5a9d3dbc..556981cb37 100644
--- a/ydb/library/yql/ast/yql_constraint.h
+++ b/ydb/library/yql/ast/yql_constraint.h
@@ -392,6 +392,7 @@ public:
 
     const TPartOfConstraintNode* ExtractField(TExprContext& ctx, const std::string_view& field) const;
     const TPartOfConstraintNode* CompleteOnly(TExprContext& ctx) const;
+    const TPartOfConstraintNode* RemoveOriginal(TExprContext& ctx, const TMainConstraint* original) const;
 
     static const TPartOfConstraintNode* MakeCommon(const std::vector<const TConstraintSet*>& constraints, TExprContext& ctx);
 
diff --git a/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp b/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
index e2f63831f5..fcdcae03d6 100644
--- a/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
+++ b/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
@@ -360,6 +360,28 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TSortedConstraintNode>(exprRoot, "LazyList", "Sorted(x[asc];z[asc])");
     }
 
+    Y_UNIT_TEST(SortedOverNestedOrderedFlatMapWithSort) {
+        const auto s = R"((
+            (let mr_sink (DataSink 'yt (quote plato)))
+            (let list (AsList
+                (AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'u)))
+                (AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'v)))
+                (AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'w)))
+            ))
+            (let sorted (Sort list (Bool 'True) (lambda '(item) '((Member item 'key) (Member item 'value)))))
+            (let map (OrderedFlatMap sorted (lambda '(row) (OrderedFlatMap (Sort (ListFromRange (Uint8 '0) (Uint8 '5) (Uint8 '1)) (Bool 'True) (lambda '(item) item)) (lambda '(index)
+                (OptionalIf (AggrNotEquals index (Uint8 '3)) (AsStruct '('x (Member row 'key)) '('y (Concat (ToString index) (Member row 'subkey))) '('z (Member row 'value))))
+            )))))
+            (let world (Write! world mr_sink (Key '('table (String 'Output))) (LazyList map) '('('mode 'renew))))
+            (let world (Commit! world mr_sink))
+            (return world)
+        ))";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TSortedConstraintNode>(exprRoot, "LazyList", "Sorted(x[asc];z[asc])");
+    }
+
     Y_UNIT_TEST(TopSort) {
         const auto s = R"((
             (let mr_sink (DataSink 'yt (quote plato)))
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index c6829a121d..3602a28beb 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -951,11 +951,13 @@ private:
     template<class TConstraint, class TInput>
     static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, TExprContext& ctx) {
         if (const auto lambda = handler.GetConstraint<TConstraint>()) {
-            if (const auto original = input.template GetConstraint<typename TConstraint::TMainConstraint>()) {
+            const auto original = input.template GetConstraint<typename TConstraint::TMainConstraint>();
+            if (original) {
                 if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
                     output.AddConstraint(complete);
                 }
-            } else if (const auto part = input.template GetConstraint<TConstraint>()) {
+            }
+            if (const auto part = input.template GetConstraint<TConstraint>()) {
                 auto mapping = lambda->GetColumnMapping();
                 for (auto it = mapping.cbegin(); mapping.cend() != it;) {
                     if (part->GetColumnMapping().contains(it->first))
@@ -967,7 +969,8 @@ private:
                     output.AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
                 }
             } else if constexpr (std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>()) {
-                output.AddConstraint(lambda);
+                if (const auto filtered = lambda->RemoveOriginal(ctx, original))
+                    output.AddConstraint(filtered);
             }
         }
     }
@@ -975,11 +978,13 @@ private:
     template<class TConstraint, bool WideOutput>
     static void GetFromMapLambda(const TExprNode::TPtr& input, TExprContext& ctx) {
         if (const auto lambda = GetConstraintFromLambda<TConstraint, WideOutput>(input->Tail(), ctx)) {
-            if (const auto original = GetDetailed(input->Head().GetConstraint<typename TConstraint::TMainConstraint>(), *input->Head().GetTypeAnn(), ctx)) {
+            const auto original = GetDetailed(input->Head().GetConstraint<typename TConstraint::TMainConstraint>(), *input->Head().GetTypeAnn(), ctx);
+            if (original) {
                 if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
                     input->AddConstraint(complete->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
                 }
-            } else if (const auto part = input->Head().GetConstraint<TConstraint>()) {
+            }
+            if (const auto part = input->Head().GetConstraint<TConstraint>()) {
                 auto mapping = lambda->GetColumnMapping();
                 for (auto it = mapping.cbegin(); mapping.cend() != it;) {
                     if (part->GetColumnMapping().contains(it->first))
@@ -991,7 +996,8 @@ private:
                     input->AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
                 }
             } else if constexpr (std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>()) {
-                input->AddConstraint(lambda);
+                if (const auto filtered = lambda->RemoveOriginal(ctx, original))
+                    input->AddConstraint(filtered);
             }
         }
     }