refactor(core): FinalizedPartitionState as a newtype

Vec<Option<Vec<ScalarValue>>> appeared in this operator's public signatures,
where it is neither readable nor searchable, and it spelled "no state for this
window expression" two ways: a missing index, and a None at a present index.
Every caller handled both. slot(window_expr_index) collapses them into one
answer, and a later change to the representation now stays internal.

Also removes two comments claiming DataFusion guarantees at most one PARTITION
BY group per partition. It does not — apache/datafusion#24035 shipped a
callback keyed by group, so that invariant is ours, and the scheduler enforces
it by rejecting any report carrying a key. A window that does have a PARTITION
BY needs nothing from this operator anyway: BoundedWindowAggExec asks for
KeyPartitioned input, so each partition's window is already independent.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: avantgardnerio

ballista/core/src/execution_plans/prefix_merge.rs

@@ -68,11 +68,15 @@
 //! the *pre-merged* [`Accumulator::state`] for each aggregate window
 //! expression (indexed by position in `BoundedWindowAggExec::window_expr()`).
 //! Only consumed by [`WindowApply::Aggregate`] entries;
-//! [`WindowApply::Scalar`] carries its own offsets inline. The DF-side getter
-//! already commits to at-most-one PARTITION BY group per DF partition (see
-//! `apache/datafusion#24007`), so no per-key dimension is exposed here. The
-//! scheduler bakes the state when it constructs the downstream stage after
-//! the upstream stage's tasks complete.
+//! [`WindowApply::Scalar`] carries its own offsets inline. No PARTITION BY
+//! dimension is exposed: DataFusion publishes state per closed group, but the
+//! rewrite that plants this operator only fires on windows without a
+//! PARTITION BY, and the scheduler rejects any report carrying a key rather
+//! than flattening two groups together. A window that does have a PARTITION
+//! BY needs no help from this operator, since `BoundedWindowAggExec` asks for
+//! `KeyPartitioned` input and each partition's window is already independent.
+//! The scheduler bakes the state when it constructs the downstream stage
+//! after the upstream stage's tasks complete.
 //!
 //! **Status.** Both apply paths are implemented.
 //!
@@ -148,12 +152,48 @@ use crate::execution_plans::plan_algebra::{
 /// value per window expression — see the `prefix_merge` module docs for the
 /// division of labor. This operator applies it; it does not compute it.
 ///
-/// The type mirrors `datafusion::physical_plan::windows::FinalizedPartitionState`
-/// from [apache/datafusion#24007]; the alias here is a local stand-in so
-/// this crate compiles against stable DataFusion 54 until that PR lands.
-///
-/// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007
-pub type FinalizedPartitionState = Vec<Option<Vec<ScalarValue>>>;
+/// A newtype rather than an alias for `Vec<Option<Vec<ScalarValue>>>`. That
+/// type appears in this operator's public signatures, where it is neither
+/// readable nor searchable, and it spells "no state for this window
+/// expression" two ways — a missing index and a `None` — which every caller
+/// then has to handle. [`Self::slot`] collapses both into one answer. Any
+/// later change to the representation also stays internal rather than
+/// breaking whoever wrote the concrete type.
+#[derive(Debug, Clone, Default, PartialEq)]
+pub struct FinalizedPartitionState {
+    /// Indexed by position in the upstream operator's `window_expr()` list.
+    per_window_expr: Vec<Option<Vec<ScalarValue>>>,
+}
+
+impl FinalizedPartitionState {
+    /// Build from one slot per window expression, in `window_expr()` order.
+    pub fn new(per_window_expr: Vec<Option<Vec<ScalarValue>>>) -> Self {
+        Self { per_window_expr }
+    }
+
+    /// Merged state for `window_expr_index`, or `None` when that expression
+    /// published none or the index is past the end.
+    pub fn slot(&self, window_expr_index: usize) -> Option<&Vec<ScalarValue>> {
+        self.per_window_expr
+            .get(window_expr_index)
+            .and_then(|slot| slot.as_ref())
+    }
+
+    /// Every slot in window-expression order, for wire encoding.
+    pub fn slots(&self) -> &[Option<Vec<ScalarValue>>] {
+        &self.per_window_expr
+    }
+
+    /// Number of window expressions this state covers.
+    pub fn len(&self) -> usize {
+        self.per_window_expr.len()
+    }
+
+    /// True when no window expression published state.
+    pub fn is_empty(&self) -> bool {
+        self.per_window_expr.is_empty()
+    }
+}
 
 /// How to combine each row's existing value in an output column with a
 /// scheduler-provided scalar offset. The result overwrites the column.
@@ -624,10 +664,10 @@ impl ExecutionPlan for PrefixMergeExec {
         let input_schema = self.input.schema();
         let output_schema = self.schema();
 
-        // The DF-side getter (BoundedWindowAggExec::finalized_partition_state)
-        // already commits to at-most-one PARTITION BY group per DF partition,
-        // so this operator receives the state indexed only by window
-        // expression — no PARTITION BY dimension to project away.
+        // Indexed by window expression only. The scheduler rejects any
+        // report carrying a PARTITION BY key, so at most one group per
+        // partition reaches here and there is no key dimension to project
+        // away.
         let key_state = &resolved[partition];
 
         let mut appliers: Vec<PreparedApply> = Vec::with_capacity(self.applies.len());
@@ -639,9 +679,7 @@ impl ExecutionPlan for PrefixMergeExec {
                     output_column,
                     window_expr_index,
                 } => {
-                    let offset_state: Option<&Vec<ScalarValue>> = key_state
-                        .get(*window_expr_index)
-                        .and_then(|slot| slot.as_ref());
+                    let offset_state = key_state.slot(*window_expr_index);
                     appliers.push(PreparedApply::Aggregate(AggregateApply::new(
                         i,
                         udf,
@@ -1042,16 +1080,22 @@ mod tests {
     }
 
     fn empty_state(partitions: usize) -> Vec<FinalizedPartitionState> {
-        (0..partitions).map(|_| Vec::new()).collect()
+        (0..partitions)
+            .map(|_| FinalizedPartitionState::default())
+            .collect()
     }
 
     /// Mismatched state length surfaces as an error rather than a panic at
     /// runtime.
     #[test]
     fn try_new_rejects_state_length_mismatch() {
         let input = partitioned_source(2, 3);
-        let err = PrefixMergeExec::try_new_resolved(input, vec![], vec![Vec::new()])
-            .expect_err("length mismatch must surface as an error");
+        let err = PrefixMergeExec::try_new_resolved(
+            input,
+            vec![],
+            vec![FinalizedPartitionState::default()],
+        )
+        .expect_err("length mismatch must surface as an error");
         assert!(
             err.to_string()
                 .contains("does not match input partition count"),
@@ -1176,8 +1220,10 @@ mod tests {
         // Partition 1 gets partition 0's HLL state at the one aggregate slot
         // (window_expr_index 0). Partition 0 gets an empty state — nothing
         // to merge in.
-        let per_partition_state: Vec<FinalizedPartitionState> =
-            vec![vec![], vec![Some(p0_state)]];
+        let per_partition_state: Vec<FinalizedPartitionState> = vec![
+            FinalizedPartitionState::default(),
+            FinalizedPartitionState::new(vec![Some(p0_state)]),
+        ];
 
         let apply = WindowApply::Aggregate {
             udf: Arc::clone(&udf),
@@ -1300,8 +1346,10 @@ mod tests {
         };
         // Partition 1 gets partition 0's (sum, count) state at the one
         // aggregate slot; partition 0 has nothing to merge in.
-        let per_partition_state: Vec<FinalizedPartitionState> =
-            vec![vec![], vec![Some(p0_state)]];
+        let per_partition_state: Vec<FinalizedPartitionState> = vec![
+            FinalizedPartitionState::default(),
+            FinalizedPartitionState::new(vec![Some(p0_state)]),
+        ];
 
         let apply = WindowApply::Aggregate {
             udf: Arc::clone(&udf),

ballista/core/src/execution_plans/window_state.rs

@@ -294,8 +294,11 @@ pub fn prefix_merge_window_state(
             window_state_from_proto(&tagged.report)?;
         // `FinalizedPartitionState` carries no PARTITION BY dimension, so a
         // second group in one partition has nowhere to go. The prefix rule
-        // only plants no-PARTITION-BY windows today; this is the assertion
-        // that keeps that gate honest rather than silently merging groups.
+        // only plants no-PARTITION-BY windows, and a window that has one
+        // needs no prefix scan anyway — `BoundedWindowAggExec` asks for
+        // `KeyPartitioned` input, so each partition's window is already
+        // independent. This assertion keeps that gate honest rather than
+        // silently merging two groups.
         if !partition_key.is_empty() {
             return internal_err!(
                 "prefix merge: window state for partition {partition} carries a \
@@ -336,7 +339,7 @@ pub fn prefix_merge_window_state(
             // function publishes nothing.
             let carried = partition
                 .checked_sub(1)
-                .and_then(|prior| prefixes[prior][expr_index].as_ref());
+                .and_then(|prior| prefixes[prior].slot(expr_index));
             let preceding = partition
                 .checked_sub(1)
                 .and_then(|prior| states.get(&(prior, expr_index)));
@@ -362,7 +365,7 @@ pub fn prefix_merge_window_state(
             }
             per_expr.push(Some(accumulator.state()?));
         }
-        prefixes.push(per_expr);
+        prefixes.push(FinalizedPartitionState::new(per_expr));
     }
     Ok(prefixes)
 }
@@ -427,7 +430,7 @@ mod tests {
     fn sums(prefixes: &[FinalizedPartitionState]) -> Vec<Option<f64>> {
         prefixes
             .iter()
-            .map(|per_expr| match &per_expr[0] {
+            .map(|per_expr| match per_expr.slot(0) {
                 None => None,
                 Some(state) => match &state[0] {
                     ScalarValue::Float64(v) => *v,

ballista/core/src/serde/mod.rs

@@ -410,7 +410,7 @@ fn encode_prefix_state(
     let mut partitions = Vec::with_capacity(state.len());
     for partition in state {
         let mut slots = Vec::with_capacity(partition.len());
-        for slot in partition {
+        for slot in partition.slots() {
             let state = match slot {
                 Some(values) => Some(protobuf::AggregateStateNode {
                     values: encode_scalars(values, "prefix state")?,
@@ -438,7 +438,7 @@ fn decode_prefix_state(
             };
             slots.push(state);
         }
-        partitions.push(slots);
+        partitions.push(FinalizedPartitionState::new(slots));
     }
     Ok(partitions)
 }
@@ -1663,7 +1663,10 @@ mod test {
         ];
         // One partition; slot 0 carries state, slot 1 is a window function
         // that published none. `None` and `Some(vec![])` must stay distinct.
-        let state = vec![vec![Some(vec![ScalarValue::Float64(Some(42.0))]), None]];
+        let state = vec![FinalizedPartitionState::new(vec![
+            Some(vec![ScalarValue::Float64(Some(42.0))]),
+            None,
+        ])];
         let original =
             PrefixMergeExec::try_new_resolved(input, applies, state.clone()).unwrap();