perf: track BoundedWindowAggExec Linear-mode watermark once per stream

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

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

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;

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>] {

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();

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,

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,