Fix overflow in take + skip

commit_hash:2e37109d1766f0f7b76363f60d3e7217686eae48

5 files changed, 75 insertions, 19 deletions

diff --git a/yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp b/yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp
index bc307c3178..7916821ca4 100644
--- a/yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp
+++ b/yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp
@@ -69,7 +69,11 @@ void TRecordsRange::Fill(const TExprNode& settingsNode) {
         if (settingName != TStringBuf("take") && settingName != TStringBuf("skip")) {
             continue;
         }
-        YQL_ENSURE(setting->Child(1)->IsCallable("Uint64"));
+        if (!setting->Child(1)->IsCallable("Uint64")) {
+            Offset.Clear();
+            Limit.Clear();
+            return;
+        }
         if (!UpdateRecordsRange(*this, settingName, NYql::FromString<ui64>(*setting->Child(1)->Child(0), NUdf::EDataSlot::Uint64))) {
             break;
         }
diff --git a/yt/yql/providers/yt/lib/res_pull/table_limiter.cpp b/yt/yql/providers/yt/lib/res_pull/table_limiter.cpp
index 98731da58f..ddd4781ed6 100644
--- a/yt/yql/providers/yt/lib/res_pull/table_limiter.cpp
+++ b/yt/yql/providers/yt/lib/res_pull/table_limiter.cpp
@@ -8,11 +8,16 @@ namespace NYql {
 
 TTableLimiter::TTableLimiter(const TRecordsRange& range)
     : Start(range.Offset.GetOrElse(0ULL))
-    , End(range.Limit.Defined() ? Start + *range.Limit : Max())
     , Current(0ULL)
     , TableStart(0ULL)
     , TableEnd(Max())
 {
+    const auto limit = range.Limit.GetOrElse(Max());
+    if (limit > Max<ui64>() - Start) {
+        End = Max();
+    } else {
+        End = Start + limit;
+    }
 }
 
 bool TTableLimiter::NextTable(ui64 recordCount) {
diff --git a/yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp b/yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp
index 0c8c9fdf6c..b3233e37f1 100644
--- a/yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp
@@ -279,31 +279,41 @@ private:
                 .Build()
             .Done();
 
-        TExprNode::TPtr limit;
-        if (const auto& limitNode = NYql::GetSetting(sort.Settings().Ref(), EYtSettingType::Limit)) {
-            limit = GetLimitExpr(limitNode, ctx);
-        }
-
+        TExprNode::TPtr work;
         auto [direct, selector] = GetOutputSortSettings(sort, ctx);
-        auto work = direct && selector ?
-            limit ?
-                Build<TCoTopSort>(ctx, sort.Pos())
+        if (direct && selector) {
+            // Don't use runtime limit for TopSort - it may have max<ui64>() value, which cause TopSort to fail
+            TMaybe<ui64> limit = GetLimit(sort.Settings().Ref());
+            work = limit
+                ? Build<TCoTopSort>(ctx, sort.Pos())
                     .Input(input)
-                    .Count(std::move(limit))
+                    .Count<TCoUint64>()
+                        .Literal()
+                            .Value(ToString(*limit), TNodeFlags::Default)
+                        .Build()
+                    .Build()
                     .SortDirections(std::move(direct))
                     .KeySelectorLambda(std::move(selector))
-                    .Done().Ptr():
-                Build<TCoSort>(ctx, sort.Pos())
+                    .Done().Ptr()
+                : Build<TCoSort>(ctx, sort.Pos())
                     .Input(input)
                     .SortDirections(std::move(direct))
                     .KeySelectorLambda(std::move(selector))
-                    .Done().Ptr():
-            limit ?
-                Build<TCoTake>(ctx, sort.Pos())
+                    .Done().Ptr()
+                ;
+        } else {
+            TExprNode::TPtr limit;
+            if (const auto& limitNode = NYql::GetSetting(sort.Settings().Ref(), EYtSettingType::Limit)) {
+                limit = GetLimitExpr(limitNode, ctx);
+            }
+
+            work = limit
+                ? Build<TCoTake>(ctx, sort.Pos())
                     .Input(input)
                     .Count(std::move(limit))
-                    .Done().Ptr():
-                input.Ptr();
+                    .Done().Ptr()
+                : input.Ptr();
+        }
 
         auto settings = NYql::AddSetting(sort.Settings().Ref(), EYtSettingType::NoDq, {}, ctx);
         auto operation = ctx.ChangeChild(sort.Ref(), TYtTransientOpBase::idx_Settings, std::move(settings));
diff --git a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
index 97b9fdeda9..d8b1f20823 100644
--- a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
@@ -905,7 +905,30 @@ TExprNode::TPtr GetLimitExpr(const TExprNode::TPtr& limitSetting, TExprContext&
         }
 
         if (skip) {
-            limitValues.push_back(ctx.NewCallable(child->Pos(), "+", { take, skip }));
+            auto uintMax = ctx.Builder(child->Pos())
+                .Callable("Uint64")
+                    .Atom(0, ToString(Max<ui64>()), TNodeFlags::Default)
+                .Seal()
+                .Build();
+            limitValues.push_back(
+                ctx.Builder(child->Pos())
+                    .Callable("If")
+                        .Callable(0, ">")
+                            .Add(0, take)
+                            .Callable(1, "-")
+                                .Add(0, uintMax)
+                                .Add(1, skip)
+                            .Seal()
+                        .Seal()
+                        .Add(1, uintMax)
+                        .Callable(2, "+")
+                            .Add(0, take)
+                            .Add(1, skip)
+                        .Seal()
+                    .Seal()
+                    .Build()
+            );
+
         } else {
             limitValues.push_back(take);
         }
diff --git a/yt/yql/tests/sql/suites/limit/dynamic_limit_offset_overflow.sql b/yt/yql/tests/sql/suites/limit/dynamic_limit_offset_overflow.sql
new file mode 100644
index 0000000000..5452aceb1c
--- /dev/null
+++ b/yt/yql/tests/sql/suites/limit/dynamic_limit_offset_overflow.sql
@@ -0,0 +1,14 @@
+-- YQL-19579
+-- Check that offset + limit don't overflow max uin64
+use plato;
+
+$limit = -1;
+$offset = 2;
+$limit = if($limit >= 0, cast($limit as uint64));
+$offset = if($offset >= 0, cast($offset as uint64));
+
+$i = select distinct key from Input;
+
+select * from $i order by key
+limit $limit offset $offset;
+