From d60929066b1fec97d338fbd121cca3aac9cc7336 Mon Sep 17 00:00:00 2001 From: msaf Date: Fri, 9 Oct 2026 10:21:47 +0100 Subject: [PATCH] fix: evaluate small window partitions per batch in PartitionAggregateWindowExec With many tiny window partitions and a frame that is not constant within a partition (for example FIRST_VALUE(x) IGNORE NULLS OVER (PARTITION BY id ORDER BY ts ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)), every partition went through the buffered reverse pass and grew and freed three memory reservations. Behind Comet's JNI memory pool each of those calls takes Spark's memory manager lock; a stage with ~2.3 rows per partition ran 4.5x slower than with DataFusion's WindowAggExec. * Partitions that begin and end inside one input batch are now evaluated together for every supported expression except RANGE frames starting at an offset: first/last/nth_value (with and without IGNORE NULLS) and aggregates over suffix frames, ntile, percent_rank and cume_dist, without buffering or reserving the rows. Partitions crossing batch boundaries keep the buffered, spilling path. * Reservations grow in 1 MiB chunks and keep up to one chunk between partitions; unused reserved memory is returned before spilling, and spilling frees the reservation as before. Micro-benchmark (2M partitions, ~2.25 rows each): 6.0 s and 8.0M pool calls before, 0.06 s and 1.3K pool calls after; WindowAggExec takes 2.0 s. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../operators/partition_aggregate_window.rs | 700 ++++++++++++++++-- 1 file changed, 637 insertions(+), 63 deletions(-) diff --git a/native/core/src/execution/operators/partition_aggregate_window.rs b/native/core/src/execution/operators/partition_aggregate_window.rs index dc6bcfa9ed2..48e9dc58f1b 100644 --- a/native/core/src/execution/operators/partition_aggregate_window.rs +++ b/native/core/src/execution/operators/partition_aggregate_window.rs @@ -125,9 +125,16 @@ impl Spec { matches!(self.kind, Kind::CumeDist | Kind::Suffix { .. }) } - /// Whether the value is the same for every row of the partition. - fn constant(&self) -> bool { - matches!(self.kind, Kind::Aggregate(_) | Kind::Value { .. }) + /// Whether whole partitions within one batch can be evaluated together: everything but + /// RANGE frames starting at an offset from the current row. + fn batched(&self) -> bool { + !matches!( + self.kind, + Kind::Suffix { + start: FrameStart::Range { delta: Some(_), .. }, + .. + } + ) } } @@ -136,8 +143,8 @@ impl Spec { enum Pending { /// Rows of one partition that may extend over other batches, processed row by row. Rows(Vec, RecordBatch), - /// Whole partitions, all within this batch, at `ranges` of it. Evaluated at once when - /// every expression is constant within a partition. + /// Whole partitions, all within this batch, at `ranges` of it. Evaluated together when + /// every expression supports it. Partitions(RecordBatch, Vec>), } @@ -441,18 +448,23 @@ impl ExecutionPlan for PartitionAggregateWindowExec { suffix_files: vec![], cursors: vec![], empty: None, - reservation: MemoryConsumer::new("WindowSuffix") - .with_can_spill(true) - .register(&runtime.memory_pool), + reservation: ChunkedReservation::new( + MemoryConsumer::new("WindowSuffix") + .with_can_spill(true) + .register(&runtime.memory_pool), + ), } }); let state = WindowState { spill: SpillManager::new(Arc::clone(&runtime), spill_metrics, input.schema()), - rows_reservation: MemoryConsumer::new("WindowRows") - .with_can_spill(true) - .register(&runtime.memory_pool), - state_reservation: MemoryConsumer::new("WindowAccumulator") - .register(&runtime.memory_pool), + rows_reservation: ChunkedReservation::new( + MemoryConsumer::new("WindowRows") + .with_can_spill(true) + .register(&runtime.memory_pool), + ), + state_reservation: ChunkedReservation::new( + MemoryConsumer::new("WindowAccumulator").register(&runtime.memory_pool), + ), baseline: BaselineMetrics::new(&self.metrics, partition), input, input_done: false, @@ -461,7 +473,7 @@ impl ExecutionPlan for PartitionAggregateWindowExec { keys: self.window.partition_by_sort_keys()?, order_by, pending: VecDeque::new(), - constant: self.specs.iter().all(Spec::constant), + batched: self.specs.iter().all(Spec::batched), target_rows: context.session_config().batch_size().max(1), buffered: vec![], buffered_rows: 0, @@ -874,6 +886,62 @@ impl SuffixState { } } +/// Memory reservations are taken from the pool in chunks and kept, up to one chunk, from +/// one window partition to the next: with Spark's memory manager behind the pool, a call +/// per small partition contends on its lock. +const RESERVATION_CHUNK: usize = 1 << 20; + +#[derive(Debug)] +struct ChunkedReservation { + reservation: MemoryReservation, + used: usize, +} + +impl ChunkedReservation { + fn new(reservation: MemoryReservation) -> Self { + Self { + reservation, + used: 0, + } + } + + fn try_resize(&mut self, used: usize) -> Result<()> { + if used > self.reservation.size() + && self + .reservation + .try_resize(used.next_multiple_of(RESERVATION_CHUNK)) + .is_err() + { + self.reservation.try_resize(used)?; + } + self.used = used; + Ok(()) + } + + fn try_grow(&mut self, additional: usize) -> Result<()> { + self.try_resize(self.used + additional) + } + + /// The tracked memory has been released; keeps up to a chunk for the next partition. + fn release(&mut self) { + self.used = 0; + if self.reservation.size() > RESERVATION_CHUNK { + self.reservation.resize(RESERVATION_CHUNK); + } + } + + fn trim(&mut self) { + if self.reservation.size() > self.used { + self.reservation.resize(self.used); + } + } + + fn free(&mut self) { + self.used = 0; + self.reservation.free(); + } +} + #[derive(Clone)] enum SuffixSource { Memory(RecordBatch), @@ -1055,7 +1123,7 @@ struct ReverseState { cursors: Vec, /// One-row suffix batch for frames starting past the partition end. empty: Option, - reservation: MemoryReservation, + reservation: ChunkedReservation, } impl ReverseState { @@ -1104,7 +1172,7 @@ impl ReverseState { self.suffix_files.clear(); self.cursors.clear(); self.empty = None; - self.reservation.free(); + self.reservation.release(); } } @@ -1124,8 +1192,8 @@ struct WindowState { rows: Vec, files: VecDeque>, spill: SpillManager, - rows_reservation: MemoryReservation, - state_reservation: MemoryReservation, + rows_reservation: ChunkedReservation, + state_reservation: ChunkedReservation, replay: Option, result: Vec>, emitting: bool, @@ -1134,9 +1202,8 @@ struct WindowState { /// ORDER BY key and start index of the current percent_rank peer group. rank: Option<(Vec, usize)>, reverse: Option, - /// Every expression is constant within a partition, so partitions inside one input batch - /// are evaluated together. - constant: bool, + /// Partitions inside one input batch are evaluated together. + batched: bool, /// Output batches smaller than half of this are concatenated up to it. target_rows: usize, buffered: Vec, @@ -1146,6 +1213,15 @@ struct WindowState { } impl WindowState { + /// Returns reserved but unused memory to the pool before spilling. + fn trim(&mut self, reverse: Option<&mut ReverseState>) { + self.rows_reservation.trim(); + self.state_reservation.trim(); + if let Some(reverse) = reverse.or(self.reverse.as_mut()) { + reverse.reservation.trim(); + } + } + fn spill_rows(&mut self) -> Result<()> { if let Some(reverse) = &mut self.reverse { reverse.spill_rows(&self.rows)?; @@ -1231,11 +1307,19 @@ impl WindowState { } let state_size = self.state_size(); if self.state_reservation.try_resize(state_size).is_err() { - self.spill_rows()?; - self.state_reservation.try_resize(state_size)?; + self.trim(None); + if self.state_reservation.try_resize(state_size).is_err() { + self.spill_rows()?; + self.state_reservation.try_resize(state_size)?; + } } let size = batch.get_array_memory_size(); if self.rows_reservation.try_grow(size).is_err() { + self.trim(None); + if self.rows_reservation.try_grow(size).is_ok() { + self.rows.push(batch); + return Ok(()); + } self.spill_rows()?; // A single input batch may itself exceed the share. Write it directly, without // retaining it or claiming an unbounded memory reservation. @@ -1287,6 +1371,10 @@ impl WindowState { *end -= rows; let size = self.state_size() + states.iter().map(|s| s.size()).sum::(); if self.state_reservation.try_resize(size).is_err() { + self.trim(Some(reverse)); + if self.state_reservation.try_resize(size).is_ok() { + return reverse.push(RecordBatch::try_new(schema, columns)?); + } reverse.flush()?; if self.state_reservation.try_resize(size).is_err() { // The reverse pass works on its own copy of in-memory rows. @@ -1390,7 +1478,7 @@ impl WindowState { self.spill_rows()?; } self.reverse_pass().await?; - self.state_reservation.free(); + self.state_reservation.release(); if self.files.is_empty() { let batches = std::mem::take(&mut self.rows); self.replay = Some(Box::pin(RecordBatchStreamAdapter::new( @@ -1480,8 +1568,8 @@ impl WindowState { Ok(columns) } - /// Queues `batch`, split at its partition `ranges`. When every expression is constant - /// within a partition, the partitions that begin and end inside the batch are queued + /// Queues `batch`, split at its partition `ranges`. When every expression supports it, + /// the partitions that begin and end inside the batch are queued /// together; only the first, which may continue the partition in progress, and the last, /// which may continue into the next batch, are processed row by row. fn split( @@ -1501,7 +1589,7 @@ impl WindowState { batch.slice(range.start, range.end - range.start), )) }; - if !self.constant || ranges.len() < 2 { + if !self.batched || ranges.len() < 2 { for range in &ranges { let pending = rows(range)?; self.pending.push_back(pending); @@ -1535,56 +1623,132 @@ impl WindowState { Ok(()) } - /// Output rows of whole partitions at `ranges` of `batch`, with the value of every - /// expression computed once per partition. + /// Output rows of whole partitions at `ranges` of `batch`, each evaluated with the rows of + /// the batch, without buffering or reserving them. fn evaluate_partitions( &self, batch: &RecordBatch, ranges: &[Range], ) -> Result { - let mut indices = Vec::with_capacity(batch.num_rows()); - for (i, range) in ranges.iter().enumerate() { - indices.extend(std::iter::repeat_n(i as u32, range.len())); + let rows = batch.num_rows(); + let mut partition = Vec::with_capacity(rows); + for range in ranges { + partition.extend(std::iter::repeat_n(range.clone(), range.len())); } - let indices = UInt32Array::from(indices); + let peers = if self.specs.iter().any(|s| { + matches!( + s.kind, + Kind::PercentRank + | Kind::CumeDist + | Kind::Suffix { + start: FrameStart::Range { .. }, + .. + } + ) + }) { + let order = self + .order_by + .iter() + .map(|o| o.evaluate_to_sort_column(batch)) + .collect::>>()?; + let mut peers = Vec::with_capacity(rows); + for range in evaluate_partition_ranges(rows, &order)? { + let mut start = range.start; + while start < range.end { + let end = partition[start].end.min(range.end); + peers.extend(std::iter::repeat_n(start..end, end - start)); + start = end; + } + } + peers + } else { + vec![] + }; let mut columns = batch.columns().to_vec(); for spec in &self.specs { let args = spec .args .iter() - .map(|e| e.evaluate(batch)?.into_array(batch.num_rows())) + .map(|e| e.evaluate(batch)?.into_array(rows)) .collect::>>()?; - let slice = |range: &Range| { - args.iter() - .map(|a| a.slice(range.start, range.len())) - .collect::>() - }; - let mut values = Vec::with_capacity(ranges.len()); - for range in ranges { - values.push(match &spec.kind { - Kind::Aggregate(aggregate) => { - let mut accumulator = aggregate.create_accumulator()?; - accumulator.update_batch(&slice(range))?; - accumulator.evaluate()? + let column: ArrayRef = match &spec.kind { + Kind::Aggregate(_) | Kind::Value { .. } => { + let mut values = Vec::with_capacity(ranges.len()); + let mut indices = Vec::with_capacity(rows); + for (i, range) in ranges.iter().enumerate() { + values.push(constant_value(spec, &args, range)?); + indices.extend(std::iter::repeat_n(i as u32, range.len())); } - Kind::Value { kind, ignore_nulls } => { - let mut value = ValueState { - kind: *kind, - ignore_nulls: *ignore_nulls, - seen: 0, - value: None, - }; - value.update(&slice(range)[0])?; - match value.value { - Some(v) => v, - None => ScalarValue::try_from(&spec.data_type)?, + let values = ScalarValue::iter_to_array(values)?; + take(values.as_ref(), &UInt32Array::from(indices), None)? + } + Kind::Ntile(n) => Arc::new(UInt64Array::from_iter_values( + partition + .iter() + .enumerate() + .map(|(row, p)| ntile(row - p.start, *n, p.len())), + )), + Kind::PercentRank => Arc::new(Float64Array::from_iter_values( + partition.iter().zip(&peers).map(|(p, peer)| { + (peer.start - p.start) as f64 / (p.len() as f64 - 1.0).max(1.0) + }), + )), + Kind::CumeDist => Arc::new(Float64Array::from_iter_values( + partition + .iter() + .zip(&peers) + .map(|(p, peer)| (peer.end - p.start) as f64 / p.len() as f64), + )), + Kind::Suffix { func, start } => { + let starts = (0..rows).map(|row| { + let p = &partition[row]; + match start { + FrameStart::Rows(delta) => (row as i64) + .saturating_add(*delta) + .clamp(p.start as i64, p.end as i64) + as usize, + FrameStart::Range { .. } => peers[row].start, + } + }); + match func { + SuffixFn::Value { kind, ignore_nulls } => { + let indices = suffix_value_indices( + &args[0], + &partition, + starts, + *kind, + *ignore_nulls, + ); + take(args[0].as_ref(), &UInt32Array::from(indices), None)? + } + SuffixFn::Aggregate(aggregate) => { + let mut suffix = vec![ScalarValue::Null; rows]; + for range in ranges { + let mut accumulator = aggregate.create_accumulator()?; + for row in range.clone().rev() { + let slice = + args.iter().map(|a| a.slice(row, 1)).collect::>(); + accumulator.update_batch(&slice)?; + suffix[row] = accumulator.evaluate()?; + } + } + let empty = empty_value(spec)?; + let values = starts + .zip(&partition) + .map(|(start, p)| { + if start < p.end { + suffix[start].clone() + } else { + empty.clone() + } + }) + .collect::>(); + ScalarValue::iter_to_array(values)? } } - _ => return Err(internal_datafusion_err!("not a constant window expression")), - }); - } - let values = ScalarValue::iter_to_array(values)?; - columns.push(take(values.as_ref(), &indices, None)?); + } + }; + columns.push(cast_to(column, &spec.data_type)?); } Ok(RecordBatch::try_new(Arc::clone(&self.schema), columns)?) } @@ -1657,7 +1821,7 @@ impl WindowState { self.replay = Some(self.spill.read_spill_as_stream_unbuffered(file, None)?); continue; } - self.rows_reservation.free(); + self.rows_reservation.release(); if let Some(reverse) = &mut self.reverse { reverse.clear(); } @@ -1723,6 +1887,112 @@ impl WindowState { } } +/// Value of a whole-partition aggregate or `first_value`/`last_value`/`nth_value` over the +/// rows of `args` at `range`. +fn constant_value(spec: &Spec, args: &[ArrayRef], range: &Range) -> Result { + let slice = args + .iter() + .map(|a| a.slice(range.start, range.len())) + .collect::>(); + match &spec.kind { + Kind::Aggregate(aggregate) => { + let mut accumulator = aggregate.create_accumulator()?; + accumulator.update_batch(&slice)?; + accumulator.evaluate() + } + Kind::Value { kind, ignore_nulls } => { + let mut value = ValueState { + kind: *kind, + ignore_nulls: *ignore_nulls, + seen: 0, + value: None, + }; + value.update(&slice[0])?; + match value.value { + Some(v) => Ok(v), + None => ScalarValue::try_from(&spec.data_type), + } + } + _ => Err(internal_datafusion_err!("not a constant window expression")), + } +} + +/// Row of `values` selected by `first_value`/`last_value`/`nth_value` over the frame from +/// each row's frame start in `starts` to the end of its partition, or `None` for NULL. +fn suffix_value_indices( + values: &ArrayRef, + partition: &[Range], + starts: impl Iterator, + kind: ValueKind, + ignore_nulls: bool, +) -> Vec> { + let rows = values.len(); + let valid = is_valid(values); + let found = |index: usize, end: usize| (index < end).then_some(index as u32); + if !ignore_nulls { + return starts + .zip(partition) + .map(|(start, p)| match kind { + _ if start >= p.end => None, + ValueKind::First => found(start, p.end), + ValueKind::Last => found(p.end - 1, p.end), + ValueKind::Nth(n) => found(start.saturating_add(n - 1), p.end), + }) + .collect(); + } + match kind { + ValueKind::First => { + let mut next = vec![rows; rows + 1]; + for row in (0..rows).rev() { + next[row] = if valid(row) { + row + } else if row + 1 < partition[row].end { + next[row + 1] + } else { + rows + }; + } + starts + .zip(partition) + .map(|(start, p)| found(next[start], p.end)) + .collect() + } + ValueKind::Last => { + let mut last = vec![None; rows]; + let mut row = 0; + while row < rows { + let p = &partition[row]; + let value = p.clone().rev().find(|&i| valid(i)); + last[p.clone()].fill(value); + row = p.end; + } + starts + .zip(&last) + .map(|(start, last)| last.filter(|&l| l >= start).map(|l| l as u32)) + .collect() + } + ValueKind::Nth(n) => { + let mut before = Vec::with_capacity(rows + 1); + let mut positions = vec![]; + for row in 0..rows { + before.push(positions.len()); + if valid(row) { + positions.push(row); + } + } + before.push(positions.len()); + starts + .zip(partition) + .map(|(start, p)| { + positions + .get(before[start].saturating_add(n - 1)) + .and_then(|&index| found(index, p.end)) + }) + .collect() + } + } +} + /// SQL NTILE: with `base = num_rows / n`, the first `num_rows % n` buckets hold `base + 1` /// rows and the rest hold `base` rows (matches DataFusion's and Spark's bucket sizes). fn ntile(row: usize, n: u64, num_rows: usize) -> u64 { @@ -2527,6 +2797,310 @@ mod tests { Ok(()) } + /// Memory pool that counts the calls reaching it, as each reaches Spark's memory manager + /// through JNI in Comet. + #[derive(Debug)] + struct CountingPool { + inner: GreedyMemoryPool, + calls: std::sync::atomic::AtomicUsize, + } + + impl std::fmt::Display for CountingPool { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CountingPool") + } + } + + impl MemoryPool for CountingPool { + fn name(&self) -> &str { + "counting" + } + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.calls + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.inner.grow(reservation, additional) + } + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.calls + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.inner.shrink(reservation, shrink) + } + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.calls + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.inner.try_grow(reservation, additional) + } + fn reserved(&self) -> usize { + self.inner.reserved() + } + } + + /// Rows of partitions with `sizes`, sorted by key and `ord`, with NULL values, NULL and + /// tied ORDER BY values, and a NULL key for the first partition. + fn layout_rows(sizes: &[usize], seed: u64) -> Vec<(Option, Option, Option)> { + let mut state = seed; + let mut next = move |n: u64| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) % n + }; + let mut rows = vec![]; + for (p, &size) in sizes.iter().enumerate() { + let mut ord = 0i64; + for _ in 0..size { + let key = (p > 0).then_some(p as i64); + let null_ord = next(10) == 0; + ord += next(3) as i64; + let value = (next(3) != 0).then(|| next(41) as i64 - 20); + rows.push((key, (!null_ord).then_some(ord), value)); + } + } + rows + } + + fn batched( + rows: &[(Option, Option, Option)], + chunk: usize, + ) -> Result> { + let mut rows = rows.to_vec(); + let options = [SortOptions { + descending: false, + nulls_first: true, + }; 2]; + rows.sort_by(|a, b| { + compare_rows( + &[ScalarValue::Int64(a.0), ScalarValue::Int64(a.1)], + &[ScalarValue::Int64(b.0), ScalarValue::Int64(b.1)], + &options, + ) + .unwrap() + }); + let batch = RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int64Array::from_iter(rows.iter().map(|r| r.0))), + Arc::new(Int64Array::from_iter(rows.iter().map(|r| r.1))), + Arc::new(Int64Array::from_iter(rows.iter().map(|r| r.2))), + Arc::new(StringArray::from(vec!["p"; rows.len()])), + ], + )?; + let batches = (0..rows.len()) + .step_by(chunk) + .map(|start| batch.slice(start, chunk.min(rows.len() - start))) + .collect::>(); + let ordering = LexOrdering::new(vec![sort("key", false), sort("ord", false)]); + let config = MemorySourceConfig::try_new(&[batches], schema(), None)? + .try_with_sort_information(vec![ordering.unwrap()])?; + Ok(Arc::new(DataSourceExec::new(Arc::new(config)))) + } + + fn every_expression(range_offset: bool) -> Vec { + let n = |n: i64| lit(ScalarValue::Int64(Some(n))); + let mut frames = vec![rows_from(0), rows_from(-2), rows_from(3), range_from(None)]; + if range_offset { + frames.push(range_from(Some(2))); + } + let mut exprs = vec![ + expr("sum", vec![col("value")], whole()), + expr("count", vec![col("value")], whole()), + ignoring_nulls(expr("last_value", vec![col("value")], whole())), + expr("nth_value", vec![col("value"), n(2)], whole()), + expr("ntile", vec![n(3)], running()), + expr("percent_rank", vec![], running()), + expr("cume_dist", vec![], running()), + ]; + for frame in frames { + exprs.push(expr("sum", vec![col("value")], frame.clone())); + exprs.push(expr("count", vec![col("value")], frame.clone())); + exprs.push(expr("max", vec![col("value")], frame.clone())); + for ignore in [false, true] { + let with = |e: Expr| if ignore { ignoring_nulls(e) } else { e }; + exprs.push(with(expr("first_value", vec![col("value")], frame.clone()))); + exprs.push(with(expr("last_value", vec![col("value")], frame.clone()))); + exprs.push(with(expr( + "nth_value", + vec![col("value"), n(2)], + frame.clone(), + ))); + exprs.push(with(expr( + "nth_value", + vec![col("value"), n(4)], + frame.clone(), + ))); + } + } + exprs + } + + /// Compares with `WindowAggExec` over partition size distributions and batch sizes that + /// put partitions within batches, across batch boundaries and spanning several batches, + /// without and with spilling. + #[tokio::test] + async fn partition_layouts_match_window_agg_exec() -> Result<()> { + let layouts: Vec> = vec![ + vec![1; 300], + vec![2; 150], + [1, 2, 3, 1, 4, 2, 1, 1, 6].repeat(25), + [1, 1, 2, 90, 1, 3, 1, 250, 2, 1, 7].repeat(4), + vec![700], + [3, 1, 1, 2].repeat(60).into_iter().chain([400]).collect(), + ]; + for range_offset in [false, true] { + let exprs = every_expression(range_offset); + let window = build(&exprs, true, false)?; + let ignore_nulls = exprs.iter().map(|e| e.ignore_nulls).collect::>(); + for (l, sizes) in layouts.iter().enumerate() { + let rows = layout_rows(sizes, l as u64 + 1); + for chunk in [1, 3, 7, 64, 1000] { + let input = batched(&rows, chunk)?; + let reference: Arc = Arc::new(WindowAggExec::try_new( + window.clone(), + Arc::clone(&input), + true, + )?); + let (ctx, _) = context(LARGE)?; + let expected = concat_batches( + &reference.schema(), + &datafusion::physical_plan::collect(reference, ctx.task_ctx()).await?, + )?; + let plan = PartitionAggregateWindowExec::try_plan( + window.clone(), + input, + true, + ignore_nulls.clone(), + )? + .expect("spilling window plan"); + for budget in [LARGE, 16_000] { + let (actual, _) = run(&plan, budget).await?; + assert_eq!(actual.num_rows(), rows.len()); + for (i, field) in expected.schema().fields().iter().enumerate() { + assert_eq!( + actual.column(i).as_ref(), + expected.column(i).as_ref(), + "column {i} {} layout={l} chunk={chunk} budget={budget} \ + range_offset={range_offset}", + field.name() + ); + } + } + } + } + } + Ok(()) + } + + /// A suffix frame over many small partitions: the memory pool sees a bounded number of + /// calls instead of several per partition, and a partition larger than the budget still + /// spills. + #[tokio::test] + async fn small_partitions_do_not_call_the_pool_per_partition() -> Result<()> { + let exprs = [ignoring_nulls(expr( + "first_value", + vec![col("value")], + rows_from(0), + ))]; + let window = build(&exprs, true, false)?; + let sizes = [1, 2, 3, 2, 4, 1, 3].repeat(1500); + let rows = layout_rows(&sizes, 7); + let plan = PartitionAggregateWindowExec::try_plan( + window.clone(), + batched(&rows, 1000)?, + true, + vec![true], + )? + .expect("spilling window plan"); + let pool = Arc::new(CountingPool { + inner: GreedyMemoryPool::new(LARGE), + calls: Default::default(), + }); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build()?, + ); + let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let output = datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()).await?; + assert_eq!( + output.iter().map(|b| b.num_rows()).sum::(), + rows.len() + ); + let calls = pool.calls.load(std::sync::atomic::Ordering::Relaxed); + let batches = rows.len().div_ceil(1000); + assert!( + calls <= 8 * batches, + "{calls} pool calls for {} partitions in {batches} batches", + sizes.len() + ); + assert_eq!(pool.reserved(), 0); + + let large = layout_rows(&[1, 2, 3000, 1, 2], 3); + let plan = PartitionAggregateWindowExec::try_plan( + window, + batched(&large, 100)?, + true, + vec![true], + )? + .expect("spilling window plan"); + let (_, spills) = run(&plan, 16_000).await?; + assert!(spills > 0); + Ok(()) + } + + /// Measures `FIRST_VALUE(x) IGNORE NULLS` over `ROWS BETWEEN CURRENT ROW AND UNBOUNDED + /// FOLLOWING` with about 2.3 rows per partition, against DataFusion's `WindowAggExec`. + #[tokio::test] + #[ignore] + async fn bench_small_partitions_suffix_frame() -> Result<()> { + let exprs = [ignoring_nulls(expr( + "first_value", + vec![col("value")], + rows_from(0), + ))]; + let window = build(&exprs, true, false)?; + let sizes = [1, 2, 3, 2, 4, 1, 3, 2].repeat(250_000); + let rows = layout_rows(&sizes, 11); + let input = batched(&rows, 8192)?; + let plans: Vec<(&str, Arc)> = vec![ + ( + "PartitionAggregateWindowExec", + PartitionAggregateWindowExec::try_plan( + window.clone(), + Arc::clone(&input), + true, + vec![true], + )? + .expect("planned"), + ), + ( + "WindowAggExec", + Arc::new(WindowAggExec::try_new(window, input, true)?), + ), + ]; + for (name, plan) in plans { + let pool = Arc::new(CountingPool { + inner: GreedyMemoryPool::new(usize::MAX / 2), + calls: Default::default(), + }); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build()?, + ); + let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let started = std::time::Instant::now(); + let output = datafusion::physical_plan::collect(plan, ctx.task_ctx()).await?; + println!( + "BENCH {name}: {:?}, {} rows, {} partitions, {} pool calls", + started.elapsed(), + output.iter().map(|b| b.num_rows()).sum::(), + sizes.len(), + pool.calls.load(std::sync::atomic::Ordering::Relaxed) + ); + } + Ok(()) + } + fn col_in(name: &str, schema: &SchemaRef) -> Arc { datafusion::physical_expr::expressions::col(name, schema).unwrap() }