Skip to content

Commit 2457666

Browse files
committed
Fix overflow in take + skip
commit_hash:2e37109d1766f0f7b76363f60d3e7217686eae48
1 parent 67fd8ae commit 2457666

File tree

5 files changed

+75
-19
lines changed

5 files changed

+75
-19
lines changed

yt/yql/providers/yt/lib/mkql_helpers/mkql_helpers.cpp

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,11 @@ void TRecordsRange::Fill(const TExprNode& settingsNode) {
6969
if (settingName != TStringBuf("take") && settingName != TStringBuf("skip")) {
7070
continue;
7171
}
72-
YQL_ENSURE(setting->Child(1)->IsCallable("Uint64"));
72+
if (!setting->Child(1)->IsCallable("Uint64")) {
73+
Offset.Clear();
74+
Limit.Clear();
75+
return;
76+
}
7377
if (!UpdateRecordsRange(*this, settingName, NYql::FromString<ui64>(*setting->Child(1)->Child(0), NUdf::EDataSlot::Uint64))) {
7478
break;
7579
}

yt/yql/providers/yt/lib/res_pull/table_limiter.cpp

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,16 @@ namespace NYql {
88

99
TTableLimiter::TTableLimiter(const TRecordsRange& range)
1010
: Start(range.Offset.GetOrElse(0ULL))
11-
, End(range.Limit.Defined() ? Start + *range.Limit : Max())
1211
, Current(0ULL)
1312
, TableStart(0ULL)
1413
, TableEnd(Max())
1514
{
15+
const auto limit = range.Limit.GetOrElse(Max());
16+
if (limit > Max<ui64>() - Start) {
17+
End = Max();
18+
} else {
19+
End = Start + limit;
20+
}
1621
}
1722

1823
bool TTableLimiter::NextTable(ui64 recordCount) {

yt/yql/providers/yt/provider/yql_yt_dq_hybrid.cpp

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -279,31 +279,41 @@ class TYtDqHybridTransformer : public TOptimizeTransformerBase {
279279
.Build()
280280
.Done();
281281

282-
TExprNode::TPtr limit;
283-
if (const auto& limitNode = NYql::GetSetting(sort.Settings().Ref(), EYtSettingType::Limit)) {
284-
limit = GetLimitExpr(limitNode, ctx);
285-
}
286-
282+
TExprNode::TPtr work;
287283
auto [direct, selector] = GetOutputSortSettings(sort, ctx);
288-
auto work = direct && selector ?
289-
limit ?
290-
Build<TCoTopSort>(ctx, sort.Pos())
284+
if (direct && selector) {
285+
// Don't use runtime limit for TopSort - it may have max<ui64>() value, which cause TopSort to fail
286+
TMaybe<ui64> limit = GetLimit(sort.Settings().Ref());
287+
work = limit
288+
? Build<TCoTopSort>(ctx, sort.Pos())
291289
.Input(input)
292-
.Count(std::move(limit))
290+
.Count<TCoUint64>()
291+
.Literal()
292+
.Value(ToString(*limit), TNodeFlags::Default)
293+
.Build()
294+
.Build()
293295
.SortDirections(std::move(direct))
294296
.KeySelectorLambda(std::move(selector))
295-
.Done().Ptr():
296-
Build<TCoSort>(ctx, sort.Pos())
297+
.Done().Ptr()
298+
: Build<TCoSort>(ctx, sort.Pos())
297299
.Input(input)
298300
.SortDirections(std::move(direct))
299301
.KeySelectorLambda(std::move(selector))
300-
.Done().Ptr():
301-
limit ?
302-
Build<TCoTake>(ctx, sort.Pos())
302+
.Done().Ptr()
303+
;
304+
} else {
305+
TExprNode::TPtr limit;
306+
if (const auto& limitNode = NYql::GetSetting(sort.Settings().Ref(), EYtSettingType::Limit)) {
307+
limit = GetLimitExpr(limitNode, ctx);
308+
}
309+
310+
work = limit
311+
? Build<TCoTake>(ctx, sort.Pos())
303312
.Input(input)
304313
.Count(std::move(limit))
305-
.Done().Ptr():
306-
input.Ptr();
314+
.Done().Ptr()
315+
: input.Ptr();
316+
}
307317

308318
auto settings = NYql::AddSetting(sort.Settings().Ref(), EYtSettingType::NoDq, {}, ctx);
309319
auto operation = ctx.ChangeChild(sort.Ref(), TYtTransientOpBase::idx_Settings, std::move(settings));

yt/yql/providers/yt/provider/yql_yt_helpers.cpp

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -905,7 +905,30 @@ TExprNode::TPtr GetLimitExpr(const TExprNode::TPtr& limitSetting, TExprContext&
905905
}
906906

907907
if (skip) {
908-
limitValues.push_back(ctx.NewCallable(child->Pos(), "+", { take, skip }));
908+
auto uintMax = ctx.Builder(child->Pos())
909+
.Callable("Uint64")
910+
.Atom(0, ToString(Max<ui64>()), TNodeFlags::Default)
911+
.Seal()
912+
.Build();
913+
limitValues.push_back(
914+
ctx.Builder(child->Pos())
915+
.Callable("If")
916+
.Callable(0, ">")
917+
.Add(0, take)
918+
.Callable(1, "-")
919+
.Add(0, uintMax)
920+
.Add(1, skip)
921+
.Seal()
922+
.Seal()
923+
.Add(1, uintMax)
924+
.Callable(2, "+")
925+
.Add(0, take)
926+
.Add(1, skip)
927+
.Seal()
928+
.Seal()
929+
.Build()
930+
);
931+
909932
} else {
910933
limitValues.push_back(take);
911934
}
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
-- YQL-19579
2+
-- Check that offset + limit don't overflow max uin64
3+
use plato;
4+
5+
$limit = -1;
6+
$offset = 2;
7+
$limit = if($limit >= 0, cast($limit as uint64));
8+
$offset = if($offset >= 0, cast($offset as uint64));
9+
10+
$i = select distinct key from Input;
11+
12+
select * from $i order by key
13+
limit $limit offset $offset;
14+

0 commit comments

Comments
 (0)