Implement ListTop/ListTopSort builtins (ydb-platform#4231)

shnikd · May 24, 2024 · 13b74c9 · 13b74c9
1 parent ae4e8d5
commit 13b74c9
Show file tree

Hide file tree

Showing 9 changed files with 156 additions and 0 deletions.
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -12471,6 +12471,12 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["ListFlatten"] = &ListFlattenWrapper;
         Functions["ListUniq"] = &ListUniqWrapper;
         Functions["ListUniqStable"] = &ListUniqStableWrapper;
+        Functions["ListTop"] = &ListTopSortWrapper;
+        Functions["ListTopAsc"] = &ListTopSortWrapper;
+        Functions["ListTopDesc"] = &ListTopSortWrapper;
+        Functions["ListTopSort"] = &ListTopSortWrapper;
+        Functions["ListTopSortAsc"] = &ListTopSortWrapper;
+        Functions["ListTopSortDesc"] = &ListTopSortWrapper;
 
         Functions["ExpandMap"] = &ExpandMapWrapper;
         Functions["WideMap"] = &WideMapWrapper;

diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -1431,6 +1431,45 @@ namespace {
         return OptListWrapperImpl<3U>(input, output, ctx, "Sort");
     }
 
+    IGraphTransformer::TStatus ListTopSortWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        if (!EnsureMinMaxArgsCount(*input, 2, 3, ctx.Expr)) { 
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        TStringBuf newName = input->Content();
+        newName.Skip(4);
+        bool desc = false;
+        if (newName.EndsWith("Asc")) {
+            newName.Chop(3);
+        } else if (newName.EndsWith("Desc")) {
+            newName.Chop(4);
+            desc = true;
+        }
+
+        TExprNode::TPtr sortLambda = nullptr;
+        if (input->ChildrenSize() == 3) {
+            sortLambda = input->ChildPtr(2);
+        } else {
+            sortLambda = ctx.Expr.Builder(input->Pos())
+                .Lambda()
+                    .Param("item")
+                    .Arg("item")
+                .Seal()
+            .Build();
+        }
+
+        return OptListWrapperImpl<4U, 4U>(ctx.Expr.Builder(input->Pos())
+            .Callable(newName)
+                .Add(0, input->ChildPtr(0))
+                .Add(1, input->ChildPtr(1))
+                .Callable(2, "Bool")
+                    .Atom(0, desc ? "false" : "true")
+                .Seal()
+                .Add(3, sortLambda)
+            .Seal()
+        .Build(), output, ctx, newName);
+    }
+
     IGraphTransformer::TStatus ListExtractWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         return OptListWrapperImpl<2U>(input, output, ctx, "OrderedExtract");
     }

diff --git a/ydb/library/yql/core/type_ann/type_ann_list.h b/ydb/library/yql/core/type_ann/type_ann_list.h
@@ -37,6 +37,8 @@ namespace NTypeAnnImpl {
     IGraphTransformer::TStatus ListEnumerateWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListReverseWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListSortWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+    IGraphTransformer::TStatus ListTopWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+    IGraphTransformer::TStatus ListTopSortWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListExtractWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus ListCollectWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
     IGraphTransformer::TStatus FoldMapWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
@@ -2904,6 +2904,12 @@ struct TBuiltinFuncData {
             {"listcollect", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListCollect", 1, 1) },
             {"listnotnull", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListNotNull", 1, 1)},
             {"listflatten", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListFlatten", 1, 1)},
+            {"listtop", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTop", 2, 3)},
+            {"listtopasc", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTopAsc", 2, 3)},
+            {"listtopdesc", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTopDesc", 2, 3)},
+            {"listtopsort", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTopSort", 2, 3)},
+            {"listtopsortasc", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTopSortAsc", 2, 3)},
+            {"listtopsortdesc", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ListTopSortDesc", 2, 3)},
 
             // Dict builtins
             {"dictlength", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("Length", 1, 1)},

diff --git a/ydb/library/yql/tests/sql/dq_file/part12/canondata/result.json b/ydb/library/yql/tests/sql/dq_file/part12/canondata/result.json
@@ -948,6 +948,28 @@
         }
     ],
     "test.test[expr-expr_null-default.txt-Results]": [],
+    "test.test[expr-list_top_sort-default.txt-Analyze]": [
+        {
+            "checksum": "01775e7c945a56ebf0edc2d478f4f68d",
+            "size": 2567,
+            "uri": "https://{canondata_backend}/1130705/02b49b8f7de27e2bd653274fbb3967717e0d56f2/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Analyze_/plan.txt"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Debug]": [
+        {
+            "checksum": "df5a5689f7d16f67cdde86a519852534",
+            "size": 1805,
+            "uri": "https://{canondata_backend}/1777230/f087d19aefc64f43c561b1716c8824e128ac8093/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Debug_/opt.yql_patched"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Plan]": [
+        {
+            "checksum": "01775e7c945a56ebf0edc2d478f4f68d",
+            "size": 2567,
+            "uri": "https://{canondata_backend}/1130705/02b49b8f7de27e2bd653274fbb3967717e0d56f2/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Plan_/plan.txt"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Results]": [],
     "test.test[expr-many_opt_comp-default.txt-Analyze]": [
         {
             "checksum": "01775e7c945a56ebf0edc2d478f4f68d",

diff --git a/ydb/library/yql/tests/sql/hybrid_file/part10/canondata/result.json b/ydb/library/yql/tests/sql/hybrid_file/part10/canondata/result.json
@@ -1049,6 +1049,20 @@
             "uri": "https://{canondata_backend}/1775319/3515b86fb929979a6751f93bd43a0291eaa01262/resource.tar.gz#test.test_expr-lds_empty_compare-default.txt-Plan_/plan.txt"
         }
     ],
+    "test.test[expr-list_top_sort-default.txt-Debug]": [
+        {
+            "checksum": "049a4840aed77024dcb5f9caf73a672f",
+            "size": 1804,
+            "uri": "https://{canondata_backend}/1937424/ef4272c0e98c55575149317381e8efaa85a26157/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Debug_/opt.yql_patched"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Plan]": [
+        {
+            "checksum": "01775e7c945a56ebf0edc2d478f4f68d",
+            "size": 2567,
+            "uri": "https://{canondata_backend}/1903885/804adfc45546340f178ce737d7f1d1e8feb56e81/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Plan_/plan.txt"
+        }
+    ],
     "test.test[expr-to_hashed_dict_list_key-default.txt-Debug]": [
         {
             "checksum": "044532517685890dde36188f8d1567cc",

diff --git a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
@@ -5837,6 +5837,13 @@
             "uri": "https://{canondata_backend}/1784117/d56ae82ad9d30397a41490647be1bd2124718f98/resource.tar.gz#test_sql2yql.test_expr-list_to_from_tuple_/sql.yql"
         }
     ],
+    "test_sql2yql.test[expr-list_top_sort]": [
+        {
+            "checksum": "b2acc119076239ea4f4ada66e164330b",
+            "size": 15716,
+            "uri": "https://{canondata_backend}/995452/0828572c263b2cebd15b1dcd234d39fd5315b924/resource.tar.gz#test_sql2yql.test_expr-list_top_sort_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[expr-list_uniq]": [
         {
             "checksum": "d15d661e6b25dd80b507f95fb057c207",
@@ -24408,6 +24415,13 @@
             "uri": "https://{canondata_backend}/1925842/ffeacd65677f75bcab2445d237890283afc5e8f6/resource.tar.gz#test_sql_format.test_expr-list_to_from_tuple_/formatted.sql"
         }
     ],
+    "test_sql_format.test[expr-list_top_sort]": [
+        {
+            "checksum": "a0af54a6ec4d671db9f8a91cc09877bf",
+            "size": 1594,
+            "uri": "https://{canondata_backend}/1689644/47813afec92f7bbc83af4667f9074aac865e7ffe/resource.tar.gz#test_sql_format.test_expr-list_top_sort_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[expr-list_uniq]": [
         {
             "checksum": "74f6f47d443fa97917b2d52760056649",

diff --git a/ydb/library/yql/tests/sql/suites/expr/list_top_sort.sql b/ydb/library/yql/tests/sql/suites/expr/list_top_sort.sql
@@ -0,0 +1,32 @@
+/* syntax version 1 */
+
+$list = [45, 20, 34, 16];
+$empty = [];
+
+$digit = ($x) -> { return $x % 10; };
+
+SELECT 
+    ListTopSort($list, 2), ListTopSort($list, 4), ListTopSort($list, 5), 
+    ListTopSort($list, 2, $digit), 
+    ListTopSort(NULL, 2), ListTopSort(Just($list), 2), ListTopSort($empty, 0);
+SELECT
+    ListTopSortAsc($list, 2), ListTopSortAsc($list, 4), ListTopSortAsc($list, 5), 
+    ListTopSortAsc($list, 2, $digit), 
+    ListTopSortAsc(NULL, 2), ListTopSortAsc(Just($list), 2), ListTopSortAsc($empty, 0);
+SELECT 
+    ListTopSortDesc($list, 2), ListTopSortDesc($list, 4), ListTopSortDesc($list, 5),
+    ListTopSortDesc($list, 2, $digit),
+    ListTopSortDesc(NULL, 2), ListTopSortDesc(Just($list), 2), ListTopSortDesc($empty, 0);
+
+SELECT 
+    ListSort(ListTop($list, 2)), ListSort(ListTop($list, 4)), ListSort(ListTop($list, 5)), 
+    ListSort(ListTop($list, 2, $digit)), 
+    ListSort(ListTop(NULL, 2)), ListSort(ListTop(Just($list), 2)), ListSort(ListTop($empty, 0));
+SELECT
+    ListSort(ListTopAsc($list, 2)), ListSort(ListTopAsc($list, 4)), ListSort(ListTopAsc($list, 5)), 
+    ListSort(ListTopAsc($list, 2, $digit)), 
+    ListSort(ListTopAsc(NULL, 2)), ListSort(ListTopAsc(Just($list), 2)), ListSort(ListTopAsc($empty, 0));
+SELECT 
+    ListSort(ListTopDesc($list, 2)), ListSort(ListTopDesc($list, 4)), ListSort(ListTopDesc($list, 5)),
+    ListSort(ListTopDesc($list, 2, $digit)),
+    ListSort(ListTopDesc(NULL, 2)), ListSort(ListTopDesc(Just($list), 2)), ListSort(ListTopDesc($empty, 0));
diff --git a/ydb/library/yql/tests/sql/yt_native_file/part12/canondata/result.json b/ydb/library/yql/tests/sql/yt_native_file/part12/canondata/result.json
@@ -975,6 +975,27 @@
             "uri": "https://{canondata_backend}/1881367/f7c9625b9d66b51a2a5927015bd3050c9cb7f8d2/resource.tar.gz#test.test_expr-expr_null-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[expr-list_top_sort-default.txt-Debug]": [
+        {
+            "checksum": "7f0089c8ef7e1971bdf7f46c7d96b45c",
+            "size": 1733,
+            "uri": "https://{canondata_backend}/1923547/ddadfd08d4d18962f17ecb8528c68375e0f1cf53/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Plan]": [
+        {
+            "checksum": "01775e7c945a56ebf0edc2d478f4f68d",
+            "size": 2567,
+            "uri": "https://{canondata_backend}/1937492/5708f4083beb3611506ac1fc16491e776779d1cf/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Plan_/plan.txt"
+        }
+    ],
+    "test.test[expr-list_top_sort-default.txt-Results]": [
+        {
+            "checksum": "79ad276e0af4567e7f2c7aa689438f80",
+            "size": 23306,
+            "uri": "https://{canondata_backend}/1923547/ddadfd08d4d18962f17ecb8528c68375e0f1cf53/resource.tar.gz#test.test_expr-list_top_sort-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[expr-many_opt_comp-default.txt-Debug]": [
         {
             "checksum": "cf876c64f15a29de2edc2b23e3b21565",