feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec

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

@@ -609,6 +609,17 @@ pub struct WindowState {
     pub state: WindowAggState,
     pub window_fn: WindowFn,
 }
+
+impl WindowState {
+    /// `Accumulator::state()` if this is an aggregate window; `None` otherwise.
+    pub fn aggregate_state(&mut self) -> Result<Option<Vec<ScalarValue>>> {
+        match &mut self.window_fn {
+            WindowFn::Aggregate(acc) => acc.state().map(Some),
+            WindowFn::Builtin(_) => Ok(None),
+        }
+    }
+}
+
 pub type PartitionWindowAggStates = IndexMap<PartitionKey, WindowState>;
 
 /// The IndexMap (i.e. an ordered HashMap) where record batches are separated for each partition.

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

@@ -23,7 +23,7 @@
 use std::cmp::{Ordering, min};
 use std::collections::VecDeque;
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 use std::task::{Context, Poll};
 
 use super::utils::create_schema;
@@ -54,7 +54,7 @@ use datafusion_common::utils::{
     evaluate_partition_ranges, get_at_indices, get_row_at_idx,
 };
 use datafusion_common::{
-    HashMap, Result, arrow_datafusion_err, exec_datafusion_err, exec_err,
+    HashMap, Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err,
 };
 use datafusion_execution::TaskContext;
 use datafusion_expr::ColumnarValue;
@@ -75,6 +75,11 @@ use hashbrown::hash_table::HashTable;
 use indexmap::IndexMap;
 use log::debug;
 
+/// Per-output-partition slots for finalized `Accumulator::state()`, indexed by
+/// window expression. Outer `Vec` is one entry per window expression; inner
+/// `Option<Vec<ScalarValue>>` is `Some` for aggregates, `None` otherwise.
+type FinalizedPartitionStates = Arc<[OnceLock<Vec<Option<Vec<ScalarValue>>>>]>;
+
 /// Window execution plan
 #[derive(Debug, Clone)]
 pub struct BoundedWindowAggExec {
@@ -99,6 +104,9 @@ pub struct BoundedWindowAggExec {
     cache: Arc<PlanProperties>,
     /// If `can_rerepartition` is false, partition_keys is always empty.
     can_repartition: bool,
+    /// Populated by the stream at partition-close; read via
+    /// [`Self::finalized_partition_state`].
+    finalized_state: FinalizedPartitionStates,
 }
 
 impl BoundedWindowAggExec {
@@ -130,6 +138,10 @@ impl BoundedWindowAggExec {
             }
         };
         let cache = Self::compute_properties(&input, &schema, &window_expr)?;
+        let finalized_state: FinalizedPartitionStates =
+            (0..cache.partitioning.partition_count())
+                .map(|_| OnceLock::new())
+                .collect();
         Ok(Self {
             input,
             window_expr,
@@ -139,6 +151,7 @@ impl BoundedWindowAggExec {
             ordered_partition_by_indices,
             cache: Arc::new(cache),
             can_repartition,
+            finalized_state,
         })
     }
 
@@ -152,6 +165,24 @@ impl BoundedWindowAggExec {
         &self.input
     }
 
+    /// Finalized `Accumulator::state()` per aggregate window expression on
+    /// output partition `partition`. `None` for the outer option means the
+    /// partition is out of range or the stream hasn't drained; inner `None`
+    /// slots are non-aggregate window functions.
+    ///
+    /// Intended to support prefix scans over ordered window aggregates.
+    /// Assumes at most one PARTITION BY group closes per output partition;
+    /// additional groups are silently dropped.
+    pub fn finalized_partition_state(
+        &self,
+        partition: usize,
+    ) -> Option<&[Option<Vec<ScalarValue>>]> {
+        self.finalized_state
+            .get(partition)?
+            .get()
+            .map(Vec::as_slice)
+    }
+
     /// Return the output sort order of partition keys: For example
     /// OVER(PARTITION BY a, ORDER BY b) -> would give sorting of the column a
     // We are sure that partition by columns are always at the beginning of sort_keys
@@ -377,6 +408,8 @@ impl ExecutionPlan for BoundedWindowAggExec {
             input,
             BaselineMetrics::new(&self.metrics, partition),
             search_mode,
+            Arc::clone(&self.finalized_state),
+            partition,
         )?);
         Ok(stream)
     }
@@ -1010,6 +1043,11 @@ pub struct BoundedWindowAggStream {
     /// Search mode for partition columns. This determines the algorithm with
     /// which we group each partition.
     search_mode: Box<dyn PartitionSearcher>,
+    /// Shared with [`BoundedWindowAggExec::finalized_state`]; write once,
+    /// first-close-wins.
+    finalized_state: FinalizedPartitionStates,
+    /// Which slot in [`Self::finalized_state`] this stream writes to.
+    partition: usize,
 }
 
 impl BoundedWindowAggStream {
@@ -1021,6 +1059,9 @@ impl BoundedWindowAggStream {
     // For instance, if `n_out` number of rows are calculated, we can remove
     // first `n_out` rows from `self.input_buffer`.
     fn prune_state(&mut self, n_out: usize) -> Result<()> {
+        // Snapshot `Accumulator::state()` for any partition/window-expr pair
+        // whose entry is about to be dropped, before the retains below fire.
+        self.publish_finalized_state()?;
         // Prune `self.window_agg_states`:
         self.prune_out_columns();
         // Prune `self.partition_batches`:
@@ -1031,6 +1072,37 @@ impl BoundedWindowAggStream {
         self.search_mode.prune(n_out);
         Ok(())
     }
+
+    /// Bundle each window expression's `Accumulator::state()` for the closing
+    /// group into one `Vec<Option<Vec<ScalarValue>>>` and try to publish it to
+    /// our slot. `OnceLock::set` is first-write-wins, so subsequent groups on
+    /// the same output partition are silently dropped.
+    fn publish_finalized_state(&mut self) -> Result<()> {
+        let Some(slot) = self.finalized_state.get(self.partition) else {
+            return Ok(());
+        };
+        if slot.get().is_some() {
+            return Ok(());
+        }
+        let n_exprs = self.window_expr.len();
+        let mut per_expr: Vec<Option<Vec<ScalarValue>>> = vec![None; n_exprs];
+        let mut any_end = false;
+        for (expr_idx, window_agg_state) in self.window_agg_states.iter_mut().enumerate()
+        {
+            for (_key, ws) in window_agg_state.iter_mut() {
+                if !ws.state.is_end {
+                    continue;
+                }
+                per_expr[expr_idx] = ws.aggregate_state()?;
+                any_end = true;
+                break;
+            }
+        }
+        if any_end {
+            let _ = slot.set(per_expr);
+        }
+        Ok(())
+    }
 }
 
 impl Stream for BoundedWindowAggStream {
@@ -1053,6 +1125,8 @@ impl BoundedWindowAggStream {
         input: SendableRecordBatchStream,
         baseline_metrics: BaselineMetrics,
         search_mode: Box<dyn PartitionSearcher>,
+        finalized_state: FinalizedPartitionStates,
+        partition: usize,
     ) -> Result<Self> {
         let state = window_expr.iter().map(|_| IndexMap::new()).collect();
         let empty_batch = RecordBatch::new_empty(Arc::clone(&schema));
@@ -1066,6 +1140,8 @@ impl BoundedWindowAggStream {
             window_expr,
             baseline_metrics,
             search_mode,
+            finalized_state,
+            partition,
         })
     }
 
@@ -1924,4 +2000,81 @@ mod tests {
         ));
         Ok(())
     }
+
+    async fn run_sum_over_pk(rows: &[(i64, i64)]) -> Result<Arc<BoundedWindowAggExec>> {
+        use arrow::array::Int64Array;
+        use datafusion_functions_aggregate::sum::sum_udaf;
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("pk", DataType::Int64, false),
+            Field::new("v", DataType::Int64, false),
+        ]));
+        let pks: Vec<i64> = rows.iter().map(|(pk, _)| *pk).collect();
+        let vs: Vec<i64> = rows.iter().map(|(_, v)| *v).collect();
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int64Array::from(pks)),
+                Arc::new(Int64Array::from(vs)),
+            ],
+        )?;
+
+        let sort_info: LexOrdering = [PhysicalSortExpr {
+            expr: col("pk", &schema)?,
+            options: SortOptions::default(),
+        }]
+        .into();
+        let input: Arc<dyn ExecutionPlan> = Arc::new(
+            TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?
+                .try_with_sort_information(vec![sort_info])?,
+        );
+
+        let window_fn = WindowFunctionDefinition::AggregateUDF(sum_udaf());
+        let args = vec![col("v", &schema)?];
+        let partitionby = vec![col("pk", &schema)?];
+        let orderby: Vec<PhysicalSortExpr> = vec![];
+        let window_frame = Arc::new(WindowFrame::new(None));
+        let window_expr = create_window_expr(
+            &window_fn,
+            "sum(v)".to_string(),
+            &args,
+            &partitionby,
+            &orderby,
+            window_frame,
+            Arc::clone(&schema),
+            false,
+            false,
+            None,
+        )?;
+
+        let exec: Arc<BoundedWindowAggExec> = Arc::new(BoundedWindowAggExec::try_new(
+            vec![window_expr],
+            input,
+            InputOrderMode::Sorted,
+            false,
+        )?);
+
+        let ctx = Arc::new(TaskContext::default());
+        let stream = ExecutionPlan::execute(exec.as_ref(), 0, ctx)?;
+        let _batches = collect(stream).await?;
+        Ok(exec)
+    }
+
+    #[tokio::test]
+    async fn finalized_partition_state_returns_single_group_state() -> Result<()> {
+        let exec = run_sum_over_pk(&[(1, 10), (1, 20), (1, 30)]).await?;
+
+        let state = exec
+            .finalized_partition_state(0)
+            .expect("stream drained; slot should be populated");
+        assert_eq!(state.len(), 1, "one entry per window expression");
+        let inner = state[0]
+            .as_ref()
+            .expect("SUM is an aggregate — state should be Some");
+        assert_eq!(inner.len(), 1, "SumAccumulator::state() is a single scalar");
+        assert_eq!(inner[0], ScalarValue::Int64(Some(60)));
+
+        assert!(exec.finalized_partition_state(99).is_none());
+        Ok(())
+    }
 }