Handle overflow in DqBuildTakeSkipStage optimizer (#14491)

11 files changed, 32 insertions, 14 deletions

diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index ab4951a1c79..ed62d85445d 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -762,7 +762,7 @@ TExprBase DqBuildFlatmapStage(TExprBase node, TExprContext& ctx, IOptimizationCo
     const TParentsMap& parentsMap, bool allowStageMultiUsage)
 {
     Y_UNUSED(optCtx);
-    
+
     if (!node.Maybe<TCoFlatMapBase>().Input().Maybe<TDqCnUnionAll>()) {
         return node;
     }
@@ -2014,11 +2014,29 @@ TExprBase DqBuildTakeSkipStage(TExprBase node, TExprContext& ctx, IOptimizationC
         .Args({"stream"})
         .Body<TCoTake>()
             .Input("stream")
-            .Count<TCoAggrAdd>()
-                .Left(take.Count())
-                .Right(skip.Count())
+            .Count<TCoIf>()
+                .Predicate<TCoCmpGreater>()
+                    .Left(take.Count())
+                    .Right<TCoSub>()
+                        .Left<TCoUint64>()
+                            .Literal()
+                                .Value(ToString(Max<ui64>()), TNodeFlags::Default)
+                            .Build()
+                        .Build()
+                        .Right(skip.Count())
+                    .Build()
+                .Build()
+                .ThenValue<TCoUint64>()
+                    .Literal()
+                        .Value(ToString(Max<ui64>()), TNodeFlags::Default)
+                    .Build()
+                .Build()
+                .ElseValue<TCoAggrAdd>()
+                    .Left(take.Count())
+                    .Right(skip.Count())
                 .Build()
             .Build()
+        .Build()
         .Done();
 
     auto newDqUnion = DqPushLambdaToStageUnionAll(dqUnion, lambda, {}, ctx, optCtx);
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-38 b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-38
index be0737d9d3e..29b7d9d47e7 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-38
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-38
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-39 b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-39
index 3c0e77c8fa1..5b78aaf964c 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-39
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-39
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-40 b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-40
index 85e3ed054df..a87ce4a6ce6 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-40
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-40
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,100)",
+                                                        "Limit": "Min(If,SUM(10,100))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-41 b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-41
index 8983bd07309..f67cc254e2d 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-41
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-41
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,10000)",
+                                                        "Limit": "Min(If,SUM(10,10000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-42 b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-42
index 775a293527d..ccb1a23253b 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-42
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_column_/queries-original-plan-column-42
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-38 b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-38
index 4225a0c3276..16e9bcdbffe 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-38
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-38
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-39 b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-39
index 285b1b9d5d4..a97ba87e6c2 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-39
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-39
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-40 b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-40
index 4eda1d3bb92..c5955001f27 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-40
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-40
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,100)",
+                                                        "Limit": "Min(If,SUM(10,100))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-41 b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-41
index 7f94be1904a..67b6710e7c4 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-41
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-41
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,10000)",
+                                                        "Limit": "Min(If,SUM(10,10000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],
diff --git a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-42 b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-42
index 64136571d70..771b01d4fb3 100644
--- a/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-42
+++ b/ydb/tests/functional/clickbench/canondata/test.test_plans_row_/queries-original-plan-row-42
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],