diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index acba75268b50d..354d0938cb744 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -258,8 +258,9 @@ impl WindowFrameContext { "Frame start cannot be UNBOUNDED FOLLOWING '{window_frame:?}'" ); } + // `n` comes from the query and can be as large as `u64::MAX` WindowFrameBound::Following(ScalarValue::UInt64(Some(n))) => { - std::cmp::min(idx + n as usize, length) + std::cmp::min(idx.saturating_add(n as usize), length) } // ERRONEOUS FRAMES WindowFrameBound::Preceding(_) | WindowFrameBound::Following(_) => { @@ -284,7 +285,7 @@ impl WindowFrameContext { // UNBOUNDED FOLLOWING WindowFrameBound::Following(ScalarValue::UInt64(None)) => length, WindowFrameBound::Following(ScalarValue::UInt64(Some(n))) => { - std::cmp::min(idx + n as usize + 1, length) + std::cmp::min(idx.saturating_add(n as usize).saturating_add(1), length) } // ERRONEOUS FRAMES WindowFrameBound::Preceding(_) | WindowFrameBound::Following(_) => { diff --git a/datafusion/physical-optimizer/src/limit_pushdown_past_window.rs b/datafusion/physical-optimizer/src/limit_pushdown_past_window.rs index 092570b051979..1f50e8f3d2c51 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown_past_window.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown_past_window.rs @@ -177,7 +177,10 @@ fn grow_limit(window: &BoundedWindowAggExec, ctx: &mut TraverseState) -> bool { } // finish grow - ctx.max_lookahead(ctx.lookahead + max_rel); + // Frame bounds come straight from the query and can be as large as + // `u64::MAX`, so saturate instead of overflowing. A saturated limit is + // still correct: it only means the sort below keeps every row. + ctx.max_lookahead(ctx.lookahead.saturating_add(max_rel)); true } @@ -194,16 +197,21 @@ fn apply_limit( ctx.lookahead = 0; return Some(Transformed::no(Arc::clone(node))); }; + let fetch = fetch.saturating_add(ctx.lookahead); let fetch = match node.fetch() { - None => fetch + ctx.lookahead, - Some(existing) => cmp::min(existing, fetch + ctx.lookahead), + None => fetch, + Some(existing) => cmp::min(existing, fetch), }; Some(Transformed::complete(node.with_fetch(Some(fetch)).unwrap())) } fn get_limit(node: &Arc, ctx: &mut TraverseState) -> bool { if let Some(limit) = node.downcast_ref::() { - ctx.reset_limit(limit.fetch().map(|fetch| fetch + limit.skip())); + ctx.reset_limit( + limit + .fetch() + .map(|fetch| fetch.saturating_add(limit.skip())), + ); return true; } // In distributed execution, GlobalLimitExec becomes LocalLimitExec diff --git a/datafusion/sqllogictest/test_files/window_limits.slt b/datafusion/sqllogictest/test_files/window_limits.slt index 11c96904f4ea2..7100f14cbf2e3 100644 --- a/datafusion/sqllogictest/test_files/window_limits.slt +++ b/datafusion/sqllogictest/test_files/window_limits.slt @@ -862,3 +862,26 @@ physical_plan 03)----BoundedWindowAggExec: wdw=[lead(employees.salary,Int64(2)) ORDER BY [employees.empno ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lead(employees.salary,Int64(2)) ORDER BY [employees.empno ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 04)------SortExec: TopK(fetch=5), expr=[empno@0 ASC NULLS LAST], preserve_partitioning=[false] 05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100_with_dates.csv]]}, projection=[empno, salary], file_type=csv, has_header=true + +# A huge FOLLOWING bound must not overflow when the limit is grown by the +# window's lookahead (https://github.com/apache/datafusion/issues/26658) +statement ok +CREATE TABLE huge_frame_t(a INT, x INT) AS VALUES (1, 1), (2, 2), (3, 3), (4, 4), (5, 5); + +query II +SELECT a, sum(x) OVER (ORDER BY a ROWS BETWEEN CURRENT ROW AND 18446744073709551615 FOLLOWING) s +FROM huge_frame_t ORDER BY a LIMIT 2; +---- +1 15 +2 14 + +# Same, with a window function that adds its own lookahead on top of the frame +query II +SELECT a, lead(x, 1) OVER (ORDER BY a ROWS BETWEEN CURRENT ROW AND 18446744073709551615 FOLLOWING) l +FROM huge_frame_t ORDER BY a LIMIT 2; +---- +1 2 +2 3 + +statement ok +DROP TABLE huge_frame_t;