perf: skip eval of standard window functions for unchanged partitions

datafusion/expr/src/window_state.rs

@@ -38,16 +38,23 @@ pub struct WindowAggState {
     /// The range that we calculate the window function
     pub window_frame_range: Range<usize>,
     pub window_frame_ctx: Option<WindowFrameContext>,
-    /// The index of the last row that its result is calculated inside the partition record batch buffer.
+    /// The index of the first row in the partition's buffered record batch
+    /// whose result has not been calculated yet; equivalently, the number
+    /// of buffered rows with calculated results. The next evaluation pass
+    /// resumes from this index.
     pub last_calculated_index: usize,
     /// The offset of the deleted row number
     pub offset_pruned_rows: usize,
     /// Stores the results calculated by window frame
     pub out_col: ArrayRef,
-    /// Keeps track of how many rows should be generated to be in sync with input record_batch.
+    /// Keeps track of how many rows should be generated to be in sync with input record_batch
     // (For each row in the input record batch we need to generate a window result).
     pub n_row_result_missing: usize,
-    /// Flag indicating whether we have received all data for this partition
+    /// Snapshot of [`PartitionBatchState::n_rows_received`] as of this
+    /// partition's last evaluation pass
+    pub n_rows_received: usize,
+    /// Snapshot of [`PartitionBatchState::is_end`] as of this partition's
+    /// last evaluation pass
     pub is_end: bool,
 }
 
@@ -97,6 +104,7 @@ impl WindowAggState {
         }
         self.n_row_result_missing =
             partition_batch_state.record_batch.num_rows() - self.last_calculated_index;
+        self.n_rows_received = partition_batch_state.n_rows_received;
         self.is_end = partition_batch_state.is_end;
         Ok(())
     }
@@ -126,6 +134,24 @@ impl WindowAggState {
         all_rows_have_results && !partition_just_ended
     }
 
+    /// Returns true when the partition's input is unchanged since the last
+    /// evaluation pass: no new rows have arrived and the end flag has not
+    /// changed.
+    ///
+    /// Standard (non-aggregate) window functions read nothing outside their
+    /// partition, so we can skip re-evaluating them when their input hasn't
+    /// changed. Aggregate window functions cannot use this test to skip
+    /// evaluation, because they also consult the ORDER BY values of the most
+    /// recent input row across *all* partitions.
+    #[inline]
+    pub fn is_input_unchanged(
+        &self,
+        partition_batch_state: &PartitionBatchState,
+    ) -> bool {
+        self.n_rows_received == partition_batch_state.n_rows_received
+            && self.is_end == partition_batch_state.is_end
+    }
+
     pub fn new(out_type: &DataType) -> Result<Self> {
         let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?;
         Ok(Self {
@@ -135,6 +161,7 @@ impl WindowAggState {
             offset_pruned_rows: 0,
             out_col: empty_out_col,
             n_row_result_missing: 0,
+            n_rows_received: 0,
             is_end: false,
         })
     }
@@ -273,6 +300,8 @@ impl WindowFrameContext {
 pub struct PartitionBatchState {
     /// The record batch belonging to current partition
     pub record_batch: RecordBatch,
+    /// Total number of rows this partition has ever received
+    pub n_rows_received: usize,
     /// Flag indicating whether we have received all data for this partition
     pub is_end: bool,
     /// Number of rows emitted for this partition since the last pruning pass
@@ -283,14 +312,17 @@ impl PartitionBatchState {
     pub fn new(schema: SchemaRef) -> Self {
         Self {
             record_batch: RecordBatch::new_empty(schema),
+            n_rows_received: 0,
             is_end: false,
             n_out_row: 0,
         }
     }
 
     pub fn new_with_batch(batch: RecordBatch) -> Self {
+        let n_rows_received = batch.num_rows();
         Self {
             record_batch: batch,
+            n_rows_received,
             is_end: false,
             n_out_row: 0,
         }
@@ -299,6 +331,7 @@ impl PartitionBatchState {
     pub fn extend(&mut self, batch: &RecordBatch) -> Result<()> {
         self.record_batch =
             concat_batches(&self.record_batch.schema(), [&self.record_batch, batch])?;
+        self.n_rows_received += batch.num_rows();
         Ok(())
     }
 }
@@ -697,7 +730,8 @@ fn check_equality(current: &[ScalarValue], target: &[ScalarValue]) -> Result<boo
 mod tests {
     use super::*;
 
-    use arrow::array::Float64Array;
+    use arrow::array::{Float64Array, UInt64Array};
+    use arrow::datatypes::{Field, Schema};
 
     fn get_test_data() -> (Vec<ArrayRef>, Vec<SortOptions>) {
         let range_columns: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![
@@ -711,6 +745,62 @@ mod tests {
         (range_columns, sort_options)
     }
 
+    #[test]
+    fn test_is_input_unchanged() -> Result<()> {
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("ts", DataType::UInt64, false)]));
+        let batch = |values: &[u64]| -> Result<RecordBatch> {
+            Ok(RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(UInt64Array::from(values.to_vec()))],
+            )?)
+        };
+        let results = |n: usize| ScalarValue::UInt64(Some(0)).to_array_of_size(n);
+
+        let mut partition = PartitionBatchState::new_with_batch(batch(&[1, 2, 3])?);
+        assert_eq!(partition.n_rows_received, 3);
+
+        // A freshly created state has seen none of the partition's rows, so
+        // it must not report the input as unchanged.
+        let mut state = WindowAggState::new(&DataType::UInt64)?;
+        assert!(!state.is_input_unchanged(&partition));
+
+        // After an evaluation pass (here producing results for two of the
+        // three rows), the state has seen the partition's current input.
+        state.update(&results(2)?, &partition)?;
+        assert!(state.is_input_unchanged(&partition));
+
+        // New rows must be detected...
+        partition.extend(&batch(&[4])?)?;
+        assert_eq!(partition.n_rows_received, 4);
+        assert!(!state.is_input_unchanged(&partition));
+        // ...until the next evaluation pass catches up again.
+        state.update(&results(2)?, &partition)?;
+        assert!(state.is_input_unchanged(&partition));
+
+        // Pruning rows whose results are already calculated shrinks the
+        // buffer but not the received-row count, so it must not make the
+        // input appear changed.
+        let n_prune = 2;
+        let batch_ref = &partition.record_batch;
+        partition.record_batch = batch_ref.slice(n_prune, batch_ref.num_rows() - n_prune);
+        // An evaluator would have advanced the frame at least past the
+        // pruned rows; prune_state requires this.
+        state.window_frame_range = n_prune..n_prune;
+        state.prune_state(n_prune);
+        assert_eq!(partition.n_rows_received, 4);
+        assert!(state.is_input_unchanged(&partition));
+
+        // The end-of-partition transition must trigger one final pass, even
+        // though no rows arrived.
+        partition.is_end = true;
+        assert!(!state.is_input_unchanged(&partition));
+        state.update(&results(0)?, &partition)?;
+        assert!(state.is_input_unchanged(&partition));
+
+        Ok(())
+    }
+
     fn assert_group_ranges(
         window_frame: &Arc<WindowFrame>,
         expected_results: Vec<(Range<usize>, usize)>,

datafusion/physical-expr/src/window/standard.rs

@@ -179,6 +179,11 @@ impl WindowExpr for StandardWindowExpr {
                             published: false,
                         })
                 };
+            // Skip partitions whose input is unchanged since the last
+            // evaluation pass.
+            if window_state.state.is_input_unchanged(partition_batch_state) {
+                continue;
+            }
             let WindowFn::Builtin(evaluator) = &mut window_state.window_fn else {
                 unreachable!()
             };

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -1642,8 +1642,10 @@ mod tests {
     };
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_functions_aggregate::sum::sum_udaf;
+    use datafusion_functions_window::lead_lag::lead_udwf;
     use datafusion_functions_window::nth_value::last_value_udwf;
     use datafusion_functions_window::nth_value::nth_value_udwf;
+    use datafusion_functions_window::row_number::row_number_udwf;
     use datafusion_physical_expr::expressions::{Column, Literal, col};
     use datafusion_physical_expr::window::{PartitionKey, StandardWindowExpr};
     use datafusion_physical_expr::{LexOrdering, PhysicalExpr};
@@ -2183,6 +2185,115 @@ mod tests {
         Ok(())
     }
 
+    // In `Linear` mode, a partition may receive no new rows for several
+    // input batches while other partitions keep growing. The evaluation
+    // sweep skips a partition whose input is unchanged, so this test
+    // drives a partition through quiet batches and then resumes it: the
+    // results after the gap must continue from the retained evaluator
+    // state. ROW_NUMBER is causal, so the quiet partition is fully
+    // calculated while it waits; LEAD is not, so its result for the
+    // partition's last buffered row stays pending across the quiet
+    // batches and must materialize once the partition receives another
+    // row.
+    #[tokio::test]
+    async fn bounded_window_linear_quiet_partition_resume_standard() -> 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<RecordBatch> {
+            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(),
+        }];
+        // Both functions use the default frame of a window with an ORDER BY
+        // clause (RANGE UNBOUNDED PRECEDING..CURRENT ROW).
+        let row_number_expr = create_window_expr(
+            &WindowFunctionDefinition::WindowUDF(row_number_udwf()),
+            "row_number".to_string(),
+            &[],
+            &partition_by,
+            &order_by,
+            Arc::new(WindowFrame::new(Some(false))),
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+        let lead_expr = create_window_expr(
+            &WindowFunctionDefinition::WindowUDF(lead_udwf()),
+            "lead".to_string(),
+            &[col("ts", &schema)?],
+            &partition_by,
+            &order_by,
+            Arc::new(WindowFrame::new(Some(false))),
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+        let physical_plan = BoundedWindowAggExec::try_new(
+            vec![row_number_expr, lead_expr],
+            memory_exec,
+            InputOrderMode::Linear,
+            true,
+        )
+        .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)?;
+
+        let batches = collect(physical_plan.execute(0, task_context())?).await?;
+
+        // The skip must not delay results that are ready to be finalized;
+        // they stream out as soon as every window expression has produced
+        // them. LEAD holds back only the last buffered row of a partition,
+        // so one row unblocks right after the first input batch, five more
+        // when partition 0 resumes in the last input batch, and the final
+        // two (whose LEAD results need the end of the input) in the flush
+        // after the input is exhausted.
+        assert_eq!(
+            batches.iter().map(|b| b.num_rows()).collect::<Vec<_>>(),
+            vec![1, 5, 2],
+            "expected results to stream as they become final"
+        );
+
+        assert_snapshot!(batches_to_string(&batches), @r"
+        +----+----+------------+------+
+        | pk | ts | row_number | lead |
+        +----+----+------------+------+
+        | 0  | 0  | 1          | 1    |
+        | 0  | 1  | 2          | 6    |
+        | 1  | 2  | 1          | 3    |
+        | 1  | 3  | 2          | 4    |
+        | 1  | 4  | 3          | 5    |
+        | 1  | 5  | 4          | 7    |
+        | 0  | 6  | 3          |      |
+        | 1  | 7  | 5          |      |
+        +----+----+------------+------+
+        ");
+        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