YQL-16254 Pass Sorted/Chopped constraints throught nested multiplexing OrderedFlatMap.

2 files changed, 52 insertions, 8 deletions

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 57a38854ac..e2f63831f5 100644
--- a/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
+++ b/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
@@ -315,6 +315,51 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TSortedConstraintNode>(exprRoot, "Collect", "Sorted(z[desc];x[asc])");
     }
 
+    Y_UNIT_TEST(SortedOverOrderedMultiMap) {
+        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 (OrderedMultiMap sorted (lambda '(row)
+                (AsStruct '('x (Member row 'key)) '('y (Member row 'subkey)) '('z (Member row 'value)))
+                (AsStruct '('x (Member row 'key)) '('y (String 'subkey_stub)) '('z (Member row 'value)))
+            )))
+            (let world (Write! world mr_sink (Key '('table (String 'Output))) map '('('mode 'renew))))
+            (let world (Commit! world mr_sink))
+            (return world)
+        ))";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TSortedConstraintNode>(exprRoot, "OrderedMultiMap", "Sorted(x[asc];z[asc])");
+    }
+
+    Y_UNIT_TEST(SortedOverNestedOrderedFlatMap) {
+        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 (ListFromRange (Uint8 '0) (Uint8 '5) (Uint8 '1)) (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 05473c8dd9..fc5dacc8bb 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -936,8 +936,7 @@ private:
                 if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
                     output.AddConstraint(complete);
                 }
-            }
-            if (const auto part = input.template GetConstraint<TConstraint>()) {
+            } else 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))
@@ -948,6 +947,8 @@ private:
                 if (!mapping.empty()) {
                     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);
             }
         }
     }
@@ -957,13 +958,9 @@ private:
         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)) {
                 if (const auto complete = TConstraint::MakeComplete(ctx, lambda->GetColumnMapping(), original)) {
-                    if constexpr (std::is_same<typename TConstraint::TMainConstraint, TUniqueConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TDistinctConstraintNode>())
-                        input->AddConstraint(complete->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
-                    else
-                        input->AddConstraint(complete);
+                    input->AddConstraint(complete->GetSimplifiedForType(*input->GetTypeAnn(), ctx));
                 }
-            }
-            if (const auto part = input->Head().GetConstraint<TConstraint>()) {
+            } else 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))
@@ -974,6 +971,8 @@ private:
                 if (!mapping.empty()) {
                     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);
             }
         }
     }