Handle overflow in DqBuildTakeSkipStage optimizer (#14491)

rvu1024 · web-flow · commit b62ee74dc98f · 2025-02-13T18:24:49.000+03:00
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -52,7 +52,7 @@
                                                                 "ExternalPlanNodeId": 4
                                                             }
                                                         ],
-                                                        "Limit": "SUM(10,1000)",
+                                                        "Limit": "Min(If,SUM(10,1000))",
                                                         "Name": "Limit"
                                                     }
                                                 ],

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@`
`52`	`52`	`"ExternalPlanNodeId": 4`
`53`	`53`	`}`
`54`	`54`	`],`
`55`		`- "Limit": "SUM(10,1000)",`
	`55`	`+ "Limit": "Min(If,SUM(10,1000))",`
`56`	`56`	`"Name": "Limit"`
`57`	`57`	`}`
`58`	`58`	`],`