From 3cbaae8fb5e63e0b16e3f935b2af47df4706820b Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Tue, 4 Aug 2026 12:13:09 -0400 Subject: [PATCH] perf: skip evaluating fully calculated window partitions In Linear mode, BoundedWindowAggStream's evaluation sweep visits every live partition for every window expression on every input batch. A partition that received no new rows and already has a result for every buffered row cannot produce anything new can be safely skipped. This avoids a bunch of redundant work: re-evaluating the window function arguments and ORDER BY columns against the retained batch, building an empty result array, and other bookkeeping. This is particularly expensive for workloads with many partitions where only a few of those partitions receive rows in a given batch, as in the "32k sparse" benchmark below. Benchmarks: - linear / range / single / 100 dense: 42.3 ms -> 42.0 ms (~noise) - linear / range / single / 10000 dense: 158.7 ms -> 152.5 ms (-3.9%) - linear / range / single / 32768 sparse: 161.1 ms -> 108.0 ms (-33.0%) - linear / rows / single / 10000 dense: 132.0 ms -> 127.5 ms (-3.4%) - linear / range / multi / 10000 dense: 255.9 ms -> 236.4 ms (-7.6%) - sorted / range / single / 10000: 33.1 ms -> 33.7 ms (~noise) --- datafusion/expr/src/window_state.rs | 25 ++++ .../physical-expr/src/window/window_expr.rs | 6 + .../src/windows/bounded_window_agg_exec.rs | 115 ++++++++++++++++++ 3 files changed, 146 insertions(+) diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index ece07e5b09c4d..1fe5ea4791fe8 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -101,6 +101,31 @@ impl WindowAggState { Ok(()) } + /// Returns true when this state is fully up to date with the partition's + /// buffered batch, meaning another evaluation pass over the partition could + /// not produce any new results or change any state: + /// + /// - `last_calculated_index` has reached the end of the partition's + /// buffered batch, so every row of this partition that has arrived so + /// far already has a result. + /// - When a partition ends, a final evaluation pass is needed to bring + /// derived state up to date. + #[inline] + pub fn is_up_to_date_with( + &self, + partition_batch_state: &PartitionBatchState, + ) -> bool { + let all_rows_have_results = + self.last_calculated_index == partition_batch_state.record_batch.num_rows(); + if all_rows_have_results { + debug_assert_eq!(self.n_row_result_missing, 0); + } + + // `self.is_end` holds the flag as of the previous evaluation pass. + let partition_just_ended = !self.is_end && partition_batch_state.is_end; + all_rows_have_results && !partition_just_ended + } + pub fn new(out_type: &DataType) -> Result { let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?; Ok(Self { diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 47147b909d342..bbced4c25d0fc 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -263,6 +263,12 @@ pub trait AggregateWindowExpr: WindowExpr { let state = &mut window_state.state; let record_batch = &partition_batch_state.record_batch; + // Skip partitions that cannot produce anything new until they + // either receive rows or reach their end. + if state.is_up_to_date_with(partition_batch_state) { + continue; + } + // If there is no window state context, initialize it. let window_frame_ctx = state.window_frame_ctx.get_or_insert_with(|| { let sort_options = self.order_by().iter().map(|o| o.options).collect(); diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 03a8e9867c170..d197941c76ef6 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -1194,6 +1194,18 @@ impl BoundedWindowAggStream { /// Prunes the sections of the record batch (for each partition) /// that we no longer need to calculate the window function result. fn prune_partition_batches(&mut self) { + // Check that per-state and per-partition end-flags are consistent; + // otherwise, the pruning code below might produce inconsistent state. + #[cfg(debug_assertions)] + for window_agg_state in self.window_agg_states.iter() { + for (partition_row, WindowState { state, .. }) in window_agg_state.iter() { + debug_assert_eq!( + state.is_end, self.partition_buffers[partition_row].is_end, + "window state's recorded end flag is out of sync with its partition" + ); + } + } + // Remove partitions which we know already ended (is_end flag is true). // Since the retain method preserves insertion order, we still have // ordering in between partitions after removal. @@ -1395,6 +1407,7 @@ mod tests { WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, }; use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_functions_window::nth_value::last_value_udwf; use datafusion_functions_window::nth_value::nth_value_udwf; use datafusion_physical_expr::expressions::{Column, Literal, col}; @@ -1834,6 +1847,108 @@ mod tests { Ok(()) } + // In `Linear` mode, a partition may receive no new rows for several + // input batches while other partitions keep growing. Once all of a + // partition's buffered rows have results, the evaluation sweep skips + // it until it receives rows again, so this test drives a partition + // through quiet batches and then resumes it: the results after the + // gap must continue from the retained accumulator state. Both frames + // are causal, so results finalize in the batch their row arrives in + // and the quiet partition is fully calculated while it waits. + #[tokio::test] + async fn bounded_window_linear_quiet_partition_resume() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])); + let make_batch = |rows: &[(u64, u64)]| -> Result { + let mut pk = UInt64Builder::with_capacity(rows.len()); + let mut ts = UInt64Builder::with_capacity(rows.len()); + for (p, t) in rows { + pk.append_value(*p); + ts.append_value(*t); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(pk.finish()), Arc::new(ts.finish())], + )?) + }; + // `ts` ascends globally; partition 0 is absent from the middle batches. + let batches = vec![ + make_batch(&[(0, 0), (0, 1), (1, 2)])?, + make_batch(&[(1, 3), (1, 4)])?, + make_batch(&[(1, 5)])?, + make_batch(&[(0, 6), (1, 7)])?, + ]; + let memory_exec = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let partition_by = vec![col("pk", &schema)?]; + let order_by = [PhysicalSortExpr { + expr: col("ts", &schema)?, + options: SortOptions::default(), + }]; + // A running COUNT (plain aggregate) and a SUM over the previous and + // current row (sliding aggregate). + let count_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let sum_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let physical_plan = BoundedWindowAggExec::try_new( + vec![count_expr, sum_expr], + memory_exec, + InputOrderMode::Linear, + true, + ) + .map(|e| Arc::new(e) as Arc)?; + + let batches = collect(physical_plan.execute(0, task_context())?).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-------+-----+ + | pk | ts | count | sum | + +----+----+-------+-----+ + | 0 | 0 | 1 | 0 | + | 0 | 1 | 2 | 1 | + | 1 | 2 | 1 | 2 | + | 1 | 3 | 2 | 5 | + | 1 | 4 | 3 | 7 | + | 1 | 5 | 4 | 9 | + | 0 | 6 | 3 | 7 | + | 1 | 7 | 5 | 12 | + +----+----+-------+-----+ + "); + Ok(()) + } + // This test, tests whether most recent row guarantee by the input batch of the `BoundedWindowAggExec` // helps `BoundedWindowAggExec` to generate low latency result in the `Linear` mode. // Input data generated at the source is