Fix ToFlow optimizator with multiple usage

ydb/library/yql/core/common_opt/yql_co_flow1.cpp

@@ -1298,91 +1298,6 @@ TExprNode::TPtr OptimizeLookup(const TExprNode::TPtr& node, TExprContext& ctx, T
     return node;
 }
 
-constexpr std::initializer_list<std::string_view> FlowPriority = {
-    "AssumeSorted", "AssumeUnique", "AssumeDistinct",
-    "Map", "OrderedMap", "MapNext",
-    "Filter", "OrderedFilter",
-    "FlatMap", "OrderedFlatMap",
-    "MultiMap", "OrderedMultiMap",
-    "FoldMap", "Fold1Map", "Chain1Map",
-    "Take", "Skip",
-    "TakeWhile", "SkipWhile",
-    "TakeWhileInclusive", "SkipWhileInclusive",
-    "SkipNullMembers", "FilterNullMembers",
-    "SkipNullElements", "FilterNullElements",
-    "Condense", "Condense1",
-    "MapJoinCore", "CommonJoinCore",
-    "CombineCore", "ExtractMembers",
-    "PartitionByKey", "SqueezeToDict"
-};
-
-TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
-    if (!optCtx.IsSingleUsage(node->Head())) {
-        return node;
-    }
-
-    if (node->Head().IsCallable(FlowPriority)) {
-        YQL_CLOG(DEBUG, Core) << "Swap " << node->Content() << " with " << node->Head().Content();
-        return ctx.SwapWithHead(*node);
-    }
-
-    if (node->Head().IsCallable("FromFlow")) {
-        YQL_CLOG(DEBUG, Core) << "Drop " << node->Content() << " with " << node->Head().Content();
-        return node->Head().HeadPtr();
-    }
-
-    if (node->Head().IsCallable("ForwardList")) {
-        YQL_CLOG(DEBUG, Core) << "Drop " << node->Head().Content() << " under " << node->Content();
-        return ctx.ChangeChild(*node, 0U,  node->Head().HeadPtr());
-    }
-
-    if (node->Head().IsCallable("Chopper")) {
-        YQL_CLOG(DEBUG, Core) << "Swap " << node->Head().Content() << " with " << node->Content();
-        auto children = node->Head().ChildrenList();
-        children.front() = ctx.ChangeChildren(*node, {std::move(children.front())});
-        children.back() = ctx.Builder(children.back()->Pos())
-            .Lambda()
-                .Param("key")
-                .Param("flow")
-                .Callable("ToFlow")
-                    .Apply(0, *children.back())
-                        .With(0, "key")
-                        .With(1)
-                            .Callable("FromFlow")
-                                .Arg(0, "flow")
-                            .Seal()
-                        .Done()
-                    .Seal()
-                .Seal()
-            .Seal().Build();
-        return ctx.ChangeChildren(node->Head(), std::move(children));
-    }
-
-    if (node->Head().IsCallable("Switch")) {
-        YQL_CLOG(DEBUG, Core) << "Swap " << node->Head().Content() << " with " << node->Content();
-        auto children = node->Head().ChildrenList();
-        children.front() = ctx.ChangeChildren(*node, {std::move(children.front())});
-        for (auto i = 3U; i < children.size(); ++++i) {
-            children[i] = ctx.Builder(children[i]->Pos())
-                .Lambda()
-                    .Param("flow")
-                    .Callable("ToFlow")
-                        .Apply(0, *children[i])
-                            .With(0)
-                                .Callable("FromFlow")
-                                    .Arg(0, "flow")
-                                .Seal()
-                            .Done()
-                        .Seal()
-                    .Seal()
-                .Seal().Build();
-        }
-        return ctx.ChangeChildren(node->Head(), std::move(children));
-    }
-
-    return node;
-}
-
 template <bool Ordered>
 TExprNode::TPtr OptimizeFlatMap(const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
     const std::conditional_t<Ordered, TCoOrderedFlatMap, TCoFlatMap> self(node);
@@ -1498,8 +1413,6 @@ TExprNode::TPtr OptimizeFlatMap(const TExprNode::TPtr& node, TExprContext& ctx,
 void RegisterCoFlowCallables1(TCallableOptimizerMap& map) {
     using namespace std::placeholders;
 
-    map["ToFlow"] = std::bind(&OptimizeToFlow, _1, _2, _3);
-
     map["FlatMap"] = std::bind(&OptimizeFlatMap<false>, _1, _2, _3);
     map["OrderedFlatMap"] = std::bind(&OptimizeFlatMap<true>, _1, _2, _3);
 

ydb/library/yql/core/common_opt/yql_co_simple1.cpp

@@ -2271,6 +2271,24 @@ TExprNode::TPtr HasNullOverVariant(const TExprNode::TPtr& node, TExprContext& ct
 
 }
 
+constexpr std::initializer_list<std::string_view> FlowPriority = {
+    "AssumeSorted", "AssumeUnique", "AssumeDistinct",
+    "Map", "OrderedMap", "MapNext",
+    "Filter", "OrderedFilter",
+    "FlatMap", "OrderedFlatMap",
+    "MultiMap", "OrderedMultiMap",
+    "FoldMap", "Fold1Map", "Chain1Map",
+    "Take", "Skip",
+    "TakeWhile", "SkipWhile",
+    "TakeWhileInclusive", "SkipWhileInclusive",
+    "SkipNullMembers", "FilterNullMembers",
+    "SkipNullElements", "FilterNullElements",
+    "Condense", "Condense1",
+    "MapJoinCore", "CommonJoinCore",
+    "CombineCore", "ExtractMembers",
+    "PartitionByKey", "SqueezeToDict"
+};
+
 TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext& ctx) {
     if (node->Head().IsCallable("Nothing")) {
         YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
@@ -2297,6 +2315,65 @@ TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext& ctx) {
         return ctx.ChangeChildren(*node, node->Head().ChildrenList());
     }
 
+    if (node->Head().IsCallable(FlowPriority)) {
+        YQL_CLOG(DEBUG, Core) << "Swap " << node->Content() << " with " << node->Head().Content();
+        return ctx.SwapWithHead(*node);
+    }
+
+    if (node->Head().IsCallable("FromFlow")) {
+        YQL_CLOG(DEBUG, Core) << "Drop " << node->Content() << " with " << node->Head().Content();
+        return node->Head().HeadPtr();
+    }
+
+    if (node->Head().IsCallable("ForwardList")) {
+        YQL_CLOG(DEBUG, Core) << "Drop " << node->Head().Content() << " under " << node->Content();
+        return ctx.ChangeChild(*node, 0U,  node->Head().HeadPtr());
+    }
+
+    if (node->Head().IsCallable("Chopper")) {
+        YQL_CLOG(DEBUG, Core) << "Swap " << node->Head().Content() << " with " << node->Content();
+        auto children = node->Head().ChildrenList();
+        children.front() = ctx.ChangeChildren(*node, {std::move(children.front())});
+        children.back() = ctx.Builder(children.back()->Pos())
+            .Lambda()
+                .Param("key")
+                .Param("flow")
+                .Callable("ToFlow")
+                    .Apply(0, *children.back())
+                        .With(0, "key")
+                        .With(1)
+                            .Callable("FromFlow")
+                                .Arg(0, "flow")
+                            .Seal()
+                        .Done()
+                    .Seal()
+                .Seal()
+            .Seal().Build();
+        return ctx.ChangeChildren(node->Head(), std::move(children));
+    }
+
+    if (node->Head().IsCallable("Switch")) {
+        YQL_CLOG(DEBUG, Core) << "Swap " << node->Head().Content() << " with " << node->Content();
+        auto children = node->Head().ChildrenList();
+        children.front() = ctx.ChangeChildren(*node, {std::move(children.front())});
+        for (auto i = 3U; i < children.size(); ++++i) {
+            children[i] = ctx.Builder(children[i]->Pos())
+                .Lambda()
+                    .Param("flow")
+                    .Callable("ToFlow")
+                        .Apply(0, *children[i])
+                            .With(0)
+                                .Callable("FromFlow")
+                                    .Arg(0, "flow")
+                                .Seal()
+                            .Done()
+                        .Seal()
+                    .Seal()
+                .Seal().Build();
+        }
+        return ctx.ChangeChildren(node->Head(), std::move(children));
+    }
+
     return node;
 }
 

ydb/library/yql/tests/s-expressions/yt_native_file/part0/canondata/result.json

@@ -1496,9 +1496,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt1Key2-Debug]": [
         {
-            "checksum": "e697c3541a74a899e6c2bc1d205fab7e",
-            "size": 12260,
-            "uri": "https://{canondata_backend}/1936997/eddc12f46844dcc083aab92a4d7028395471bd1e/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1Key2-Debug_/opt.yql"
+            "checksum": "1c6930d45679382ee4b199b3595abfa5",
+            "size": 12103,
+            "uri": "https://{canondata_backend}/1784826/280894b68d0f84365de02760bedba1323759c046/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1Key2-Debug_/opt.yql"
         },
         {
             "checksum": "d41d8cd98f00b204e9800998ecf8427e",
@@ -1603,9 +1603,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt1Key2-Plan]": [
         {
-            "checksum": "da186e7214b63700e7a3faae84dfd196",
-            "size": 33255,
-            "uri": "https://{canondata_backend}/1936997/eddc12f46844dcc083aab92a4d7028395471bd1e/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1Key2-Plan_/plan.txt"
+            "checksum": "f11d89ba700d9c98735408ba83965cae",
+            "size": 32683,
+            "uri": "https://{canondata_backend}/1784826/280894b68d0f84365de02760bedba1323759c046/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1Key2-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMapOpt1Key2-Results]": [

ydb/library/yql/tests/s-expressions/yt_native_file/part1/canondata/result.json

@@ -1444,9 +1444,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt2Key2-Debug]": [
         {
-            "checksum": "abec116368ed31e5dfbb0138d63c2600",
-            "size": 11570,
-            "uri": "https://{canondata_backend}/1936997/3fb71bce2d9d914e6b4ccf17de5db7fd43bb076a/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt2Key2-Debug_/opt.yql"
+            "checksum": "afae353d31d065d6727addee75e2bc59",
+            "size": 11484,
+            "uri": "https://{canondata_backend}/937458/0f57afbef9a099d0cf1b5ea3bcb5ab98632998ce/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt2Key2-Debug_/opt.yql"
         },
         {
             "checksum": "d41d8cd98f00b204e9800998ecf8427e",
@@ -1551,9 +1551,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt2Key2-Plan]": [
         {
-            "checksum": "2cf3f104a9e25a6cb06a52e034d11e7a",
-            "size": 32089,
-            "uri": "https://{canondata_backend}/1936997/3fb71bce2d9d914e6b4ccf17de5db7fd43bb076a/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt2Key2-Plan_/plan.txt"
+            "checksum": "d2289aaff2b1dcb15b42421a2a2f7078",
+            "size": 31803,
+            "uri": "https://{canondata_backend}/937458/0f57afbef9a099d0cf1b5ea3bcb5ab98632998ce/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt2Key2-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMapOpt2Key2-Results]": [

ydb/library/yql/tests/s-expressions/yt_native_file/part2/canondata/result.json

@@ -969,9 +969,9 @@
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U8-Debug]": [
         {
-            "checksum": "d87c1a23542fc29a1977933752acc149",
-            "size": 53157,
-            "uri": "https://{canondata_backend}/1923547/a0cc9646b39377ed13e6bfcf3d7a99167b242f62/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U8-Debug_/opt.yql"
+            "checksum": "0da4ac3e3cba548f3830955be75158e0",
+            "size": 53154,
+            "uri": "https://{canondata_backend}/1916746/b18a726d33b8d46cddb25e10e1350d7c6bd71311/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U8-Debug_/opt.yql"
         },
         {
             "checksum": "de438775b4de1b6e9e8e2a94418df106",
@@ -1856,9 +1856,9 @@
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U8-Plan]": [
         {
-            "checksum": "6725f33a77d357132f00caec43edb6d6",
-            "size": 154783,
-            "uri": "https://{canondata_backend}/1924537/2317e67e431b5a3d1571c00417bd8e0fbe0d40be/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U8-Plan_/plan.txt"
+            "checksum": "1ad8ccbd613815cef42b307779cacf0a",
+            "size": 153067,
+            "uri": "https://{canondata_backend}/1814674/895c74707cb831a9b51e52c29c51842f99f1f427/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U8-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U8-Results]": [

ydb/library/yql/tests/s-expressions/yt_native_file/part3/canondata/result.json

@@ -2579,9 +2579,9 @@
     ],
     "test.test[EquiJoinMap-JoinMap-Debug]": [
         {
-            "checksum": "06ba73f0eb6818cb88916897bb70f450",
-            "size": 11475,
-            "uri": "https://{canondata_backend}/1937492/20dbbb65997057ee913a7da1643b851dbebff2ca/resource.tar.gz#test.test_EquiJoinMap-JoinMap-Debug_/opt.yql"
+            "checksum": "22454220ef107b51e9b5a8b81dc82bb5",
+            "size": 11355,
+            "uri": "https://{canondata_backend}/1597364/857311e8ec9f08c9a1767079dbcac8e2136948e4/resource.tar.gz#test.test_EquiJoinMap-JoinMap-Debug_/opt.yql"
         },
         {
             "checksum": "0f80b02b64f2aa4a39a9765ecefab0f1",
@@ -2686,9 +2686,9 @@
     ],
     "test.test[EquiJoinMap-JoinMap-Plan]": [
         {
-            "checksum": "f08c4567686936bd71a51cfd8b1a808e",
-            "size": 32293,
-            "uri": "https://{canondata_backend}/1937492/20dbbb65997057ee913a7da1643b851dbebff2ca/resource.tar.gz#test.test_EquiJoinMap-JoinMap-Plan_/plan.txt"
+            "checksum": "9fb035d9c766bc69e1740db85d84ff01",
+            "size": 31721,
+            "uri": "https://{canondata_backend}/1920236/df054e4be86a61a518eb2517405ce4c923bafedc/resource.tar.gz#test.test_EquiJoinMap-JoinMap-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMap-Results]": [
@@ -2700,9 +2700,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt1-Debug]": [
         {
-            "checksum": "e8b8fb8d2b981433063fbae00d64c6cf",
-            "size": 11894,
-            "uri": "https://{canondata_backend}/1937492/20dbbb65997057ee913a7da1643b851dbebff2ca/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1-Debug_/opt.yql"
+            "checksum": "8d55802f14c5e93c2147bfd3d9b49605",
+            "size": 11774,
+            "uri": "https://{canondata_backend}/1775319/da336a036c0a9d0dd564142426be501233a2e3dc/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1-Debug_/opt.yql"
         },
         {
             "checksum": "0f80b02b64f2aa4a39a9765ecefab0f1",
@@ -2807,9 +2807,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt1-Plan]": [
         {
-            "checksum": "36498252b46be0415f5c718403a54baa",
-            "size": 33193,
-            "uri": "https://{canondata_backend}/1937492/20dbbb65997057ee913a7da1643b851dbebff2ca/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1-Plan_/plan.txt"
+            "checksum": "e3d94a46aa264a34bec28c1c37e856e9",
+            "size": 32621,
+            "uri": "https://{canondata_backend}/1936273/38e8909fc0ff224bb75a20a3ee42c359fa526dbd/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt1-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMapOpt1-Results]": [

ydb/library/yql/tests/s-expressions/yt_native_file/part4/canondata/result.json

@@ -1035,9 +1035,9 @@
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U32-Debug]": [
         {
-            "checksum": "07df0647cd2720568de2acd1c2f5ac23",
-            "size": 18434,
-            "uri": "https://{canondata_backend}/1880306/3923e1a600f93542b414a543bc91520fbd3a2554/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U32-Debug_/opt.yql"
+            "checksum": "ca54ab9b5a671343d62a512410c268c1",
+            "size": 18316,
+            "uri": "https://{canondata_backend}/1600758/8ae22e8895aeca97699cbfb33bc7979a3bfb8310/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U32-Debug_/opt.yql"
         },
         {
             "checksum": "b3ad6c26babf882b76373f6084dbf37c",
@@ -1262,9 +1262,9 @@
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U32-Plan]": [
         {
-            "checksum": "f54da3e0fe9734725763ef7f3ee0a8f0",
-            "size": 43842,
-            "uri": "https://{canondata_backend}/1931696/343d3ddbf17ce205e67f4b609bbc3d6820c76625/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U32-Plan_/plan.txt"
+            "checksum": "fd4d3cbbf361c1525902f36a16ecf03e",
+            "size": 43270,
+            "uri": "https://{canondata_backend}/1937429/9ccaeaa71c2a2536036c2fc5be4d363a086f39f9/resource.tar.gz#test.test_EquiJoin-EquiJoinConvertMapI8U32-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoin-EquiJoinConvertMapI8U32-Results]": [
@@ -2055,9 +2055,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt12-Debug]": [
         {
-            "checksum": "33945bdfd2af302b24b1d9df58e02660",
-            "size": 11419,
-            "uri": "https://{canondata_backend}/1936997/7e1c17c15eb7c14c851e4812e1d7a1d1e5ecb4c8/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12-Debug_/opt.yql"
+            "checksum": "3a7e0d7ecf1c6c01b3ce1a587d747111",
+            "size": 11370,
+            "uri": "https://{canondata_backend}/1942415/69fda1cecc1db727564c082823132d3fb504cfc5/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12-Debug_/opt.yql"
         },
         {
             "checksum": "e286e591e319776361f0538fbcb207f2",
@@ -2162,9 +2162,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt12-Plan]": [
         {
-            "checksum": "2abf7a421efb9d665ba0e0e3f74f7dd9",
-            "size": 32783,
-            "uri": "https://{canondata_backend}/1936997/7e1c17c15eb7c14c851e4812e1d7a1d1e5ecb4c8/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12-Plan_/plan.txt"
+            "checksum": "2e2e956784072772e3a835a767c5f1c2",
+            "size": 32497,
+            "uri": "https://{canondata_backend}/212715/be64088d0c950974876d778f6923606e08b13081/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMapOpt12-Results]": [

ydb/library/yql/tests/s-expressions/yt_native_file/part5/canondata/result.json

@@ -2247,9 +2247,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt12Key2-Debug]": [
         {
-            "checksum": "972f7cfb59951585e0de61f6a06f1f4d",
-            "size": 11551,
-            "uri": "https://{canondata_backend}/1937027/9b837c8118f089b56a7bc68568eb0258b2af0bc3/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12Key2-Debug_/opt.yql"
+            "checksum": "0ad6c2319bcd6b396448b9199b8c0e32",
+            "size": 11502,
+            "uri": "https://{canondata_backend}/1936842/d293e0569b2c9d76735a36d71d28bb6d3f0c9fae/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12Key2-Debug_/opt.yql"
         },
         {
             "checksum": "d41d8cd98f00b204e9800998ecf8427e",
@@ -2354,9 +2354,9 @@
     ],
     "test.test[EquiJoinMap-JoinMapOpt12Key2-Plan]": [
         {
-            "checksum": "1985c1fca9e8e1dc20b489ad59f732fe",
-            "size": 32845,
-            "uri": "https://{canondata_backend}/1937027/9b837c8118f089b56a7bc68568eb0258b2af0bc3/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12Key2-Plan_/plan.txt"
+            "checksum": "e2e310111a7a6a0d5278e5cd559579b6",
+            "size": 32559,
+            "uri": "https://{canondata_backend}/1920236/3678687df230bd85b22afa78dc31dfe1815b8f9f/resource.tar.gz#test.test_EquiJoinMap-JoinMapOpt12Key2-Plan_/plan.txt"
         }
     ],
     "test.test[EquiJoinMap-JoinMapOpt12Key2-Results]": [