more tests YQL-19494

commit_hash:5710def46c1ae0647d455212ec5f8fb10e598a8f

13 files changed, 271 insertions, 13 deletions

diff --git a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
index 1361eb1061..d9107d0cb6 100644
--- a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -3269,12 +3269,66 @@ TExprNode::TPtr OptimizeMap(const TExprNode::TPtr& node, TExprContext& ctx) {
     return node;
 }
 
+TExprNode::TPtr MakeWideTableSource(const TExprNode& tableSource, TExprContext& ctx, TVector<TString>* narrowMapColumns = nullptr) {
+    // TODO check wide limit
+    if (tableSource.GetTypeAnn()->GetKind() != ETypeAnnotationKind::List) {
+        return nullptr;
+    }
+
+    YQL_CLOG(DEBUG, CorePeepHole) << "Generate WideTableSource";
+    auto structType = tableSource.GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
+    TVector<TString> columns;
+    for (const auto& item : structType->GetItems()) {
+        columns.push_back(TString(item->GetName()));
+    }
+
+    auto innerFlow = ctx.NewCallable(tableSource.Pos(), "ToFlow", { tableSource.HeadPtr() });
+    auto expandMap = MakeExpandMap(tableSource.Pos(), columns, innerFlow, ctx);
+    auto source = ctx.NewCallable(tableSource.Pos(), "WideTableSource", { expandMap });
+    if (narrowMapColumns) {
+        *narrowMapColumns = columns;
+        return source;
+    } else {
+        return MakeNarrowMap(tableSource.Pos(), columns, source, ctx);
+    }
+}
+
+TExprNode::TPtr OptimizeExtend(const TExprNode::TPtr& node, TExprContext& ctx) {
+    if (node->ChildrenSize() > 0 && AllOf(node->Children(), [](const auto& x) { return x->IsCallable("TableSource");})) {
+        TExprNodeList wideChildren;
+        bool allWide = true;
+        TVector<TString> columns;
+        for (const auto& x : node->Children()) {
+            if (auto wide = MakeWideTableSource(*x, ctx, &columns)) {
+                wideChildren.push_back(wide);
+            } else {
+                allWide = false;
+                break;
+            }
+        }
+
+        if (allWide) {
+            return ctx.NewCallable(node->Pos(), "Collect", {
+                MakeNarrowMap(node->Pos(), columns, ctx.NewCallable(node->Pos(), "Extend", std::move(wideChildren)), ctx)
+            });
+        }
+    }
+
+    return node;
+}
+
 TExprNode::TPtr OptimizeSkip(const TExprNode::TPtr& node, TExprContext& ctx) {
     if (const auto& input = node->Head(); input.IsCallable({"Map", "OrderedMap", "ExpandMap", "WideMap", "NarrowMap"})) {
         YQL_CLOG(DEBUG, CorePeepHole) << "Swap " << node->Content() << " with " << input.Content();
         return ctx.SwapWithHead(*node);
     }
 
+    if (node->Head().IsCallable("TableSource")) {
+        if (auto wide = MakeWideTableSource(node->Head(), ctx)) {
+            return ctx.NewCallable(node->Pos(), "Collect", { ctx.ChangeChild(*node, 0, std::move(wide)) });
+        }
+    }
+
     return node;
 }
 
@@ -4766,6 +4820,12 @@ TExprNode::TPtr OptimizeTopOrSort(const TExprNode::TPtr& node, TExprContext& ctx
         }
     }
 
+    if (node->Head().IsCallable("TableSource")) {
+        if (auto wide = MakeWideTableSource(node->Head(), ctx)) {
+            return ctx.NewCallable(node->Pos(), "Collect", { ctx.ChangeChild(*node, 0, std::move(wide)) });
+        }
+    }
+
     return node;
 }
 
@@ -8328,19 +8388,8 @@ TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext& ctx) {
 
 
     if (node->Head().IsCallable("TableSource")) {
-        if (node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::List) {
-            // TODO check wide limit
-            YQL_CLOG(DEBUG, CorePeepHole) << "Generate WideTableSource";
-            auto structType = node->Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
-            TVector<TString> columns;
-            for (const auto& item : structType->GetItems()) {
-                columns.push_back(TString(item->GetName()));
-            }
-
-            auto innerFlow = ctx.NewCallable(node->Pos(), "ToFlow", { node->Head().HeadPtr() });
-            auto expandMap = MakeExpandMap(node->Pos(), columns, innerFlow, ctx);
-            auto source = ctx.NewCallable(node->Pos(), "WideTableSource", { expandMap });
-            return MakeNarrowMap(node->Pos(), columns, source, ctx);
+        if (auto wide = MakeWideTableSource(node->Head(), ctx)) {
+            return wide;
         }
     }
 
@@ -8582,6 +8631,7 @@ struct TPeepHoleRules {
     };
 
     const TPeepHoleOptimizerMap FinalStageRules = {
+        {"Extend", &OptimizeExtend},
         {"Take", &OptimizeTake},
         {"Skip", &OptimizeSkip},
         {"GroupByKey", &PeepHoleConvertGroupBySingleKey},
diff --git a/yql/essentials/tests/sql/minirun/part1/canondata/result.json b/yql/essentials/tests/sql/minirun/part1/canondata/result.json
index e9313a184a..ff0d370a35 100644
--- a/yql/essentials/tests/sql/minirun/part1/canondata/result.json
+++ b/yql/essentials/tests/sql/minirun/part1/canondata/result.json
@@ -251,6 +251,27 @@
             "uri": "https://{canondata_backend}/1881367/d69a84f4c03a928385a2fa1e80f00c900ecc2329/resource.tar.gz#test.test_blocks-empty_as_table-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[blocks-multimap-default.txt-Debug]": [
+        {
+            "checksum": "6486282c5559b4e925dd3412fdfdaa22",
+            "size": 367,
+            "uri": "https://{canondata_backend}/1817427/36dfa08eea10bb99bfc47f3f1c2f11cd961db2f4/resource.tar.gz#test.test_blocks-multimap-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[blocks-multimap-default.txt-Peephole]": [
+        {
+            "checksum": "8ffd922a67911b41b53ca0d75214de42",
+            "size": 460,
+            "uri": "https://{canondata_backend}/1817427/36dfa08eea10bb99bfc47f3f1c2f11cd961db2f4/resource.tar.gz#test.test_blocks-multimap-default.txt-Peephole_/opt.yql"
+        }
+    ],
+    "test.test[blocks-multimap-default.txt-Results]": [
+        {
+            "checksum": "03f400ad58adee187aa5d3c45f42965e",
+            "size": 1282,
+            "uri": "https://{canondata_backend}/1817427/36dfa08eea10bb99bfc47f3f1c2f11cd961db2f4/resource.tar.gz#test.test_blocks-multimap-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[blocks-or_scalar-default.txt-Debug]": [
         {
             "checksum": "9153a8e9df67920deba9d895e8d4a8b1",
diff --git a/yql/essentials/tests/sql/minirun/part2/canondata/result.json b/yql/essentials/tests/sql/minirun/part2/canondata/result.json
index 39f596db22..93a0b6e07e 100644
--- a/yql/essentials/tests/sql/minirun/part2/canondata/result.json
+++ b/yql/essentials/tests/sql/minirun/part2/canondata/result.json
@@ -258,6 +258,48 @@
             "uri": "https://{canondata_backend}/1942671/13a781c07395a2ba482afc02d12bb4a6fa60d5ef/resource.tar.gz#test.test_blocks-exists-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[blocks-sort-default.txt-Debug]": [
+        {
+            "checksum": "f5d0c3ad8b655c3a6684a13ff19885e6",
+            "size": 437,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-sort-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[blocks-sort-default.txt-Peephole]": [
+        {
+            "checksum": "ac26701a5abd8c7670e227bb534127ea",
+            "size": 879,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-sort-default.txt-Peephole_/opt.yql"
+        }
+    ],
+    "test.test[blocks-sort-default.txt-Results]": [
+        {
+            "checksum": "0827720cfdc7bfc868110137b98d6ad6",
+            "size": 1535,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-sort-default.txt-Results_/results.txt"
+        }
+    ],
+    "test.test[blocks-take_skip-default.txt-Debug]": [
+        {
+            "checksum": "14f19770f1c340e0707f94dd4c9351f4",
+            "size": 377,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-take_skip-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[blocks-take_skip-default.txt-Peephole]": [
+        {
+            "checksum": "6de1c3b7d8decabd30376dc756d53d17",
+            "size": 879,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-take_skip-default.txt-Peephole_/opt.yql"
+        }
+    ],
+    "test.test[blocks-take_skip-default.txt-Results]": [
+        {
+            "checksum": "0b490567ef58097bc590e7c286c55804",
+            "size": 1449,
+            "uri": "https://{canondata_backend}/1599023/429c43efbf699a2907bed12a3ec8e6c4b74d39ce/resource.tar.gz#test.test_blocks-take_skip-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[compute_range-in_literal_nulls-default.txt-Debug]": [
         {
             "checksum": "32fad62c2c315c6dad7fe390ba7f40a8",
diff --git a/yql/essentials/tests/sql/minirun/part5/canondata/result.json b/yql/essentials/tests/sql/minirun/part5/canondata/result.json
index 3bbd5eafd1..148c9d957e 100644
--- a/yql/essentials/tests/sql/minirun/part5/canondata/result.json
+++ b/yql/essentials/tests/sql/minirun/part5/canondata/result.json
@@ -366,6 +366,27 @@
             "uri": "https://{canondata_backend}/1775319/488597481a390ab007602d4f2c7cb21f0d5833de/resource.tar.gz#test.test_blocks-decimal_binop-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[blocks-extend-default.txt-Debug]": [
+        {
+            "checksum": "d01511a7bd9623a15c633d05a4f1bdd9",
+            "size": 314,
+            "uri": "https://{canondata_backend}/1937001/b704bbc48624073924a714c5207966ffdfdddc62/resource.tar.gz#test.test_blocks-extend-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[blocks-extend-default.txt-Peephole]": [
+        {
+            "checksum": "424973dad53667a534c277cbb74c20f7",
+            "size": 624,
+            "uri": "https://{canondata_backend}/1937001/b704bbc48624073924a714c5207966ffdfdddc62/resource.tar.gz#test.test_blocks-extend-default.txt-Peephole_/opt.yql"
+        }
+    ],
+    "test.test[blocks-extend-default.txt-Results]": [
+        {
+            "checksum": "c7d2990b6653e16a29caab47b63e78b7",
+            "size": 981,
+            "uri": "https://{canondata_backend}/1937001/b704bbc48624073924a714c5207966ffdfdddc62/resource.tar.gz#test.test_blocks-extend-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[blocks-filter-default.txt-Debug]": [
         {
             "checksum": "3dea830ec45f78ed78603576b054b4ca",
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/result.json b/yql/essentials/tests/sql/sql2yql/canondata/result.json
index 0d7fe3bcf9..4576fae869 100644
--- a/yql/essentials/tests/sql/sql2yql/canondata/result.json
+++ b/yql/essentials/tests/sql/sql2yql/canondata/result.json
@@ -1385,6 +1385,13 @@
             "uri": "https://{canondata_backend}/1916746/15a617669fb2b654b6be1fc493d61b99fb1f6f81/resource.tar.gz#test_sql2yql.test_blocks-exists_/sql.yql"
         }
     ],
+    "test_sql2yql.test[blocks-extend]": [
+        {
+            "checksum": "d79b2f801b2e1d96d54c1f1b6047d3a4",
+            "size": 1722,
+            "uri": "https://{canondata_backend}/1871102/7c02af7ee8a00a03832f80f641cc353cdcdeefd4/resource.tar.gz#test_sql2yql.test_blocks-extend_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[blocks-filter]": [
         {
             "checksum": "e09ada2d8790a1c8d43bdc76683eaa40",
@@ -1427,6 +1434,13 @@
             "uri": "https://{canondata_backend}/1900335/c447765ddbde200b8fe3ee8091f4d625b36b6bc6/resource.tar.gz#test_sql2yql.test_blocks-member_/sql.yql"
         }
     ],
+    "test_sql2yql.test[blocks-multimap]": [
+        {
+            "checksum": "d331aca01555e8a4306b317f4e3a3a05",
+            "size": 1888,
+            "uri": "https://{canondata_backend}/1817427/1181987a1e6ad55ee2b11f6938ee9daecb115152/resource.tar.gz#test_sql2yql.test_blocks-multimap_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[blocks-not]": [
         {
             "checksum": "60ffa8348b787bc96b78945aa79545d2",
@@ -1504,6 +1518,20 @@
             "uri": "https://{canondata_backend}/1130705/7bfa4c7bef8ecf6ffba8b79d210708dc39bab04e/resource.tar.gz#test_sql2yql.test_blocks-projection_add_ints_filter_/sql.yql"
         }
     ],
+    "test_sql2yql.test[blocks-sort]": [
+        {
+            "checksum": "69badbf5aa282179d350e6d76c20e026",
+            "size": 2089,
+            "uri": "https://{canondata_backend}/1871102/7c02af7ee8a00a03832f80f641cc353cdcdeefd4/resource.tar.gz#test_sql2yql.test_blocks-sort_/sql.yql"
+        }
+    ],
+    "test_sql2yql.test[blocks-take_skip]": [
+        {
+            "checksum": "2c5f4c1f9a6c48b0dc0f186ed84c858b",
+            "size": 2205,
+            "uri": "https://{canondata_backend}/1871102/7c02af7ee8a00a03832f80f641cc353cdcdeefd4/resource.tar.gz#test_sql2yql.test_blocks-take_skip_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[blocks-topg]": [
         {
             "checksum": "0f669eec4c7c11d663db474f09902bd4",
@@ -8129,6 +8157,11 @@
             "uri": "file://test_sql_format.test_blocks-exists_/formatted.sql"
         }
     ],
+    "test_sql_format.test[blocks-extend]": [
+        {
+            "uri": "file://test_sql_format.test_blocks-extend_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[blocks-filter]": [
         {
             "uri": "file://test_sql_format.test_blocks-filter_/formatted.sql"
@@ -8159,6 +8192,11 @@
             "uri": "file://test_sql_format.test_blocks-member_/formatted.sql"
         }
     ],
+    "test_sql_format.test[blocks-multimap]": [
+        {
+            "uri": "file://test_sql_format.test_blocks-multimap_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[blocks-not]": [
         {
             "uri": "file://test_sql_format.test_blocks-not_/formatted.sql"
@@ -8214,6 +8252,16 @@
             "uri": "file://test_sql_format.test_blocks-projection_add_ints_filter_/formatted.sql"
         }
     ],
+    "test_sql_format.test[blocks-sort]": [
+        {
+            "uri": "file://test_sql_format.test_blocks-sort_/formatted.sql"
+        }
+    ],
+    "test_sql_format.test[blocks-take_skip]": [
+        {
+            "uri": "file://test_sql_format.test_blocks-take_skip_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[blocks-topg]": [
         {
             "uri": "file://test_sql_format.test_blocks-topg_/formatted.sql"
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-extend_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-extend_/formatted.sql
new file mode 100644
index 0000000000..4acd27012c
--- /dev/null
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-extend_/formatted.sql
@@ -0,0 +1,13 @@
+$data1 = [<|x: 1|>, <|x: 3|>, <|x: 2|>];
+$data2 = [<|x: 4|>, <|x: 1|>];
+
+SELECT
+    *
+FROM
+    as_table($data1)
+UNION ALL
+SELECT
+    *
+FROM
+    as_table($data2)
+;
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-multimap_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-multimap_/formatted.sql
new file mode 100644
index 0000000000..4def8cc8ed
--- /dev/null
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-multimap_/formatted.sql
@@ -0,0 +1,13 @@
+$data = [<|x: 1|>, <|x: 2|>];
+
+SELECT
+    *
+FROM (
+    SELECT
+        x,
+        [3, 4] AS y
+    FROM
+        as_table($data)
+)
+    FLATTEN LIST BY y
+;
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-sort_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-sort_/formatted.sql
new file mode 100644
index 0000000000..8a3a1f9af8
--- /dev/null
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-sort_/formatted.sql
@@ -0,0 +1,17 @@
+$data = [<|x: 1|>, <|x: 3|>, <|x: 2|>];
+
+SELECT
+    *
+FROM
+    as_table($data)
+ORDER BY
+    x
+;
+
+SELECT
+    *
+FROM
+    as_table($data)
+ORDER BY
+    x DESC
+LIMIT 2;
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-take_skip_/formatted.sql b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-take_skip_/formatted.sql
new file mode 100644
index 0000000000..179fdab916
--- /dev/null
+++ b/yql/essentials/tests/sql/sql2yql/canondata/test_sql_format.test_blocks-take_skip_/formatted.sql
@@ -0,0 +1,15 @@
+PRAGMA warning('disable', '4537');
+
+$data = [<|x: 1|>, <|x: 2|>, <|x: 3|>];
+
+SELECT
+    *
+FROM
+    as_table($data)
+LIMIT 2;
+
+SELECT
+    *
+FROM
+    as_table($data)
+LIMIT 1 OFFSET 1;
diff --git a/yql/essentials/tests/sql/suites/blocks/extend.sql b/yql/essentials/tests/sql/suites/blocks/extend.sql
new file mode 100644
index 0000000000..2e713c376a
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/blocks/extend.sql
@@ -0,0 +1,6 @@
+$data1 = [<|x:1|>,<|x:3|>,<|x:2|>];
+$data2 = [<|x:4|>,<|x:1|>];
+
+select * from as_table($data1)
+union all
+select * from as_table($data2)
diff --git a/yql/essentials/tests/sql/suites/blocks/multimap.sql b/yql/essentials/tests/sql/suites/blocks/multimap.sql
new file mode 100644
index 0000000000..d26fb1a256
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/blocks/multimap.sql
@@ -0,0 +1,3 @@
+$data = [<|x:1|>,<|x:2|>];
+
+select * from (select x,[3,4] as y from as_table($data)) flatten list by y
diff --git a/yql/essentials/tests/sql/suites/blocks/sort.sql b/yql/essentials/tests/sql/suites/blocks/sort.sql
new file mode 100644
index 0000000000..61fef4aeee
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/blocks/sort.sql
@@ -0,0 +1,4 @@
+$data = [<|x:1|>,<|x:3|>,<|x:2|>];
+
+select * from as_table($data) order by x;
+select * from as_table($data) order by x desc limit 2;
diff --git a/yql/essentials/tests/sql/suites/blocks/take_skip.sql b/yql/essentials/tests/sql/suites/blocks/take_skip.sql
new file mode 100644
index 0000000000..11a7599f5a
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/blocks/take_skip.sql
@@ -0,0 +1,5 @@
+pragma warning("disable","4537");
+$data = [<|x:1|>,<|x:2|>,<|x:3|>];
+
+select * from as_table($data) limit 2;
+select * from as_table($data) limit 1 offset 1;