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, 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`].