From 65ac32e925ebf4d9d3ef4389f48eb67d03b23827 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 10:16:20 -0600 Subject: [PATCH 1/2] feat(physical-plan): expose finalized Accumulator state on BoundedWindowAggExec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add per-output-partition slots that hold each aggregate window function's final Accumulator::state() at partition-close, plus a public getter that distributed executors can call once a task has drained to ship state via a side-channel for cross-partition prefix scans without a two-pass halo-row scheme. The stream snapshots state at the top of prune_state — the last moment `state.is_end` entries are live before either prune_out_columns or prune_partition_batches drops them — so both mid-stream partition close (as SQL PARTITION BY groups end) and EOS are covered by one write site. Non-aggregate window functions (row_number, rank, lead/lag, ...) occupy `None` at their slot in the per-partition-key Vec. No trait changes: this uses the existing Accumulator::state() surface that the group-by Partial/Final protocol already relies on. Behavior without a reader is unchanged — the slots simply hold state that no one reads. Co-Authored-By: Claude Opus 4.7 (1M context) --- datafusion/physical-expr/src/window/mod.rs | 1 + .../src/windows/bounded_window_agg_exec.rs | 217 +++++++++++++++++- 2 files changed, 215 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-expr/src/window/mod.rs b/datafusion/physical-expr/src/window/mod.rs index b45e35440ac20..c54ba332d5cad 100644 --- a/datafusion/physical-expr/src/window/mod.rs +++ b/datafusion/physical-expr/src/window/mod.rs @@ -29,4 +29,5 @@ pub use window_expr::PartitionBatches; pub use window_expr::PartitionKey; pub use window_expr::PartitionWindowAggStates; pub use window_expr::WindowExpr; +pub use window_expr::WindowFn; pub use window_expr::WindowState; diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index d5863080895f6..eaf64045d12a6 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/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, Mutex}; use std::task::{Context, Poll}; use super::utils::create_schema; @@ -54,13 +54,14 @@ 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, + internal_datafusion_err, }; use datafusion_execution::TaskContext; use datafusion_expr::ColumnarValue; use datafusion_expr::window_state::{PartitionBatchState, WindowAggState}; use datafusion_physical_expr::window::{ - PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowState, + PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowFn, WindowState, }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ @@ -75,6 +76,14 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; +/// One output partition's snapshot of finalized [`Accumulator::state`] values, +/// keyed by PARTITION BY tuple. The inner `Vec` is indexed by +/// [`BoundedWindowAggExec::window_expr`]; non-aggregate window functions +/// (row_number, rank, lead/lag, ...) occupy `None`. +/// +/// [`Accumulator::state`]: datafusion_expr::Accumulator::state +pub type FinalizedPartitionState = HashMap>>>; + /// Window execution plan #[derive(Debug, Clone)] pub struct BoundedWindowAggExec { @@ -99,6 +108,16 @@ pub struct BoundedWindowAggExec { cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, + /// Per-output-partition slots holding the finalized `Accumulator::state()` + /// for each aggregate window expression, keyed by PARTITION BY tuple. + /// The stream writes into slot `p` at partition-close, immediately before + /// internal state would otherwise be pruned. Distributed executors read + /// via [`Self::finalized_partition_state`] once a task has drained, to + /// ship state via a side-channel for cross-partition prefix scans without + /// a two-pass halo scheme. Non-aggregate window functions (row_number, + /// rank, lead/lag, ...) have no `Accumulator` and are stored as `None` + /// at their index in the per-partition-key `Vec`. + finalized_state: Arc<[Mutex]>, } impl BoundedWindowAggExec { @@ -130,6 +149,11 @@ impl BoundedWindowAggExec { } }; let cache = Self::compute_properties(&input, &schema, &window_expr)?; + let partition_count = cache.partitioning.partition_count(); + let finalized_state: Arc<[Mutex>]> = (0..partition_count) + .map(|_| Mutex::new(HashMap::new())) + .collect::>() + .into(); Ok(Self { input, window_expr, @@ -139,6 +163,7 @@ impl BoundedWindowAggExec { ordered_partition_by_indices, cache: Arc::new(cache), can_repartition, + finalized_state, }) } @@ -235,6 +260,40 @@ impl BoundedWindowAggExec { } } + /// Snapshot of the finalized [`Accumulator::state`] for every partition + /// key seen on output partition `partition`, populated as partitions close + /// during streaming and complete once the stream has drained. + /// + /// The outer `HashMap` is keyed by PARTITION BY tuple; the inner `Vec` is + /// indexed by [`Self::window_expr`], with `None` at any slot whose window + /// function is not an aggregate (row_number, rank, lead/lag, ...) and + /// therefore exposes no state. + /// + /// Returns an empty map when no partitions have closed yet. Errors when + /// `partition` is outside the exec's output partitioning or when the slot + /// mutex is poisoned. + /// + /// [`Accumulator::state`]: datafusion_expr::Accumulator::state + pub fn finalized_partition_state( + &self, + partition: usize, + ) -> Result { + let slot = self.finalized_state.get(partition).ok_or_else(|| { + internal_datafusion_err!( + "BoundedWindowAggExec: partition {} out of range (have {})", + partition, + self.finalized_state.len() + ) + })?; + let guard = slot.lock().map_err(|e| { + internal_datafusion_err!( + "BoundedWindowAggExec partition {}: finalized-state mutex poisoned: {e}", + partition + ) + })?; + Ok(guard.clone()) + } + fn statistics_helper(&self, statistics: Statistics) -> Result { let win_cols = self.window_expr.len(); let input_cols = self.input.schema().fields().len(); @@ -371,12 +430,15 @@ impl ExecutionPlan for BoundedWindowAggExec { ) -> Result { let input = self.input.execute(partition, context)?; let search_mode = self.get_search_algo()?; + let finalized_state = Arc::clone(&self.finalized_state); let stream = Box::pin(BoundedWindowAggStream::new( Arc::clone(&self.schema), self.window_expr.clone(), input, BaselineMetrics::new(&self.metrics, partition), search_mode, + finalized_state, + partition, )?); Ok(stream) } @@ -1010,6 +1072,12 @@ pub struct BoundedWindowAggStream { /// Search mode for partition columns. This determines the algorithm with /// which we group each partition. search_mode: Box, + /// Shared with [`BoundedWindowAggExec::finalized_state`]. The stream + /// writes into slot `partition` at partition-close, immediately before + /// the entry is pruned from `window_agg_states`. + finalized_state: Arc<[Mutex]>, + /// Which slot in [`Self::finalized_state`] this stream writes to. + partition: usize, } impl BoundedWindowAggStream { @@ -1021,6 +1089,11 @@ 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 every partition/window-expr pair + // whose entry is about to be dropped, before any of the retain + // calls below fire. Pruning happens both mid-stream (as SQL partitions + // close) and at EOS; this snapshot covers both. + self.publish_finalized_state()?; // Prune `self.window_agg_states`: self.prune_out_columns(); // Prune `self.partition_batches`: @@ -1053,6 +1126,8 @@ impl BoundedWindowAggStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, search_mode: Box, + finalized_state: Arc<[Mutex]>, + partition: usize, ) -> Result { let state = window_expr.iter().map(|_| IndexMap::new()).collect(); let empty_batch = RecordBatch::new_empty(Arc::clone(&schema)); @@ -1066,6 +1141,8 @@ impl BoundedWindowAggStream { window_expr, baseline_metrics, search_mode, + finalized_state, + partition, }) } @@ -1200,6 +1277,58 @@ impl BoundedWindowAggStream { } } + /// Snapshot [`Accumulator::state`] for every entry whose `state.is_end` is + /// set, i.e. every partition/window-expr pair about to be dropped by the + /// retain in [`Self::prune_partition_batches`]. The result is written into + /// this stream's slot in the exec's shared `finalized_state`. + /// + /// [`Accumulator::state`]: datafusion_expr::Accumulator::state + fn publish_finalized_state(&mut self) -> Result<()> { + let n_exprs = self.window_expr.len(); + let mut finalizing: HashMap>>> = + HashMap::new(); + for (expr_idx, window_agg_state) in self.window_agg_states.iter_mut().enumerate() + { + for (partition_row, WindowState { state, window_fn }) in + window_agg_state.iter_mut() + { + if !state.is_end { + continue; + } + let acc_state = match window_fn { + WindowFn::Aggregate(acc) => Some(acc.state()?), + WindowFn::Builtin(_) => None, + }; + let entry = finalizing + .entry(partition_row.clone()) + .or_insert_with(|| vec![None; n_exprs]); + entry[expr_idx] = acc_state; + } + } + if finalizing.is_empty() { + return Ok(()); + } + let slot = self.finalized_state.get(self.partition).ok_or_else(|| { + internal_datafusion_err!( + "BoundedWindowAggStream: partition {} out of range for \ + finalized_state (have {})", + self.partition, + self.finalized_state.len() + ) + })?; + let mut guard = slot.lock().map_err(|e| { + internal_datafusion_err!( + "BoundedWindowAggStream partition {}: finalized-state mutex \ + poisoned: {e}", + self.partition + ) + })?; + for (key, per_expr) in finalizing { + guard.insert(key, per_expr); + } + Ok(()) + } + /// Prunes the section of the input batch whose aggregate results /// are calculated and emitted. fn prune_input_batch(&mut self, n_out: usize) -> Result<()> { @@ -1908,6 +2037,88 @@ mod tests { Ok(()) } + /// After the stream drains, `finalized_partition_state` returns one + /// `Accumulator::state()` per PARTITION BY key seen — the per-partition + /// final sum for a cumulative `SUM(v) OVER (PARTITION BY pk)`. Confirms + /// state survives both the mid-stream partition-close prune (when the + /// sort keys change between rows) and the EOS flush. + #[tokio::test] + async fn finalized_partition_state_captures_sum_per_partition_key() -> Result<()> { + 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 batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1, 1, 1, 2, 2, 2])), + Arc::new(Int64Array::from(vec![10, 20, 30, 100, 200, 300])), + ], + )?; + + let sort_info: LexOrdering = [PhysicalSortExpr { + expr: col("pk", &schema)?, + options: SortOptions::default(), + }] + .into(); + let input: Arc = 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 = 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 = 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?; + + let state = exec.finalized_partition_state(0)?; + assert_eq!(state.len(), 2, "one entry per PARTITION BY key"); + let sum_of = |key: i64| -> ScalarValue { + let states = state + .get(&vec![ScalarValue::Int64(Some(key))]) + .unwrap_or_else(|| panic!("partition key {key} missing from state")); + assert_eq!(states.len(), 1, "one window expr"); + let inner = states[0] + .as_ref() + .expect("SUM is an aggregate — Accumulator state should be Some"); + assert_eq!(inner.len(), 1, "SumAccumulator::state() is a single scalar"); + inner[0].clone() + }; + assert_eq!(sum_of(1), ScalarValue::Int64(Some(60))); + assert_eq!(sum_of(2), ScalarValue::Int64(Some(600))); + + // A partition slot the exec doesn't own errors, not panics. + assert!(exec.finalized_partition_state(99).is_err()); + Ok(()) + } + #[test] fn test_bounded_window_agg_cardinality_effect() -> Result<()> { let schema = test_schema(); From 61c2aa8e58850f81c4f6b14da187548b8804caef Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 12:56:33 -0600 Subject: [PATCH 2/2] refactor: scope finalized_partition_state to at-most-one PARTITION BY group MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The caller for this getter — cross-DF-partition prefix scan — only applies when at most one PARTITION BY group is active per DF output partition (either no PARTITION BY, or a synthetic single-key one). Queries with a real multi-group PARTITION BY are handled by DataFusion's normal key-partitioned distribution and don't need this API. The previous shape leaked BWAG's internal PARTITION BY tuple keying into the public type, forcing every caller to project it away. The getter now returns Vec>> directly — outer Vec indexed by window expression, inner Vec the aggregate's state — and errors if BWAG ever observed more than one group. Internal storage mirrors that: a 3-state FinalStateSlot enum (Empty / Single / Multi) rather than a HashMap that grows with N groups. Split the existing test into single-group-success and multi-group-error to cover both paths under the new shape. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/windows/bounded_window_agg_exec.rs | 209 +++++++++++++----- 1 file changed, 154 insertions(+), 55 deletions(-) diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index eaf64045d12a6..e6316c411ec57 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -55,7 +55,7 @@ use datafusion_common::utils::{ }; use datafusion_common::{ HashMap, Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, - internal_datafusion_err, + internal_datafusion_err, internal_err, }; use datafusion_execution::TaskContext; use datafusion_expr::ColumnarValue; @@ -76,13 +76,52 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; -/// One output partition's snapshot of finalized [`Accumulator::state`] values, -/// keyed by PARTITION BY tuple. The inner `Vec` is indexed by -/// [`BoundedWindowAggExec::window_expr`]; non-aggregate window functions -/// (row_number, rank, lead/lag, ...) occupy `None`. +/// One output partition's snapshot of finalized [`Accumulator::state`] +/// values, indexed by [`BoundedWindowAggExec::window_expr`]: outer position +/// picks the window expression, `None` at a slot marks a non-aggregate window +/// function (row_number, rank, lead/lag, ...), and the inner `Vec` +/// is whatever [`Accumulator::state`] returned for that aggregate (1 for +/// SUM/COUNT/MIN/MAX, 2 for AVG's (sum, count), N for sketch-backed / higher- +/// moment aggregates). +/// +/// The public API doesn't expose the PARTITION BY tuple BWAG uses to key its +/// internal state — see [`BoundedWindowAggExec::finalized_partition_state`] +/// for why (short version: cross-DF-partition prefix scan is the caller, and +/// it only applies when at most one PARTITION BY group is active per DF +/// partition; queries with real PARTITION BY get handled by DF's normal +/// key-based distribution and don't need this getter). /// /// [`Accumulator::state`]: datafusion_expr::Accumulator::state -pub type FinalizedPartitionState = HashMap>>>; +pub type FinalizedPartitionState = Vec>>; + +/// One slot's write-side view. Holds at most one state — the getter is +/// intentionally scoped to the at-most-one-group case, so there's no point +/// carrying N entries around when we'd refuse to combine them anyway. If a +/// second distinct PARTITION BY group is observed, we transition to `Multi` +/// and drop the state (the getter will error on read). +#[derive(Debug, Default)] +enum FinalStateSlot { + /// No PARTITION BY group has closed on this DF partition yet. + #[default] + Empty, + /// Exactly one group has closed; its state is captured here. + Single(FinalizedPartitionState), + /// More than one group has closed. States discarded — this getter + /// doesn't have a well-defined answer to give and the reader will + /// error. + Multi, +} + +impl FinalStateSlot { + /// Absorb one group's finalized state. Transitions: + /// `Empty` → `Single(state)`; `Single(_)` → `Multi`; `Multi` → `Multi`. + fn observe(&mut self, state: FinalizedPartitionState) { + *self = match std::mem::take(self) { + Self::Empty => Self::Single(state), + Self::Single(_) | Self::Multi => Self::Multi, + }; + } +} /// Window execution plan #[derive(Debug, Clone)] @@ -109,15 +148,15 @@ pub struct BoundedWindowAggExec { /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, /// Per-output-partition slots holding the finalized `Accumulator::state()` - /// for each aggregate window expression, keyed by PARTITION BY tuple. - /// The stream writes into slot `p` at partition-close, immediately before - /// internal state would otherwise be pruned. Distributed executors read - /// via [`Self::finalized_partition_state`] once a task has drained, to - /// ship state via a side-channel for cross-partition prefix scans without - /// a two-pass halo scheme. Non-aggregate window functions (row_number, - /// rank, lead/lag, ...) have no `Accumulator` and are stored as `None` - /// at their index in the per-partition-key `Vec`. - finalized_state: Arc<[Mutex]>, + /// for each aggregate window expression, keyed internally by PARTITION BY + /// tuple (BWAG's natural write shape). The stream writes into slot `p` at + /// partition-close, immediately before internal state would otherwise be + /// pruned. Distributed executors read via + /// [`Self::finalized_partition_state`] once a task has drained, to ship + /// state via a side-channel for cross-partition prefix scans without a + /// two-pass halo scheme; that getter projects the internal keyed shape + /// down to a `Vec>>` indexed by window expression. + finalized_state: Arc<[Mutex]>, } impl BoundedWindowAggExec { @@ -150,8 +189,8 @@ impl BoundedWindowAggExec { }; let cache = Self::compute_properties(&input, &schema, &window_expr)?; let partition_count = cache.partitioning.partition_count(); - let finalized_state: Arc<[Mutex>]> = (0..partition_count) - .map(|_| Mutex::new(HashMap::new())) + let finalized_state: Arc<[Mutex]> = (0..partition_count) + .map(|_| Mutex::new(FinalStateSlot::default())) .collect::>() .into(); Ok(Self { @@ -260,18 +299,38 @@ impl BoundedWindowAggExec { } } - /// Snapshot of the finalized [`Accumulator::state`] for every partition - /// key seen on output partition `partition`, populated as partitions close - /// during streaming and complete once the stream has drained. + /// Snapshot of the finalized [`Accumulator::state`] for every aggregate + /// window expression on output partition `partition`, populated as + /// partitions close during streaming and complete once the stream has + /// drained. + /// + /// The returned [`FinalizedPartitionState`] is indexed by + /// [`Self::window_expr`]; slots for non-aggregate window functions + /// (`row_number`, `rank`, `lead`/`lag`, ...) are `None`. The inner + /// `Vec` is whatever `Accumulator::state()` returned for + /// that aggregate — a variable count of scalars depending on the shape + /// of its state. + /// + /// # Scoping to a single PARTITION BY group /// - /// The outer `HashMap` is keyed by PARTITION BY tuple; the inner `Vec` is - /// indexed by [`Self::window_expr`], with `None` at any slot whose window - /// function is not an aggregate (row_number, rank, lead/lag, ...) and - /// therefore exposes no state. + /// BWAG's internal storage is keyed by PARTITION BY tuple — one entry + /// per SQL partition group. This getter is intended for the + /// cross-DF-partition prefix-scan use case where at most one such group + /// is active per DF output partition (either because the query has no + /// PARTITION BY, or because the caller has arranged for a synthetic + /// single-key PARTITION BY). Queries with a real multi-group PARTITION + /// BY are handled by DataFusion's normal key-partitioned distribution + /// and don't need this API. /// - /// Returns an empty map when no partitions have closed yet. Errors when - /// `partition` is outside the exec's output partitioning or when the slot - /// mutex is poisoned. + /// Behavior: + /// - No partitions closed yet → returns `vec![None; window_expr.len()]` + /// (indistinguishable from a plan where every window expression is + /// non-aggregate — that's fine for the intended callers, since both + /// cases correctly mean "no state to merge in"). + /// - Exactly one PARTITION BY group closed → returns its state. + /// - More than one group closed → errors. + /// - `partition` outside the exec's output partitioning → errors. + /// - Slot mutex poisoned → errors. /// /// [`Accumulator::state`]: datafusion_expr::Accumulator::state pub fn finalized_partition_state( @@ -291,7 +350,15 @@ impl BoundedWindowAggExec { partition ) })?; - Ok(guard.clone()) + match &*guard { + FinalStateSlot::Empty => Ok(vec![None; self.window_expr.len()]), + FinalStateSlot::Single(state) => Ok(state.clone()), + FinalStateSlot::Multi => internal_err!( + "BoundedWindowAggExec partition {partition}: finalized state \ + observed more than one PARTITION BY group; this getter is \ + scoped to the at-most-one-group case (see docs)" + ), + } } fn statistics_helper(&self, statistics: Statistics) -> Result { @@ -1075,7 +1142,7 @@ pub struct BoundedWindowAggStream { /// Shared with [`BoundedWindowAggExec::finalized_state`]. The stream /// writes into slot `partition` at partition-close, immediately before /// the entry is pruned from `window_agg_states`. - finalized_state: Arc<[Mutex]>, + finalized_state: Arc<[Mutex]>, /// Which slot in [`Self::finalized_state`] this stream writes to. partition: usize, } @@ -1126,7 +1193,7 @@ impl BoundedWindowAggStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, search_mode: Box, - finalized_state: Arc<[Mutex]>, + finalized_state: Arc<[Mutex]>, partition: usize, ) -> Result { let state = window_expr.iter().map(|_| IndexMap::new()).collect(); @@ -1323,8 +1390,14 @@ impl BoundedWindowAggStream { self.partition ) })?; - for (key, per_expr) in finalizing { - guard.insert(key, per_expr); + // We only key the local `finalizing` map by `PartitionKey` to + // collate state across the outer per-window-expr loop; the slot + // itself just needs to know "another group closed". Passing every + // entry through `observe` lets the slot transition to `Multi` if + // more than one distinct key ever appears without holding onto the + // rest. + for (_key, per_expr) in finalizing { + guard.observe(per_expr); } Ok(()) } @@ -2037,13 +2110,11 @@ mod tests { Ok(()) } - /// After the stream drains, `finalized_partition_state` returns one - /// `Accumulator::state()` per PARTITION BY key seen — the per-partition - /// final sum for a cumulative `SUM(v) OVER (PARTITION BY pk)`. Confirms - /// state survives both the mid-stream partition-close prune (when the - /// sort keys change between rows) and the EOS flush. - #[tokio::test] - async fn finalized_partition_state_captures_sum_per_partition_key() -> Result<()> { + /// Helper: builds an exec with a single `SUM(v)` window aggregate, + /// partitioned by `pk`, over the provided `(pk, v)` rows. Runs the + /// stream to completion and returns the exec so callers can inspect + /// `finalized_partition_state`. + async fn run_sum_over_pk(rows: &[(i64, i64)]) -> Result> { use arrow::array::Int64Array; use datafusion_functions_aggregate::sum::sum_udaf; @@ -2051,11 +2122,13 @@ mod tests { Field::new("pk", DataType::Int64, false), Field::new("v", DataType::Int64, false), ])); + let pks: Vec = rows.iter().map(|(pk, _)| *pk).collect(); + let vs: Vec = rows.iter().map(|(_, v)| *v).collect(); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int64Array::from(vec![1, 1, 1, 2, 2, 2])), - Arc::new(Int64Array::from(vec![10, 20, 30, 100, 200, 300])), + Arc::new(Int64Array::from(pks)), + Arc::new(Int64Array::from(vs)), ], )?; @@ -2097,28 +2170,54 @@ mod tests { let ctx = Arc::new(TaskContext::default()); let stream = ExecutionPlan::execute(exec.as_ref(), 0, ctx)?; let _batches = collect(stream).await?; + Ok(exec) + } + + /// With a single PARTITION BY group active on the DF output partition, + /// `finalized_partition_state` returns one `Accumulator::state()` per + /// window expression — for `SUM(v)`, a single `Int64` scalar. Confirms + /// state survives the EOS flush and the getter unwraps the internal + /// keyed storage down to the caller-facing `Vec>>`. + #[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)?; - assert_eq!(state.len(), 2, "one entry per PARTITION BY key"); - let sum_of = |key: i64| -> ScalarValue { - let states = state - .get(&vec![ScalarValue::Int64(Some(key))]) - .unwrap_or_else(|| panic!("partition key {key} missing from state")); - assert_eq!(states.len(), 1, "one window expr"); - let inner = states[0] - .as_ref() - .expect("SUM is an aggregate — Accumulator state should be Some"); - assert_eq!(inner.len(), 1, "SumAccumulator::state() is a single scalar"); - inner[0].clone() - }; - assert_eq!(sum_of(1), ScalarValue::Int64(Some(60))); - assert_eq!(sum_of(2), ScalarValue::Int64(Some(600))); + 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))); // A partition slot the exec doesn't own errors, not panics. assert!(exec.finalized_partition_state(99).is_err()); Ok(()) } + /// When more than one PARTITION BY group closes on the same DF output + /// partition (i.e. the query has a real PARTITION BY and DF's normal + /// key-partitioned distribution has co-located multiple groups), + /// `finalized_partition_state` errors rather than silently picking one. + /// The getter is intentionally scoped to the cross-DF-partition + /// prefix-scan use case, which by construction has at most one group. + #[tokio::test] + async fn finalized_partition_state_errors_on_multi_group() -> Result<()> { + let exec = + run_sum_over_pk(&[(1, 10), (1, 20), (1, 30), (2, 100), (2, 200), (2, 300)]) + .await?; + + let err = exec + .finalized_partition_state(0) + .expect_err("multi-group state must surface as an error"); + assert!( + err.to_string() + .contains("more than one PARTITION BY group"), + "unexpected error: {err}" + ); + Ok(()) + } + #[test] fn test_bounded_window_agg_cardinality_effect() -> Result<()> { let schema = test_schema();