aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorudovichenko-r <udovichenko-r@yandex-team.com>2025-02-12 20:21:50 +0300
committerudovichenko-r <udovichenko-r@yandex-team.com>2025-02-12 20:55:08 +0300
commit2457666f4ce8a05e187b4995a25dcee52ba885e7 (patch)
tree536e7c8ba5cf13063d2853191515dc652187b955
parent67fd8aea0244391fdda4240c7e97bd03a8edf6e5 (diff)
downloadydb-2457666f4ce8a05e187b4995a25dcee52ba885e7.tar.gz
Fix overflow in take + skip
commit_hash:2e37109d1766f0f7b76363f60d3e7217686eae48
-rw-r--r--yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp6
-rw-r--r--yt/yql/providers/yt/lib/res_pull/table_limiter.cpp7
-rw-r--r--yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp42
-rw-r--r--yt/yql/providers/yt/provider/yql_yt_helpers.cpp25
-rw-r--r--yt/yql/tests/sql/suites/limit/dynamic_limit_offset_overflow.sql14
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;
+