[] Add optimizers for Take/Limit/Skip over literal lists

3 files changed, 43 insertions, 22 deletions

diff --git a/ydb/core/kqp/ut/kqp_stats_ut.cpp b/ydb/core/kqp/ut/kqp_stats_ut.cpp
index 2228f3c6600..61e7ccc6e6d 100644
--- a/ydb/core/kqp/ut/kqp_stats_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_stats_ut.cpp
@@ -356,7 +356,7 @@ Y_UNIT_TEST_TWIN(StatsProfile, UseSessionActor) {
     auto node1 = FindPlanNodeByKv(plan, "Node Type", "Aggregate-TableFullScan");
     UNIT_ASSERT_EQUAL(node1.GetMap().at("Stats").GetMapSafe().at("ComputeNodes").GetArraySafe().size(), 2);
 
-    auto node2 = FindPlanNodeByKv(plan, "Node Type", "Aggregate-Limit");
+    auto node2 = FindPlanNodeByKv(plan, "Node Type", "Aggregate");
     UNIT_ASSERT_EQUAL(node2.GetMap().at("Stats").GetMapSafe().at("ComputeNodes").GetArraySafe().size(), 1);
 }
 
diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index 4c70f3cc147..96319ecc28b 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -153,6 +153,13 @@ TExprNode::TPtr KeepConstraints(TExprNode::TPtr node, const TExprNode& src, TExp
     return res;
 }
 
+TExprNode::TPtr MakeEmptyCollectionWithConstraintsFrom(const TExprNode& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+    auto res = ctx.NewCallable(node.Pos(), GetEmptyCollectionName(node.GetTypeAnn()),
+        { ExpandType(node.Pos(), *node.GetTypeAnn(), ctx) });
+    res = KeepConstraints(res, node, ctx);
+    return KeepColumnOrder(res, node, ctx, *optCtx.Types);
+}
+
 template<typename TInt>
 bool ConstIntAggregate(const TExprNode::TChildrenType& values, std::function<TInt(TInt, TInt)> aggFunc,
     TInt& result)
@@ -3830,29 +3837,55 @@ void RegisterCoSimpleCallables1(TCallableOptimizerMap& map) {
     map["FlatListIf"] = std::bind(&OptimizeFlatContainerIf<true>, _1, _2);
     map["FlatOptionalIf"] = std::bind(&OptimizeFlatContainerIf<false>, _1, _2);
 
-    map["Skip"] = [](const TExprNode::TPtr& node, TExprContext& /*ctx*/, TOptimizeContext& /*optCtx*/) {
+    map["Skip"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        if (node->Head().IsCallable("List")) {
+            YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
+            return node->HeadPtr();
+        }
         if (node->Tail().IsCallable("Uint64")) {
             const auto value = FromString<ui64>(node->Tail().Head().Content());
             if (!value) {
                 YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content();
                 return node->HeadPtr();
+            } else if (value == std::numeric_limits<ui64>::max()) {
+                YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content();
+                return MakeEmptyCollectionWithConstraintsFrom(*node, ctx, optCtx);
+            } else if (node->Head().IsCallable("AsList")) {
+                YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content()
+                                      << " over " << node->Head().Content();
+                if (node->Head().ChildrenSize() > value) {
+                    auto children = node->Head().ChildrenList();
+                    children.erase(children.begin(), children.begin() + value);
+                    return ctx.ChangeChildren(node->Head(), std::move(children));
+                }
+                return MakeEmptyCollectionWithConstraintsFrom(*node, ctx, optCtx);
             }
         }
 
         return node;
     };
 
-    map["Take"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+    map["Take"] = map["Limit"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
+        if (node->Head().IsCallable("List")) {
+            YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
+            return node->HeadPtr();
+        }
         if (node->Tail().IsCallable("Uint64")) {
             const auto value = FromString<ui64>(node->Tail().Head().Content());
             if (!value) {
                 YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content();
-                auto res = ctx.NewCallable(node->Tail().Pos(), GetEmptyCollectionName(node->GetTypeAnn()), {ExpandType(node->Pos(), *node->GetTypeAnn(), ctx)});
-                res = KeepConstraints(res, *node, ctx);
-                return KeepColumnOrder(res, *node, ctx, *optCtx.Types);
+                return MakeEmptyCollectionWithConstraintsFrom(*node, ctx, optCtx);
             } else if (value == std::numeric_limits<ui64>::max()) {
                 YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content();
                 return node->HeadPtr();
+            } else if (node->Head().IsCallable("AsList")) {
+                YQL_CLOG(DEBUG, Core) << node->Content() << " with " << node->Tail().Content() << " '" << node->Tail().Head().Content() << " over " << node->Head().Content();
+                if (node->Head().ChildrenSize() <= value) {
+                    return node->HeadPtr();
+                }
+                auto children = node->Head().ChildrenList();
+                children.resize(value);
+                return ctx.ChangeChildren(node->Head(), std::move(children));
             }
         }
 
diff --git a/ydb/tests/functional/canonical/canondata/test_sql.TestCanonicalFolder1.test_case_explain.script-script_/explain.script.plan b/ydb/tests/functional/canonical/canondata/test_sql.TestCanonicalFolder1.test_case_explain.script-script_/explain.script.plan
index 335217c5ae4..e5536bc41a9 100644
--- a/ydb/tests/functional/canonical/canondata/test_sql.TestCanonicalFolder1.test_case_explain.script-script_/explain.script.plan
+++ b/ydb/tests/functional/canonical/canondata/test_sql.TestCanonicalFolder1.test_case_explain.script-script_/explain.script.plan
@@ -94,13 +94,9 @@
                                 "Operators": [
                                     {
                                         "Name": "Aggregate"
-                                    },
-                                    {
-                                        "Name": "Limit",
-                                        "Limit": "1001"
                                     }
                                 ],
-                                "Node Type": "Aggregate-Limit"
+                                "Node Type": "Aggregate"
                             }
                         ],
                         "Node Type": "Precompute_0_0",
@@ -211,13 +207,9 @@
                                 "Operators": [
                                     {
                                         "Name": "Aggregate"
-                                    },
-                                    {
-                                        "Name": "Limit",
-                                        "Limit": "1001"
                                     }
                                 ],
-                                "Node Type": "Aggregate-Limit"
+                                "Node Type": "Aggregate"
                             }
                         ],
                         "Node Type": "Precompute_0_0",
@@ -285,15 +277,11 @@
                                 "PlanNodeId": 1,
                                 "Operators": [
                                     {
-                                        "Iterator": "Take",
+                                        "Iterator": "[{column0: 1 * 2 * 3 * 4 * 5}]",
                                         "Name": "Iterator"
-                                    },
-                                    {
-                                        "Name": "Limit",
-                                        "Limit": "1001"
                                     }
                                 ],
-                                "Node Type": "ConstantExpr-Limit"
+                                "Node Type": "ConstantExpr"
                             }
                         ],
                         "Node Type": "ResultSet_0_0",