Fix Sort by StablePickle case. Add tests. (#2489)

4 files changed, 84 insertions, 19 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 f3421a50fb6..c5c3449b056 100644
--- a/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
+++ b/ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
@@ -95,6 +95,25 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TSortedConstraintNode>(exprRoot, "Sort", "Sorted(key[asc])");
     }
 
+    Y_UNIT_TEST(SortByStablePickle) {
+        const auto s = R"((
+            (let mr_sink (DataSink 'yt (quote plato)))
+            (let list (AsList
+                (AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'v)))
+                (AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'v)))
+                (AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'v)))
+            ))
+            (let sorted (Sort list '((Bool 'False) (Bool 'True)) (lambda '(item) '((Member item 'key) (StablePickle (Member item 'subkey))))))
+            (let world (Write! world mr_sink (Key '('table (String 'Output))) sorted '('('mode 'renew))))
+            (let world (Commit! world mr_sink))
+            (return world)
+        ))";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TSortedConstraintNode>(exprRoot, "Sort", "");
+    }
+
     Y_UNIT_TEST(SortByTranspentIfPresent) {
         const auto s = R"((
             (let mr_sink (DataSink 'yt (quote plato)))
@@ -1489,6 +1508,34 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TDistinctConstraintNode>(exprRoot, "Skip", "Distinct((key,subkey))");
     }
 
+    Y_UNIT_TEST(PartitionsByKeysWithCondense1AndStablePickle) {
+        const auto s = R"(
+(
+    (let mr_sink (DataSink 'yt (quote plato)))
+    (let list (AsList
+        (AsStruct '('key (Just (String '4))) '('subkey (Just (String 'c))) '('value (Just (String 'x))))
+        (AsStruct '('key (Just (String '1))) '('subkey (Just (String 'b))) '('value (Just (String 'y))))
+        (AsStruct '('key (Just (String '4))) '('subkey (Just (String 'b))) '('value (Just (String 'z))))
+    ))
+    (let extractor (lambda '(item) '((Member item 'key) (StablePickle(Member item 'subkey)))))
+    (let aggr (PartitionsByKeys list extractor (Void) (Void)
+        (lambda '(stream) (Condense1 stream (lambda '(row) row)
+            (lambda '(row state) (IsKeySwitch row state extractor extractor))
+            (lambda '(row state) (AsStruct '('key (Member row 'key)) '('subkey (Member row 'subkey)) '('value (Coalesce (Member row 'value) (Member state 'value)))))
+        ))
+    ))
+    (let world (Write! world mr_sink (Key '('table (String 'Output))) (Skip aggr (Uint64 '1)) '('('mode 'renew))))
+    (let world (Commit! world mr_sink))
+    (return world)
+)
+    )";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TUniqueConstraintNode>(exprRoot, "Skip", "Unique((key,subkey))");
+        CheckConstraint<TDistinctConstraintNode>(exprRoot, "Skip", "Distinct((key,subkey))");
+    }
+
     Y_UNIT_TEST(PartitionsByKeysWithCondense1WithSingleItemTupleKey) {
         const auto s = R"(
 (
diff --git a/ydb/library/yql/core/yql_expr_constraint.cpp b/ydb/library/yql/core/yql_expr_constraint.cpp
index 294fde2765a..f3032f33d3f 100644
--- a/ydb/library/yql/core/yql_expr_constraint.cpp
+++ b/ydb/library/yql/core/yql_expr_constraint.cpp
@@ -781,7 +781,7 @@ private:
                 columns.resize(size, std::make_pair(TPartOfConstraintBase::TPathType(), columns.back().second));
                 auto it = columns.begin();
                 for (auto i = 0U; i < size; ++i) {
-                    if (auto path = GetPathToKey(*keySelectorBody.Child(i), keySelectorArg)) {
+                    if (auto path = GetPathToKey<true>(*keySelectorBody.Child(i), keySelectorArg)) {
                         if (set.insert(*path).second)
                             it++->first = std::move(*path);
                         else if (columns.cend() != it)
@@ -792,7 +792,7 @@ private:
                 }
             } else
                 return {};
-        else if (auto path = GetPathToKey(keySelectorBody, keySelectorArg))
+        else if (auto path = GetPathToKey<true>(keySelectorBody, keySelectorArg))
             if (columns.size() == 1U)
                 columns.front().first = std::move(*path);
             else
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index 53f38546051..0d677f7b2be 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -2005,29 +2005,36 @@ void OptimizeSubsetFieldsForNodeWithMultiUsage(const TExprNode::TPtr& node, cons
 }
 
 
+template<bool Ordered>
 std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(const TExprNode& body, const TExprNode::TChildrenType& args) {
     if (body.IsArgument()) {
         for (auto i = 0U; i < args.size(); ++i)
             if (&body == args[i].Get())
                 return std::make_pair(TPartOfConstraintBase::TPathType(), i);
     } else if (body.IsCallable({"Member","Nth"})) {
-        if (auto path = GetPathToKey(body.Head(), args)) {
+        if (auto path = GetPathToKey<Ordered>(body.Head(), args)) {
             path->first.emplace_back(body.Tail().Content());
             return path;
         } else if (const auto& head = SkipCallables(body.Head(), {"CastStruct","FilterMembers"}); head.IsCallable("AsStruct") && body.IsCallable("Member")) {
-            return GetPathToKey(GetLiteralStructMember(head, body.Tail()), args);
+            return GetPathToKey<Ordered>(GetLiteralStructMember(head, body.Tail()), args);
         } else if (body.IsCallable("Nth") && body.Head().IsList()) {
-            return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), args);
+            return GetPathToKey<Ordered>(*body.Head().Child(FromString<ui32>(body.Tail().Content())), args);
         } else if (body.IsCallable({"CastStruct","FilterMembers"}))  {
-            return GetPathToKey(body.Head(), args);
+            return GetPathToKey<Ordered>(body.Head(), args);
+        }
+    } else if constexpr (!Ordered) {
+        if (body.IsCallable("StablePickle")) {
+            return GetPathToKey<Ordered>(body.Head(), args);
         }
-    } else if (body.IsCallable("StablePickle")) {
-        return GetPathToKey(body.Head(), args);
     }
 
     return std::nullopt;
 }
 
+template std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey<true>(const TExprNode& body, const TExprNode::TChildrenType& args);
+template std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey<false>(const TExprNode& body, const TExprNode::TChildrenType& args);
+
+template<bool Ordered>
 std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& body, const TExprNode& arg) {
     if (&body == &arg)
         return TPartOfConstraintBase::TPathType();
@@ -2040,11 +2047,11 @@ std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& bo
     }
 
     if (body.IsCallable({"CastStruct","FilterMembers","Just","Unwrap"}))
-        return GetPathToKey(body.Head(), arg);
+        return GetPathToKey<Ordered>(body.Head(), arg);
     if (body.IsCallable("Member") && body.Head().IsCallable("AsStruct"))
-        return GetPathToKey(GetLiteralStructMember(body.Head(), body.Tail()), arg);
+        return GetPathToKey<Ordered>(GetLiteralStructMember(body.Head(), body.Tail()), arg);
     if (body.IsCallable("Nth") && body.Head().IsList())
-        return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), arg);
+        return GetPathToKey<Ordered>(*body.Head().Child(FromString<ui32>(body.Tail().Content())), arg);
     if (body.IsList() && 1U == body.ChildrenSize() && body.Head().IsCallable("Nth") && body.Head().Tail().IsAtom("0") &&
         1U == RemoveOptionality(*body.Head().Head().GetTypeAnn()).Cast<TTupleExprType>()->GetSize())
         // Especialy for "Extract single item tuple from Condense1" optimizer.
@@ -2053,31 +2060,39 @@ std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& bo
         body.Head().Head().Content() == body.Head().Tail().Tail().Content() &&
         1U == RemoveOptionality(*body.Head().Tail().Head().GetTypeAnn()).Cast<TStructExprType>()->GetSize())
         // Especialy for "Extract single item struct from Condense1" optimizer.
-        return GetPathToKey(body.Head().Tail().Head(), arg);
+        return GetPathToKey<Ordered>(body.Head().Tail().Head(), arg);
     if (IsTransparentIfPresent(body) && &body.Head() == &arg)
-        return GetPathToKey(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
-    if (body.IsCallable("StablePickle"))
-        return GetPathToKey(body.Head(), arg);
+        return GetPathToKey<Ordered>(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
+    if constexpr (!Ordered)
+        if (body.IsCallable("StablePickle"))
+            return GetPathToKey<Ordered>(body.Head(), arg);
 
     return std::nullopt;
 }
 
+template<bool Ordered>
 TPartOfConstraintBase::TSetType GetPathsToKeys(const TExprNode& body, const TExprNode& arg) {
     TPartOfConstraintBase::TSetType keys;
     if (body.IsList()) {
         if (const auto size = body.ChildrenSize()) {
             keys.reserve(size);
             for (auto i = 0U; i < size; ++i)
-                if (auto path = GetPathToKey(*body.Child(i), arg))
+                if (auto path = GetPathToKey<Ordered>(*body.Child(i), arg))
                     keys.insert_unique(std::move(*path));
         }
-    } else if (body.IsCallable("StablePickle")) {
-        return GetPathsToKeys(body.Head(), arg);
-    } else if (auto path = GetPathToKey(body, arg)) {
+    } else if constexpr (!Ordered) {
+        if (body.IsCallable("StablePickle")) {
+            return GetPathsToKeys<Ordered>(body.Head(), arg);
+        }
+    }
+    if (auto path = GetPathToKey<Ordered>(body, arg)) {
         keys.insert_unique(std::move(*path));
     }
 
     return keys;
 }
 
+template TPartOfConstraintBase::TSetType GetPathsToKeys<true>(const TExprNode& body, const TExprNode& arg);
+template TPartOfConstraintBase::TSetType GetPathsToKeys<false>(const TExprNode& body, const TExprNode& arg);
+
 }
diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
index 2980de885f1..03cbdc352f2 100644
--- a/ydb/library/yql/core/yql_opt_utils.h
+++ b/ydb/library/yql/core/yql_opt_utils.h
@@ -149,9 +149,12 @@ void OptimizeSubsetFieldsForNodeWithMultiUsage(const TExprNode::TPtr& node, cons
     TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx,
     std::function<TExprNode::TPtr(const TExprNode::TPtr&, const TExprNode::TPtr&, const TParentsMap&, TExprContext&)> handler);
 
+template<bool Ordered = false>
 std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& body, const TExprNode& arg);
+template<bool Ordered = false>
 std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(const TExprNode& body, const TExprNode::TChildrenType& args);
 
+template<bool Ordered = false>
 TPartOfConstraintBase::TSetType GetPathsToKeys(const TExprNode& body, const TExprNode& arg);
 
 }