From 65014d5ec9ac95351da684a71b2ae8ab1919b02c Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 31 Jul 2026 12:53:38 -0600 Subject: [PATCH 1/2] Expose accumulator state to allow prefix scanning --- .../physical-expr/src/window/window_expr.rs | 13 ++ .../src/windows/bounded_window_agg_exec.rs | 206 +++++++++++++++++- 2 files changed, 215 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 8db5651346e8..a78611c3335e 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -610,6 +610,19 @@ pub struct WindowState { pub state: WindowAggState, pub window_fn: WindowFn, } + +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). + pub fn aggregate_state(&mut self) -> Result>> { + match &mut self.window_fn { + WindowFn::Aggregate(accumulator) => accumulator.state().map(Some), + WindowFn::Builtin(_) => Ok(None), + } + } +} + pub type PartitionWindowAggStates = IndexMap; /// The IndexMap (i.e. an ordered HashMap) where record batches are separated for each partition. 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 07751a70ecee..65d8de866543 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -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,8 +75,18 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; +/// Called by [`BoundedWindowAggExec`] when a PARTITION BY group closes, +/// once per (output-partition-index, PARTITION BY tuple). The third argument +/// is one entry per window expression on the exec, in the same order as +/// [`BoundedWindowAggExec::window_expr`]; `None` for built-in functions +/// (`row_number`, `rank`, `lead`/`lag`, ...), `Some(state)` for aggregates +/// via [`datafusion_expr::Accumulator::state`]. +pub type FinalizedWindowStateObserver = Arc< + dyn Fn(usize, &PartitionKey, &[Option>]) -> Result<()> + Send + Sync, +>; + /// Window execution plan -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct BoundedWindowAggExec { /// Input plan input: Arc, @@ -99,6 +109,32 @@ pub struct BoundedWindowAggExec { cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, + /// Invoked at partition-close to publish finalized per-partition window + /// state. Storage and multi-group handling are the caller's; the exec is + /// a pure event source. + finalized_state_observer: Option, +} + +impl std::fmt::Debug for BoundedWindowAggExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BoundedWindowAggExec") + .field("input", &self.input) + .field("window_expr", &self.window_expr) + .field("schema", &self.schema) + .field("metrics", &self.metrics) + .field("input_order_mode", &self.input_order_mode) + .field( + "ordered_partition_by_indices", + &self.ordered_partition_by_indices, + ) + .field("cache", &self.cache) + .field("can_repartition", &self.can_repartition) + .field( + "finalized_state_observer", + &self.finalized_state_observer.as_ref().map(|_| "..."), + ) + .finish() + } } impl BoundedWindowAggExec { @@ -139,9 +175,20 @@ impl BoundedWindowAggExec { ordered_partition_by_indices, cache: Arc::new(cache), can_repartition, + finalized_state_observer: None, }) } + /// Install a callback that receives each PARTITION BY group's finalized + /// window state at partition close. See [`FinalizedWindowStateObserver`]. + pub fn with_finalized_state_observer( + mut self, + observer: FinalizedWindowStateObserver, + ) -> Self { + self.finalized_state_observer = Some(observer); + self + } + /// Window expressions pub fn window_expr(&self) -> &[Arc] { &self.window_expr @@ -345,12 +392,16 @@ impl ExecutionPlan for BoundedWindowAggExec { children: Vec>, ) -> Result> { check_if_same_properties!(self, children); - Ok(Arc::new(BoundedWindowAggExec::try_new( + let mut new = BoundedWindowAggExec::try_new( self.window_expr.clone(), Arc::clone(&children[0]), self.input_order_mode.clone(), self.can_repartition, - )?)) + )?; + if let Some(observer) = &self.finalized_state_observer { + new = new.with_finalized_state_observer(Arc::clone(observer)); + } + Ok(Arc::new(new)) } fn with_new_children_and_same_properties( @@ -377,6 +428,8 @@ impl ExecutionPlan for BoundedWindowAggExec { input, BaselineMetrics::new(&self.metrics, partition), search_mode, + partition, + self.finalized_state_observer.clone(), )?); Ok(stream) } @@ -1010,6 +1063,13 @@ pub struct BoundedWindowAggStream { /// Search mode for partition columns. This determines the algorithm with /// which we group each partition. search_mode: Box, + /// Output partition index this stream serves; passed as the first + /// argument to `finalized_state_observer`. + partition_idx: usize, + /// If set, invoked at partition close in [`Self::prune_state`] with the + /// finalized per-window-expression state for every partition key that is + /// about to be dropped. + finalized_state_observer: Option, } impl BoundedWindowAggStream { @@ -1021,6 +1081,33 @@ 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<()> { + // `WindowAggState::is_end` is copied from `PartitionBatchState::is_end` + // during `evaluate_stateful`, so every window_expr's state map agrees + // on which partition keys have just closed. Publish those keys before + // the retains in `prune_out_columns` / `prune_partition_batches` drop + // them. + if let Some(observer) = self.finalized_state_observer.clone() + && !self.window_agg_states.is_empty() + { + let closed_keys: Vec = self.window_agg_states[0] + .iter() + .filter(|(_, ws)| ws.state.is_end) + .map(|(k, _)| k.clone()) + .collect(); + for key in closed_keys { + let mut states: Vec>> = + Vec::with_capacity(self.window_agg_states.len()); + for per_expr in self.window_agg_states.iter_mut() { + let entry = per_expr.get_mut(&key).ok_or_else(|| { + exec_datafusion_err!( + "finalized_state_observer: missing state for closed partition key" + ) + })?; + states.push(entry.aggregate_state()?); + } + observer(self.partition_idx, &key, &states)?; + } + } // Prune `self.window_agg_states`: self.prune_out_columns(); // Prune `self.partition_batches`: @@ -1053,6 +1140,8 @@ impl BoundedWindowAggStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, search_mode: Box, + partition_idx: usize, + finalized_state_observer: Option, ) -> Result { let state = window_expr.iter().map(|_| IndexMap::default()).collect(); let empty_batch = RecordBatch::new_empty(Arc::clone(&schema)); @@ -1066,6 +1155,8 @@ impl BoundedWindowAggStream { window_expr, baseline_metrics, search_mode, + partition_idx, + finalized_state_observer, }) } @@ -1908,6 +1999,113 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_finalized_state_observer_fires_at_partition_close() -> Result<()> { + use crate::windows::bounded_window_agg_exec::FinalizedWindowStateObserver; + use datafusion_physical_expr::window::PartitionKey; + use std::sync::Mutex; + + let task_ctx = Arc::new(TaskContext::default()); + let schema = test_schema(); + + // Two PARTITION BY groups: hash=1 [sn=1,2,3] then hash=2 [sn=4,5,6]. + // Input is sorted by (hash, sn) so we can run in Sorted mode; in that + // mode `mark_partition_end` closes the leading group mid-stream and + // EOS closes the tail — both should fire the observer. + let mut sn_b = UInt64Builder::with_capacity(6); + let mut hash_b = Int64Builder::with_capacity(6); + for (sn, hash) in [(1u64, 1i64), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)] { + sn_b.append_value(sn); + hash_b.append_value(hash); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [ + PhysicalSortExpr { + expr: col("hash", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }, + ] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let window_fn = WindowFunctionDefinition::AggregateUDF(count_udaf()); + let args = vec![col("sn", &schema)?]; + let partition_by = vec![col("hash", &schema)?]; + let order_by = vec![PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }]; + // CURRENT ROW → UNBOUNDED FOLLOWING forces each row's output to wait + // for partition close (is_causal = false), so hash=2's rows are held + // until EOS marks its buffer is_end — which is the path we want to + // exercise for the observer. + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(None)), + ); + let expr = create_window_expr( + &window_fn, + "cnt".to_string(), + &args, + &partition_by, + &order_by, + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + type Observation = (usize, PartitionKey, Vec>>); + let observations: Arc>> = Arc::new(Mutex::new(vec![])); + let sink = Arc::clone(&observations); + let observer: FinalizedWindowStateObserver = Arc::new(move |idx, key, states| { + sink.lock() + .unwrap() + .push((idx, key.clone(), states.to_vec())); + Ok(()) + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_finalized_state_observer(observer); + + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let obs = observations.lock().unwrap(); + assert_eq!(obs.len(), 2, "one observation per PARTITION BY group"); + let keys: Vec = obs + .iter() + .map(|(_, k, _)| match &k[0] { + ScalarValue::Int64(Some(v)) => *v, + other => panic!("unexpected partition-key element: {other:?}"), + }) + .collect(); + assert_eq!(keys, vec![1, 2]); + for (idx, _, states) in obs.iter() { + assert_eq!(*idx, 0, "single output partition"); + assert_eq!(states.len(), 1, "one window expression"); + assert!(states[0].is_some(), "count() is an aggregate → Some(state)"); + } + Ok(()) + } + #[test] fn test_bounded_window_agg_cardinality_effect() -> Result<()> { let schema = test_schema(); From 0767f739762012d49d57b11e34e913c66c866cb0 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Sat, 1 Aug 2026 10:11:20 -0600 Subject: [PATCH 2/2] docs --- datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 65d8de866543..ac501a0de14c 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -180,7 +180,7 @@ impl BoundedWindowAggExec { } /// Install a callback that receives each PARTITION BY group's finalized - /// window state at partition close. See [`FinalizedWindowStateObserver`]. + /// window state at partition close. pub fn with_finalized_state_observer( mut self, observer: FinalizedWindowStateObserver,