diff options
author | udovichenko-r <udovichenko-r@yandex-team.com> | 2025-02-12 20:21:50 +0300 |
---|---|---|
committer | udovichenko-r <udovichenko-r@yandex-team.com> | 2025-02-12 20:55:08 +0300 |
commit | 2457666f4ce8a05e187b4995a25dcee52ba885e7 (patch) | |
tree | 536e7c8ba5cf13063d2853191515dc652187b955 | |
parent | 67fd8aea0244391fdda4240c7e97bd03a8edf6e5 (diff) | |
download | ydb-2457666f4ce8a05e187b4995a25dcee52ba885e7.tar.gz |
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; + |