This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 3ef9a8c221 perf: track `BoundedWindowAggExec` Linear-mode watermark
once per stream (#24033)
3ef9a8c221 is described below
commit 3ef9a8c2216bf0761580946748c241bc2dc9b816
Author: Neil Conway <[email protected]>
AuthorDate: Tue Aug 4 01:47:58 2026 -0400
perf: track `BoundedWindowAggExec` Linear-mode watermark once per stream
(#24033)
## Which issue does this PR close?
- Related to #23982
## Rationale for this change
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<RecordBatch>` 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.
Benchmarks: (using #24032)
- 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)
## What changes are included in this PR?
* Add `WindowEvalContext` parameter to `aggregate_evaluate_stateful`
* Remove `PartitionBatchState::most_recent_row` and its setter; instead
pass row via `WindowEvalContext`
## Are these changes tested?
Yes, covered by existing tests.
## Are there any user-facing changes?
Yes, API change to `datafusion-expr`.
---
datafusion/expr/src/window_state.rs | 13 ----
datafusion/physical-expr/src/window/aggregate.rs | 7 +-
datafusion/physical-expr/src/window/mod.rs | 1 +
.../physical-expr/src/window/sliding_aggregate.rs | 7 +-
datafusion/physical-expr/src/window/standard.rs | 3 +-
datafusion/physical-expr/src/window/window_expr.rs | 55 ++++++++++++---
.../src/windows/bounded_window_agg_exec.rs | 52 ++++++++------
docs/source/library-user-guide/upgrading/55.0.0.md | 81 ++++++++++++++++++++++
8 files changed, 171 insertions(+), 48 deletions(-)
diff --git a/datafusion/expr/src/window_state.rs
b/datafusion/expr/src/window_state.rs
index f8d4609d36..ece07e5b09 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<RecordBatch>,
/// 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 1ff13d107c..7cfdcb167f 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 b45e35440a..79b9a9580a 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 a71df3ec88..29e569363a 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<dyn PhysicalExpr>] {
diff --git a/datafusion/physical-expr/src/window/standard.rs
b/datafusion/physical-expr/src/window/standard.rs
index 6f61174ee0..2de080ec9a 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 8db5651346..47147b909d 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<dyn Accumulator>,
record_batch: &RecordBatch,
- most_recent_row: Option<&RecordBatch>,
+ most_recent_row_order_bys: Option<&[ArrayRef]>,
last_range: &mut Range<usize>,
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<ScalarValue>;
+/// 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 680d5b657e..0ed751f506 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<dyn PartitionSearcher>,
+ /// 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<RecordBatch>,
}
impl BoundedWindowAggStream {
@@ -1066,15 +1066,22 @@ impl BoundedWindowAggStream {
window_expr,
baseline_metrics,
search_mode,
+ most_recent_row: None,
})
}
fn compute_aggregates(&mut self) -> Result<Option<RecordBatch>> {
// 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 803362d9d5..9a65651fc3 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`].
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
