Expose accumulator state to allow prefix scanning

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

@@ -176,6 +176,7 @@ impl WindowExpr for StandardWindowExpr {
                         .or_insert(WindowState {
                             state: new_state.clone(),
                             window_fn: WindowFn::Builtin(evaluator),
+                            published: false,
                         })
                 };
             let evaluator = match &mut window_state.window_fn {

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

@@ -33,7 +33,8 @@ use datafusion_common::cast::as_boolean_array;
 use datafusion_common::hash_utils::RandomState;
 use datafusion_common::utils::compare_rows;
 use datafusion_common::{
-    Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, internal_err,
+    Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err,
+    internal_err,
 };
 use datafusion_expr::window_state::{
     PartitionBatchState, WindowAggState, WindowFrameContext, WindowFrameStateGroups,
@@ -250,6 +251,7 @@ pub trait AggregateWindowExpr: WindowExpr {
                     WindowState {
                         state: WindowAggState::new(out_type)?,
                         window_fn: WindowFn::Aggregate(accumulator),
+                        published: false,
                     },
                 );
             };
@@ -646,7 +648,49 @@ impl<'a> WindowEvalContext<'a> {
 pub struct WindowState {
     pub state: WindowAggState,
     pub window_fn: WindowFn,
+    /// True once [`Self::aggregate_state`] has been called on this entry.
+    /// Guards against a second destructive [`Accumulator::state`] read: the
+    /// method itself errors on second call, and the observer loop in
+    /// `BoundedWindowAggStream::publish_finalized_states` uses this as an
+    /// early-skip so it doesn't attempt one. Independent of `state.is_end`,
+    /// which is a group-closed signal that the pruning path also reads.
+    pub published: bool,
 }
+
+impl WindowState {
+    /// [`Accumulator::state`] if this window function is an aggregate, `None`
+    /// otherwise (built-in functions like `row_number`, `rank`, `lead`/`lag`
+    /// have no serializable accumulator state).
+    ///
+    /// [`Accumulator::state`] takes `&mut self` and its trait doc calls out
+    /// that "this function should not be called twice, otherwise it will
+    /// result in potentially non-deterministic behavior." Several built-in
+    /// impls (`median`, `percentile_cont`, `string_agg`,
+    /// `min_max_bytes`/`min_max_struct`) `std::mem::take` their internal
+    /// buffers on call — a second call returns *empty* state, not the same
+    /// state, so a downstream prefix-merge would silently lose every value
+    /// the accumulator had ingested.
+    ///
+    /// Enforced at this layer: on first call we set [`Self::published`] and
+    /// return the state; any later call errors rather than performing a
+    /// destructive re-read.
+    pub fn aggregate_state(&mut self) -> Result<Option<Vec<ScalarValue>>> {
+        if self.published {
+            return exec_err!(
+                "WindowState::aggregate_state called more than once; \
+                 Accumulator::state is a destructive read for several \
+                 built-in aggregates and a second call would silently lose data"
+            );
+        }
+        let state = match &mut self.window_fn {
+            WindowFn::Aggregate(accumulator) => Some(accumulator.state()?),
+            WindowFn::Builtin(_) => None,
+        };
+        self.published = true;
+        Ok(state)
+    }
+}
+
 pub type PartitionWindowAggStates = IndexMap<PartitionKey, WindowState, RandomState>;
 
 /// The IndexMap (i.e. an ordered HashMap) where record batches are separated for each partition.
@@ -656,11 +700,66 @@ pub type PartitionBatches = IndexMap<PartitionKey, PartitionBatchState, RandomSt
 mod tests {
     use std::sync::Arc;
 
-    use crate::window::window_expr::is_row_ahead;
+    use crate::window::window_expr::{WindowFn, WindowState, is_row_ahead};
 
     use arrow::array::{ArrayRef, Float64Array};
     use arrow::compute::SortOptions;
-    use datafusion_common::Result;
+    use arrow::datatypes::DataType;
+    use datafusion_common::{Result, ScalarValue};
+    use datafusion_expr::{Accumulator, window_state::WindowAggState};
+
+    /// Minimal [`Accumulator`] whose `state()` records how many times it was
+    /// called by returning the count as its single state element. Any second
+    /// call would surface (were it allowed to happen) as `[UInt64(2)]`
+    /// instead of `[UInt64(1)]`.
+    #[derive(Debug)]
+    struct CallCountingAccumulator {
+        calls: usize,
+    }
+
+    impl Accumulator for CallCountingAccumulator {
+        fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {
+            Ok(())
+        }
+        fn evaluate(&mut self) -> Result<ScalarValue> {
+            Ok(ScalarValue::Null)
+        }
+        fn size(&self) -> usize {
+            size_of::<Self>()
+        }
+        fn state(&mut self) -> Result<Vec<ScalarValue>> {
+            self.calls += 1;
+            Ok(vec![ScalarValue::UInt64(Some(self.calls as u64))])
+        }
+        fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> {
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn aggregate_state_errors_on_second_call() -> Result<()> {
+        // `Accumulator::state()` is a destructive read for several built-in
+        // aggregates (median, percentile_cont, string_agg, min_max_bytes/
+        // min_max_struct all `mem::take` their internal buffers). Its trait
+        // doc says "should not be called twice"; `WindowState::aggregate_state`
+        // enforces that at this layer by returning an error rather than
+        // performing the second read.
+        let acc: Box<dyn Accumulator> = Box::new(CallCountingAccumulator { calls: 0 });
+        let mut ws = WindowState {
+            state: WindowAggState::new(&DataType::UInt64)?,
+            window_fn: WindowFn::Aggregate(acc),
+            published: false,
+        };
+        let first = ws.aggregate_state()?;
+        assert_eq!(first, Some(vec![ScalarValue::UInt64(Some(1))]));
+        assert!(ws.published, "published must flip on successful publish");
+        let err = ws.aggregate_state().unwrap_err().to_string();
+        assert!(
+            err.contains("called more than once"),
+            "expected second-call error, got: {err}"
+        );
+        Ok(())
+    }
 
     #[test]
     fn test_is_row_ahead() -> Result<()> {

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -1178,6 +1178,7 @@ pub fn ensure_distribution(
             exec.window_expr(),
             exec.input(),
             &exec.partition_keys(),
+            None,
         )? {
             plan = updated_window;
         }
@@ -1186,6 +1187,7 @@ pub fn ensure_distribution(
             exec.window_expr(),
             exec.input(),
             &exec.partition_keys(),
+            exec.state_observer().cloned(),
         )?
     {
         plan = updated_window;

datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs

@@ -556,17 +556,33 @@ fn adjust_window_sort_removal(
     window_tree.children.push(child_node);
 
     let child_plan = &window_tree.children[0].plan;
+    // Captured up-front so the fallback `BoundedWindowAggExec::try_new` below
+    // can reinstall the observer that was on the source exec. `None` when
+    // the source is a `WindowAggExec` (no observer) or when no observer was
+    // installed on the source `BoundedWindowAggExec`.
+    let state_observer = window_tree
+        .plan
+        .downcast_ref::<BoundedWindowAggExec>()
+        .and_then(|exec| exec.state_observer().cloned());
     let (window_expr, new_window) = if let Some(exec) =
         window_tree.plan.downcast_ref::<WindowAggExec>()
     {
         let window_expr = exec.window_expr();
-        let new_window =
-            get_best_fitting_window(window_expr, child_plan, &exec.partition_keys())?;
+        let new_window = get_best_fitting_window(
+            window_expr,
+            child_plan,
+            &exec.partition_keys(),
+            None,
+        )?;
         (window_expr, new_window)
     } else if let Some(exec) = window_tree.plan.downcast_ref::<BoundedWindowAggExec>() {
         let window_expr = exec.window_expr();
-        let new_window =
-            get_best_fitting_window(window_expr, child_plan, &exec.partition_keys())?;
+        let new_window = get_best_fitting_window(
+            window_expr,
+            child_plan,
+            &exec.partition_keys(),
+            state_observer.clone(),
+        )?;
         (window_expr, new_window)
     } else {
         return plan_err!("Expected WindowAggExec or BoundedWindowAggExec");
@@ -589,12 +605,15 @@ fn adjust_window_sort_removal(
         window_tree.children.push(child_node);
 
         if window_expr.iter().all(|e| e.uses_bounded_memory()) {
-            Arc::new(BoundedWindowAggExec::try_new(
-                window_expr.to_vec(),
-                child_plan,
-                InputOrderMode::Sorted,
-                !window_expr[0].partition_by().is_empty(),
-            )?) as _
+            Arc::new(
+                BoundedWindowAggExec::try_new(
+                    window_expr.to_vec(),
+                    child_plan,
+                    InputOrderMode::Sorted,
+                    !window_expr[0].partition_by().is_empty(),
+                )?
+                .with_state_observer(state_observer)?,
+            ) as _
         } else {
             Arc::new(WindowAggExec::try_new(
                 window_expr.to_vec(),