Repository navigation
BoundedWindowAggExec can emit batches larger than batch_size #26959
Description
Describe the bug
BoundedWindowAggExec emits every row that has become final in one RecordBatch, so its output can be larger than datafusion.execution.batch_size. Rows pile up whenever the frame or function needs rows past the current one: a long RANGE peer group, a ROWS ... N FOLLOWING frame, or lead(x, N). When the needed rows arrive, all the held rows come out together.
compute_aggregates builds the output from the first n_out rows of the input buffer and returns it as one batch, with no split by batch size (bounded_window_agg_exec.rs, lines 1305 to 1346 on main at 611311c).
#26772 and #26788 fix the same problem for WindowAggExec and grouped TopK, but they don't touch BoundedWindowAggExec. #19481 is the general issue.
To Reproduce
DataFusion 55.2.0, batch_size = 2, one partition, 10 rows. compute_aggregates is the same on main at 611311c.
use datafusion::physical_plan::collect;
use datafusion::prelude::{SessionConfig, SessionContext};
#[tokio::main]
async fn main() -> datafusion::common::Result<()> {
let config = SessionConfig::new().with_batch_size(2).with_target_partitions(1);
let ctx = SessionContext::new_with_config(config);
ctx.sql(
"CREATE TABLE t AS SELECT 1 AS k, CAST(v AS BIGINT) AS v \
FROM (VALUES (1),(2),(3),(4),(5),(6),(7),(8),(9),(10)) AS s(v)",
)
.await?
.collect()
.await?;
for query in [
"SELECT v, sum(v) OVER (ORDER BY k RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) FROM t",
"SELECT v, sum(v) OVER (ORDER BY v ROWS BETWEEN CURRENT ROW AND 6 FOLLOWING) FROM t",
"SELECT v, lead(v, 6) OVER (ORDER BY v) FROM t",
"SELECT v, sum(v) OVER (ORDER BY v ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) FROM t",
] {
let plan = ctx.sql(query).await?.create_physical_plan().await?;
let sizes: Vec<_> = collect(plan, ctx.task_ctx())
.await?
.iter()
.map(|b| b.num_rows())
.collect();
println!("{sizes:?} {query}");
}
Ok(())
}Every plan is BoundedWindowAggExec (mode=[Sorted]) over SortExec. The SortExec emits batches of 2 rows.
| Query | Output batch sizes |
|---|---|
RANGE frame over a single 10-row peer group |
[10] |
ROWS BETWEEN CURRENT ROW AND 6 FOLLOWING |
[1, 2, 7] |
lead(v, 6) |
[2, 2, 6] |
ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW (no lookahead) |
[2, 2, 2, 2, 2] |
Expected behavior
No output batch larger than batch_size, so [2, 2, 2, 2, 2] for every query above. The extra batches could be zero-copy RecordBatch::slice chunks, the way #26788 does it for WindowAggExec, and the way TopK::emit already splits its output.
Activity