From 8bf70634b1678ea2f6b703bd8893390d6c9a99f5 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Thu, 30 Jul 2026 10:50:14 -0400 Subject: [PATCH 1/2] perf: track BoundedWindowAggExec Linear-mode watermark once per stream In Linear mode, rows arrive ordered on a prefix of the window's ORDER BY expressions, so each new row bounds every row that will arrive in the future, even for other partitions. We exploit this by using the last row to arrive in a batch to close window frames for all live partitions, not just the partition to which that row belongs. The most-recent-row-in-the-batch is conceptually per-batch state, but it was previously implemented as per-partition state: - update_partition_batch copied the last row of each incoming batch into every live partition's PartitionBatchState. - Each copy is ~(40 + 16 * n_cols) bytes (Option plus a Vec of ArrayRefs); for example, withn 100k live partitions over 10 columns, that is ~20MB of duplicate state. - Every partition visit re-evaluated the row's ORDER BY expressions. Instead, just store the row once. Rather than copying the row into every partition, we pass the row to window expression evaluation. This removes `most_recent_row` and its setter from `PartitionBatchState`, which is a breaking API change for datafusion-expr. We can also arrange to evaluate the ORDER BY once per batch instead of once per partition. This is a modest performance improvement and memory savings, but also a conceptual cleanup/refactor. Benchmark results (benches/bounded_window.rs): - linear 100 partitions: 44.4 ms -> 44.5 ms (within noise) - linear 10000 partitions: 203.3 ms -> 199.7 ms (-1.7%) - linear sparse 32768 partitions: 236.9 ms -> 224.5 ms (-5.2%) - linear rows 10000 partitions: 174.2 ms -> 169.4 ms (-2.4%) - linear multi 10000 partitions: 304.6 ms -> 295.7 ms (-3.0%) - sorted 10000 partitions: 34.2 ms -> 34.4 ms (within noise) --- datafusion/expr/src/window_state.rs | 13 ----- .../physical-expr/src/window/aggregate.rs | 7 ++- datafusion/physical-expr/src/window/mod.rs | 1 + .../src/window/sliding_aggregate.rs | 7 ++- .../physical-expr/src/window/standard.rs | 3 +- .../physical-expr/src/window/window_expr.rs | 55 ++++++++++++++++--- .../src/windows/bounded_window_agg_exec.rs | 52 +++++++++++------- 7 files changed, 90 insertions(+), 48 deletions(-) diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index f8d4609d3690c..ece07e5b09c4d 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -248,11 +248,6 @@ impl WindowFrameContext { pub struct PartitionBatchState { /// The record batch belonging to current partition pub record_batch: RecordBatch, - /// The record batch that contains the most recent row at the input. - /// Please note that this batch doesn't necessarily have the same partitioning - /// with `record_batch`. Keeping track of this batch enables us to prune - /// `record_batch` when cardinality of the partition is sparse. - pub most_recent_row: Option, /// Flag indicating whether we have received all data for this partition pub is_end: bool, /// Number of rows emitted for each partition @@ -263,7 +258,6 @@ impl PartitionBatchState { pub fn new(schema: SchemaRef) -> Self { Self { record_batch: RecordBatch::new_empty(schema), - most_recent_row: None, is_end: false, n_out_row: 0, } @@ -272,7 +266,6 @@ impl PartitionBatchState { pub fn new_with_batch(batch: RecordBatch) -> Self { Self { record_batch: batch, - most_recent_row: None, is_end: false, n_out_row: 0, } @@ -283,12 +276,6 @@ impl PartitionBatchState { concat_batches(&self.record_batch.schema(), [&self.record_batch, batch])?; Ok(()) } - - pub fn set_most_recent_row(&mut self, batch: RecordBatch) { - // It is enough for the batch to contain only a single row (the rest - // are not necessary). - self.most_recent_row = Some(batch); - } } /// This structure encapsulates all the state information we require as we scan diff --git a/datafusion/physical-expr/src/window/aggregate.rs b/datafusion/physical-expr/src/window/aggregate.rs index 1ff13d107c036..7cfdcb167f80a 100644 --- a/datafusion/physical-expr/src/window/aggregate.rs +++ b/datafusion/physical-expr/src/window/aggregate.rs @@ -23,7 +23,9 @@ use std::sync::Arc; use crate::aggregate::AggregateFunctionExpr; use crate::window::standard::add_new_ordering_expr_with_partition_by; -use crate::window::window_expr::{AggregateWindowExpr, WindowFn, filter_array}; +use crate::window::window_expr::{ + AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array, +}; use crate::window::{ PartitionBatches, PartitionWindowAggStates, SlidingAggregateWindowExpr, WindowExpr, }; @@ -148,8 +150,9 @@ impl WindowExpr for PlainAggregateWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { - self.aggregate_evaluate_stateful(partition_batches, window_agg_state)?; + self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx)?; // Update window frame range for each partition. As we know that // non-sliding aggregations will never call `retract_batch`, this value diff --git a/datafusion/physical-expr/src/window/mod.rs b/datafusion/physical-expr/src/window/mod.rs index b45e35440ac20..79b9a9580af89 100644 --- a/datafusion/physical-expr/src/window/mod.rs +++ b/datafusion/physical-expr/src/window/mod.rs @@ -28,5 +28,6 @@ pub use standard_window_function_expr::StandardWindowFunctionExpr; pub use window_expr::PartitionBatches; pub use window_expr::PartitionKey; pub use window_expr::PartitionWindowAggStates; +pub use window_expr::WindowEvalContext; pub use window_expr::WindowExpr; pub use window_expr::WindowState; diff --git a/datafusion/physical-expr/src/window/sliding_aggregate.rs b/datafusion/physical-expr/src/window/sliding_aggregate.rs index a71df3ec88472..29e569363ae2b 100644 --- a/datafusion/physical-expr/src/window/sliding_aggregate.rs +++ b/datafusion/physical-expr/src/window/sliding_aggregate.rs @@ -22,7 +22,9 @@ use std::ops::Range; use std::sync::Arc; use crate::aggregate::AggregateFunctionExpr; -use crate::window::window_expr::{AggregateWindowExpr, WindowFn, filter_array}; +use crate::window::window_expr::{ + AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array, +}; use crate::window::{ PartitionBatches, PartitionWindowAggStates, PlainAggregateWindowExpr, WindowExpr, }; @@ -102,8 +104,9 @@ impl WindowExpr for SlidingAggregateWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { - self.aggregate_evaluate_stateful(partition_batches, window_agg_state) + self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx) } fn partition_by(&self) -> &[Arc] { diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index 6f61174ee089b..2de080ec9a132 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -22,7 +22,7 @@ use std::ops::Range; use std::sync::Arc; use super::{StandardWindowFunctionExpr, WindowExpr}; -use crate::window::window_expr::{WindowFn, get_orderby_values}; +use crate::window::window_expr::{WindowEvalContext, WindowFn, get_orderby_values}; use crate::window::{PartitionBatches, PartitionWindowAggStates, WindowState}; use crate::{EquivalenceProperties, PhysicalExpr}; @@ -157,6 +157,7 @@ impl WindowExpr for StandardWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + _eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { let field = self.expr.field()?; let out_type = field.data_type(); diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 8db5651346e8f..47147b909d342 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -99,10 +99,14 @@ pub trait WindowExpr: Send + Sync + Debug { /// Evaluate the window function against the batch. This function facilitates /// stateful, bounded-memory implementations. + /// + /// `eval_ctx` carries stream-level (cross-partition) information; see + /// [`WindowEvalContext`]. fn evaluate_stateful( &self, _partition_batches: &PartitionBatches, _window_agg_state: &mut PartitionWindowAggStates, + _eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { internal_err!("evaluate_stateful is not implemented for {}", self.name()) } @@ -226,9 +230,18 @@ pub trait AggregateWindowExpr: WindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { let field = self.field()?; let out_type = field.data_type(); + // Every partition consults the same most recent input row, so its + // ORDER BY values can be evaluated once, outside the per-partition + // loop. + let most_recent_row_order_bys = eval_ctx + .most_recent_row + .map(|batch| self.order_by_columns(batch)) + .transpose()? + .map(get_orderby_values); for (partition_row, partition_batch_state) in partition_batches.iter() { if !window_agg_state.contains_key(partition_row) { let accumulator = self.get_accumulator()?; @@ -249,7 +262,6 @@ pub trait AggregateWindowExpr: WindowExpr { }; let state = &mut window_state.state; let record_batch = &partition_batch_state.record_batch; - let most_recent_row = partition_batch_state.most_recent_row.as_ref(); // If there is no window state context, initialize it. let window_frame_ctx = state.window_frame_ctx.get_or_insert_with(|| { @@ -259,7 +271,7 @@ pub trait AggregateWindowExpr: WindowExpr { let out_col = self.get_result_column( accumulator, record_batch, - most_recent_row, + most_recent_row_order_bys.as_deref(), // Start search from the last range &mut state.window_frame_range, window_frame_ctx, @@ -277,7 +289,8 @@ pub trait AggregateWindowExpr: WindowExpr { /// # Arguments /// * `accumulator`: The accumulator to use for the calculation. /// * `record_batch`: batch belonging to the current partition (see [`PartitionBatchState`]). - /// * `most_recent_row`: the batch that contains the most recent row, if available (see [`PartitionBatchState`]). + /// * `most_recent_row_order_bys`: ORDER BY values of the most recent input + /// row, if available (see [`WindowExpr::evaluate_stateful`]). /// * `last_range`: The last range of rows that were processed (see [`WindowAggState`]). /// * `window_frame_ctx`: Details about the window frame (see [`WindowFrameContext`]). /// * `idx`: The index of the current row in the record batch. @@ -287,7 +300,7 @@ pub trait AggregateWindowExpr: WindowExpr { &self, accumulator: &mut Box, record_batch: &RecordBatch, - most_recent_row: Option<&RecordBatch>, + most_recent_row_order_bys: Option<&[ArrayRef]>, last_range: &mut Range, window_frame_ctx: &mut WindowFrameContext, mut idx: usize, @@ -327,10 +340,6 @@ pub trait AggregateWindowExpr: WindowExpr { return value.to_array_of_size(record_batch.num_rows()); } let order_bys = get_orderby_values(self.order_by_columns(record_batch)?); - let most_recent_row_order_bys = most_recent_row - .map(|batch| self.order_by_columns(batch)) - .transpose()? - .map(get_orderby_values); // We iterate on each row to perform a running calculation. let length = values[0].len(); @@ -347,7 +356,7 @@ pub trait AggregateWindowExpr: WindowExpr { && !is_end_bound_safe( window_frame_ctx, &order_bys, - most_recent_row_order_bys.as_deref(), + most_recent_row_order_bys, self.order_by(), idx, )? @@ -605,6 +614,34 @@ pub enum WindowFn { /// PartitionKey would consist of unique `[a,b]` pairs pub type PartitionKey = Vec; +/// Stream-level context passed to [`WindowExpr::evaluate_stateful`]. +/// +/// This carries information that spans all partitions of the input, as +/// opposed to the per-partition state in [`PartitionBatches`] and +/// [`PartitionWindowAggStates`]. It is `non_exhaustive` so that fields can +/// be added without breaking implementors; construct it with +/// [`Default::default`] and the `with_*` builder methods. +#[derive(Debug, Clone, Copy, Default)] +#[non_exhaustive] +pub struct WindowEvalContext<'a> { + /// A single-row batch containing the most recent input row, whichever + /// partition that row belongs to. It is `Some` only when the input is + /// ordered by the first ORDER BY column across partitions (`Linear` + /// mode), in which case no future input row -- in any partition -- can + /// precede it in that column; implementations can use this bound to + /// decide whether pending window frames can be finalized before their + /// partition receives more data. + pub most_recent_row: Option<&'a RecordBatch>, +} + +impl<'a> WindowEvalContext<'a> { + /// Sets the most recent input row (see [`Self::most_recent_row`]). + pub fn with_most_recent_row(mut self, batch: Option<&'a RecordBatch>) -> Self { + self.most_recent_row = batch; + self + } +} + #[derive(Debug)] pub struct WindowState { pub state: WindowAggState, 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 07751a70eceeb..ef14dacd17a34 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -60,7 +60,8 @@ use datafusion_execution::TaskContext; use datafusion_expr::ColumnarValue; use datafusion_expr::window_state::{PartitionBatchState, WindowAggState}; use datafusion_physical_expr::window::{ - PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowState, + PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowEvalContext, + WindowState, }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ @@ -526,25 +527,6 @@ trait PartitionSearcher: Send { } } - if self.is_mode_linear() { - // In `Linear` mode, it is guaranteed that the first ORDER BY column - // is sorted across partitions. Note that only the first ORDER BY - // column is guaranteed to be ordered. As a counter example, consider - // the case, `PARTITION BY b, ORDER BY a, c` when the input is sorted - // by `[a, b, c]`. In this case, `BoundedWindowAggExec` mode will be - // `Linear`. However, we cannot guarantee that the last row of the - // input data will be the "last" data in terms of the ordering requirement - // `[a, c]` -- it will be the "last" data in terms of `[a, b, c]`. - // Hence, only column `a` should be used as a guarantee of the "last" - // data across partitions. For other modes (`Sorted`, `PartiallySorted`), - // we do not need to keep track of the most recent row guarantee across - // partitions. Since leading ordering separates partitions, guaranteed - // by the most recent row, already prune the previous partitions completely. - let last_row = get_last_row_batch(&record_batch)?; - for (_, partition_batch) in partition_buffers.iter_mut() { - partition_batch.set_most_recent_row(last_row.clone()); - } - } self.mark_partition_end(partition_buffers); *input_buffer = if input_buffer.num_rows() == 0 { @@ -1010,6 +992,24 @@ pub struct BoundedWindowAggStream { /// Search mode for partition columns. This determines the algorithm with /// which we group each partition. search_mode: Box, + /// In `Linear` mode, a single-row batch containing the most recent input + /// row (whichever partition that row belongs to); `None` in other modes + /// and before the first non-empty batch arrives. Since in `Linear` mode + /// the input is sorted by the first ORDER BY column, no future input row + /// -- in any partition -- can precede this row in that column. Every + /// partition's evaluation consults this bound to decide whether pending + /// window frames can be finalized before the partition receives more + /// data (which in turn allows buffered state to be pruned). Note that + /// only the first ORDER BY column provides this guarantee. As a counter + /// example, consider `PARTITION BY b, ORDER BY a, c` when the input is + /// sorted by `[a, b, c]`: the mode will be `Linear`, but the last row of + /// the input is the "last" data in terms of `[a, b, c]`, not in terms of + /// the ordering requirement `[a, c]`. Hence, only column `a` can serve + /// as a guarantee of the "last" data across partitions. In the `Sorted` + /// and `PartiallySorted` modes, the leading ordering separates + /// partitions, so finished partitions are pruned eagerly instead and no + /// such bound is needed. + most_recent_row: Option, } impl BoundedWindowAggStream { @@ -1066,15 +1066,22 @@ impl BoundedWindowAggStream { window_expr, baseline_metrics, search_mode, + most_recent_row: None, }) } fn compute_aggregates(&mut self) -> Result> { // calculate window cols + let eval_ctx = WindowEvalContext::default() + .with_most_recent_row(self.most_recent_row.as_ref()); for (cur_window_expr, state) in self.window_expr.iter().zip(&mut self.window_agg_states) { - cur_window_expr.evaluate_stateful(&self.partition_buffers, state)?; + cur_window_expr.evaluate_stateful( + &self.partition_buffers, + state, + &eval_ctx, + )?; } let schema = Arc::clone(&self.schema); @@ -1118,6 +1125,9 @@ impl BoundedWindowAggStream { // stopped when dropped. let _timer = elapsed_compute.timer(); + if self.search_mode.is_mode_linear() && batch.num_rows() > 0 { + self.most_recent_row = Some(get_last_row_batch(&batch)?); + } self.search_mode.update_partition_batch( &mut self.input_buffer, batch, From b337cec0f00b371604d6dd9329ff52f5b80c360f Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Fri, 31 Jul 2026 13:29:02 -0400 Subject: [PATCH 2/2] docs: add upgrade guide entries for the window watermark API changes --- .../library-user-guide/upgrading/55.0.0.md | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index d97082e881e0e..d9bc55326c92a 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -943,6 +943,87 @@ let plan = deserialize_bytes(&proto_bytes)?; See [PR #23827](https://github.com/apache/datafusion/pull/23827) for details. +### `WindowExpr::evaluate_stateful` now takes a `WindowEvalContext` + +`WindowExpr::evaluate_stateful` (and the provided +`AggregateWindowExpr::aggregate_evaluate_stateful` method) take a new +`WindowEvalContext` argument carrying stream-level information that is shared +by all partitions: + +```rust,ignore +// Before +fn evaluate_stateful( + &self, + partition_batches: &PartitionBatches, + window_agg_state: &mut PartitionWindowAggStates, +) -> Result<()> + +// After +fn evaluate_stateful( + &self, + partition_batches: &PartitionBatches, + window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, +) -> Result<()> +``` + +`WindowEvalContext` currently carries the most recent input row, which +previously lived in each partition's `PartitionBatchState` (see the next +section). The struct is `#[non_exhaustive]` so that fields can be added +without further signature changes: construct it with +`WindowEvalContext::default()` and set fields through its builder methods. + +**Who is affected:** + +- Implementations of the `WindowExpr` trait that override `evaluate_stateful` + must add the new parameter. +- Callers of `evaluate_stateful` or `aggregate_evaluate_stateful` must pass a + context. + +**Migration guide:** + +```rust,ignore +use datafusion_physical_expr::window::WindowEvalContext; + +// Before +window_expr.evaluate_stateful(&partition_batches, &mut window_agg_state)?; + +// After +let eval_ctx = WindowEvalContext::default() + .with_most_recent_row(most_recent_row.as_ref()); +window_expr.evaluate_stateful( + &partition_batches, + &mut window_agg_state, + &eval_ctx, +)?; +``` + +Pass `WindowEvalContext::default()` when no most-recent-row watermark is +available (for example, when the input is sorted by the partition keys and +partition ends are detected directly). + +### `PartitionBatchState::most_recent_row` removed + +The `most_recent_row` field and the `set_most_recent_row` method have been +removed from `datafusion_expr::window_state::PartitionBatchState`. The most +recent input row is a property of the whole input stream rather than +per-partition state: every partition observed the same value. It is now +tracked once by the operator driving the evaluation and passed to window +expressions through the new `WindowEvalContext` argument of +`WindowExpr::evaluate_stateful` described above. + +**Who is affected:** + +- Code that read `PartitionBatchState::most_recent_row` or called + `set_most_recent_row`, such as custom streaming window operators. + +**Migration guide:** + +Track the most recent input row once per stream (for example, a one-row +slice of the last non-empty input batch) and pass it to window expressions +via `WindowEvalContext::with_most_recent_row` instead of copying it into +each partition's state. + ### `MSRV` updated to 1.94.0 The Minimum Supported Rust Version (MSRV) has been updated to [`1.94.0`].