perf: skip evaluating fully calculated window partitions

datafusion/expr/src/window_state.rs

@@ -101,6 +101,31 @@ impl WindowAggState {
         Ok(())
     }
 
+    /// Returns true when this state is fully up to date with the partition's
+    /// buffered batch, meaning another evaluation pass over the partition could
+    /// not produce any new results or change any state:
+    ///
+    /// - `last_calculated_index` has reached the end of the partition's
+    ///   buffered batch, so every row of this partition that has arrived so
+    ///   far already has a result.
+    /// - When a partition ends, a final evaluation pass is needed to bring
+    ///   derived state up to date.
+    #[inline]
+    pub fn is_up_to_date_with(
+        &self,
+        partition_batch_state: &PartitionBatchState,
+    ) -> bool {
+        let all_rows_have_results =
+            self.last_calculated_index == partition_batch_state.record_batch.num_rows();
+        if all_rows_have_results {
+            debug_assert_eq!(self.n_row_result_missing, 0);
+        }
+
+        // `self.is_end` holds the flag as of the previous evaluation pass.
+        let partition_just_ended = !self.is_end && partition_batch_state.is_end;
+        all_rows_have_results && !partition_just_ended
+    }
+
     pub fn new(out_type: &DataType) -> Result<Self> {
         let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?;
         Ok(Self {

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

@@ -263,6 +263,12 @@ pub trait AggregateWindowExpr: WindowExpr {
             let state = &mut window_state.state;
             let record_batch = &partition_batch_state.record_batch;
 
+            // Skip partitions that cannot produce anything new until they
+            // either receive rows or reach their end.
+            if state.is_up_to_date_with(partition_batch_state) {
+                continue;
+            }
+
             // If there is no window state context, initialize it.
             let window_frame_ctx = state.window_frame_ctx.get_or_insert_with(|| {
                 let sort_options = self.order_by().iter().map(|o| o.options).collect();

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

@@ -1194,6 +1194,18 @@ impl BoundedWindowAggStream {
     /// Prunes the sections of the record batch (for each partition)
     /// that we no longer need to calculate the window function result.
     fn prune_partition_batches(&mut self) {
+        // Check that per-state and per-partition end-flags are consistent;
+        // otherwise, the pruning code below might produce inconsistent state.
+        #[cfg(debug_assertions)]
+        for window_agg_state in self.window_agg_states.iter() {
+            for (partition_row, WindowState { state, .. }) in window_agg_state.iter() {
+                debug_assert_eq!(
+                    state.is_end, self.partition_buffers[partition_row].is_end,
+                    "window state's recorded end flag is out of sync with its partition"
+                );
+            }
+        }
+
         // Remove partitions which we know already ended (is_end flag is true).
         // Since the retain method preserves insertion order, we still have
         // ordering in between partitions after removal.
@@ -1395,6 +1407,7 @@ mod tests {
         WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
     };
     use datafusion_functions_aggregate::count::count_udaf;
+    use datafusion_functions_aggregate::sum::sum_udaf;
     use datafusion_functions_window::nth_value::last_value_udwf;
     use datafusion_functions_window::nth_value::nth_value_udwf;
     use datafusion_physical_expr::expressions::{Column, Literal, col};
@@ -1834,6 +1847,108 @@ mod tests {
         Ok(())
     }
 
+    // In `Linear` mode, a partition may receive no new rows for several
+    // input batches while other partitions keep growing. Once all of a
+    // partition's buffered rows have results, the evaluation sweep skips
+    // it until it receives rows again, so this test drives a partition
+    // through quiet batches and then resumes it: the results after the
+    // gap must continue from the retained accumulator state. Both frames
+    // are causal, so results finalize in the batch their row arrives in
+    // and the quiet partition is fully calculated while it waits.
+    #[tokio::test]
+    async fn bounded_window_linear_quiet_partition_resume() -> 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(),
+        }];
+        // A running COUNT (plain aggregate) and a SUM over the previous and
+        // current row (sliding aggregate).
+        let count_expr = create_window_expr(
+            &WindowFunctionDefinition::AggregateUDF(count_udaf()),
+            "count".to_string(),
+            &[col("ts", &schema)?],
+            &partition_by,
+            &order_by,
+            Arc::new(WindowFrame::new_bounds(
+                WindowFrameUnits::Rows,
+                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+                WindowFrameBound::CurrentRow,
+            )),
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+        let sum_expr = create_window_expr(
+            &WindowFunctionDefinition::AggregateUDF(sum_udaf()),
+            "sum".to_string(),
+            &[col("ts", &schema)?],
+            &partition_by,
+            &order_by,
+            Arc::new(WindowFrame::new_bounds(
+                WindowFrameUnits::Rows,
+                WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))),
+                WindowFrameBound::CurrentRow,
+            )),
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+        let physical_plan = BoundedWindowAggExec::try_new(
+            vec![count_expr, sum_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?;
+
+        assert_snapshot!(batches_to_string(&batches), @r"
+        +----+----+-------+-----+
+        | pk | ts | count | sum |
+        +----+----+-------+-----+
+        | 0  | 0  | 1     | 0   |
+        | 0  | 1  | 2     | 1   |
+        | 1  | 2  | 1     | 2   |
+        | 1  | 3  | 2     | 5   |
+        | 1  | 4  | 3     | 7   |
+        | 1  | 5  | 4     | 9   |
+        | 0  | 6  | 3     | 7   |
+        | 1  | 7  | 5     | 12  |
+        +----+----+-------+-----+
+        ");
+        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