Fix StablePickle constraints for complex types

commit_hash:3d15defe18fb19d2185bd1798b181d6f30277d68

4 files changed, 102 insertions, 1 deletions

diff --git a/yql/essentials/core/ut/yql_expr_constraint_ut.cpp b/yql/essentials/core/ut/yql_expr_constraint_ut.cpp
index e9b3efd69f1..45d9f6285de 100644
--- a/yql/essentials/core/ut/yql_expr_constraint_ut.cpp
+++ b/yql/essentials/core/ut/yql_expr_constraint_ut.cpp
@@ -3316,6 +3316,52 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
         CheckConstraint<TDistinctConstraintNode>(exprRoot, "PartitionsByKeys", "Distinct((data,group0))");
         CheckConstraint<TUniqueConstraintNode>(exprRoot, "PartitionsByKeys", "Unique((data,group0))");
     }
+
+    Y_UNIT_TEST(StablePickleOfComplexUnique) {
+        const TStringBuf s = R"(
+(
+    (let config (DataSource 'config))
+    (let res_sink (DataSink 'result))
+
+    (let list (AsList
+        (AsStruct '('key (Uint32 '1)) '('value (Uint32 '2)))
+        (AsStruct '('key (Uint32 '2)) '('value (Uint32 '3)))
+    ))
+
+    (let res (Aggregate list '('key 'value) '() '()))
+    (let res (Map res (lambda '(item)
+        (AsStruct
+            '('composite (AsStruct
+                '('k (Member item 'key))
+                '('v (Member item 'value))
+            ))
+            '('key (Member item 'key))
+            '('value (Member item 'value))
+        )
+    )))
+    (let res (FlatMap res (lambda '(item)
+        (Just (AsStruct
+            '('packed (StablePickle (Member item 'composite)))
+            '('composite (Member item 'composite))
+            '('key (Member item 'key))
+            '('value (Member item 'value))
+        ))
+    )))
+    (let world (Write! world res_sink (Key) res '('('type))))
+    (let world (Commit! world res_sink))
+    (return world)
+)
+        )";
+
+        TExprContext exprCtx;
+        const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+        CheckConstraint<TDistinctConstraintNode>(exprRoot, "StablePickle", "");
+        CheckConstraint<TUniqueConstraintNode>(exprRoot, "StablePickle", "");
+        CheckConstraint<TDistinctConstraintNode>(exprRoot, "Map", "Distinct(({composite/k,key},{composite/v,value}))");
+        CheckConstraint<TUniqueConstraintNode>(exprRoot, "Map", "Unique(({composite/k,key},{composite/v,value}))");
+        CheckConstraint<TDistinctConstraintNode>(exprRoot, "FlatMap", "Distinct(({composite/k,key},{composite/v,value}))");
+        CheckConstraint<TUniqueConstraintNode>(exprRoot, "FlatMap", "Unique(({composite/k,key},{composite/v,value}))");
+    }
 }
 
 } // namespace NYql
diff --git a/yql/essentials/core/yql_expr_constraint.cpp b/yql/essentials/core/yql_expr_constraint.cpp
index 49997179f55..7921e31f97e 100644
--- a/yql/essentials/core/yql_expr_constraint.cpp
+++ b/yql/essentials/core/yql_expr_constraint.cpp
@@ -248,7 +248,7 @@ public:
         Functions["BlockMergeFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
         Functions["BlockMergeManyFinalizeHashed"] = &TCallableConstraintTransformer::AggregateWrap<true>;
         Functions["MultiHoppingCore"] = &TCallableConstraintTransformer::MultiHoppingCoreWrap;
-        Functions["StablePickle"] = &TCallableConstraintTransformer::FromFirst<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
+        Functions["StablePickle"] = &TCallableConstraintTransformer::PickleWrap;
         Functions["Unpickle"] = &TCallableConstraintTransformer::FromSecond<TUniqueConstraintNode, TPartOfUniqueConstraintNode, TDistinctConstraintNode, TPartOfDistinctConstraintNode, TPartOfChoppedConstraintNode, TVarIndexConstraintNode>;
     }
 
@@ -404,6 +404,14 @@ private:
         return FromFirst<TEmptyConstraintNode, TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
     }
 
+    TStatus PickleWrap(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) const {
+        if (IsDataOrOptionalOfDataOrPg(input->Head().GetTypeAnn())) {
+            TApplyConstraintFromInput<0, TPartOfChoppedConstraintNode, TPartOfUniqueConstraintNode, TPartOfDistinctConstraintNode>::Do(input);
+        }
+
+        return FromFirst<TUniqueConstraintNode, TDistinctConstraintNode, TVarIndexConstraintNode>(input, output, ctx);
+    }
+
     TStatus AssumeConstraintsWrap(const TExprNode::TPtr& input, TExprNode::TPtr& /*output*/, TExprContext& ctx) const {
         TConstraintSet set;
         try {
diff --git a/yql/essentials/tests/s-expressions/minirun/part9/canondata/result.json b/yql/essentials/tests/s-expressions/minirun/part9/canondata/result.json
index 352f3fa6c6d..f839ab81fd5 100644
--- a/yql/essentials/tests/s-expressions/minirun/part9/canondata/result.json
+++ b/yql/essentials/tests/s-expressions/minirun/part9/canondata/result.json
@@ -321,6 +321,20 @@
             "uri": "https://{canondata_backend}/1937367/e77ba21bdee3073d5ac816e83de0c7af4a7daf05/resource.tar.gz#test.test_Lib-Agg_List_Compare_By_Opts-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[Optimizers-EmptyOverListWithComplexUniq-default.txt-Debug]": [
+        {
+            "checksum": "fc513f9ab77aa25c63bb1a43467fe29c",
+            "size": 396,
+            "uri": "https://{canondata_backend}/995452/07ee5aca68bd02fba1bfa8356eaea86f0e5a7870/resource.tar.gz#test.test_Optimizers-EmptyOverListWithComplexUniq-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[Optimizers-EmptyOverListWithComplexUniq-default.txt-Results]": [
+        {
+            "checksum": "8a557f375bf97e6adf4e8c1e40cd18c0",
+            "size": 2123,
+            "uri": "https://{canondata_backend}/995452/07ee5aca68bd02fba1bfa8356eaea86f0e5a7870/resource.tar.gz#test.test_Optimizers-EmptyOverListWithComplexUniq-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[Optimizers-Fold1OneItemTuple-default.txt-Debug]": [
         {
             "checksum": "d92be86870a01fdf575d08165de8f8a0",
diff --git a/yql/essentials/tests/s-expressions/suites/Optimizers/EmptyOverListWithComplexUniq.yqls b/yql/essentials/tests/s-expressions/suites/Optimizers/EmptyOverListWithComplexUniq.yqls
new file mode 100644
index 00000000000..601fc0d5489
--- /dev/null
+++ b/yql/essentials/tests/s-expressions/suites/Optimizers/EmptyOverListWithComplexUniq.yqls
@@ -0,0 +1,33 @@
+(
+(let config (DataSource 'config))
+(let res_sink (DataSink 'result))
+
+(let list (AsList
+    (AsStruct '('key (Uint32 '1)) '('value (Uint32 '2)))
+    (AsStruct '('key (Uint32 '2)) '('value (Uint32 '3)))
+))
+
+(let res (Aggregate list '('key 'value) '() '()))
+(let res (Map res (lambda '(item)
+    (AsStruct
+        '('composite (AsStruct
+            '('k (Member item 'key))
+            '('v (Member item 'value))
+        ))
+        '('key (Member item 'key))
+        '('value (Member item 'value))
+    )
+)))
+(let res (Take res (Uint64 '0)))
+(let res (FlatMap res (lambda '(item)
+    (Just (AsStruct
+        '('packed (StablePickle (Member item 'composite)))
+        '('composite (Member item 'composite))
+        '('key (Member item 'key))
+        '('value (Member item 'value))
+    ))
+)))
+(let world (Write! world res_sink (Key) res '('('type))))
+(let world (Commit! world res_sink))
+(return world)
+)