diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 47147b909d342..bcd657c21ee92 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -647,6 +647,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 03a8e9867c170..21f6b7143bab2 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; @@ -76,8 +76,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, @@ -100,6 +110,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 { @@ -140,9 +176,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. + 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 @@ -346,12 +393,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( @@ -378,6 +429,8 @@ impl ExecutionPlan for BoundedWindowAggExec { input, BaselineMetrics::new(&self.metrics, partition), search_mode, + partition, + self.finalized_state_observer.clone(), )?); Ok(stream) } @@ -1042,9 +1095,46 @@ pub struct BoundedWindowAggStream { /// partitions, so finished partitions are pruned eagerly instead and no /// such bound is needed. most_recent_row: Option, + /// Output partition index this stream serves; passed as the first + /// argument to `finalized_state_observer`. + partition_idx: usize, + /// If set, invoked from [`Self::publish_finalized_states`] with the + /// finalized per-window-expression state for every partition key that is + /// about to be dropped. + finalized_state_observer: Option, } impl BoundedWindowAggStream { + /// Fire `finalized_state_observer` for every partition key whose + /// `WindowAggState::is_end` is true. + fn publish_finalized_states(&mut self) -> Result<()> { + let Some(observer) = self.finalized_state_observer.clone() else { + return Ok(()); + }; + if self.window_agg_states.is_empty() { + return Ok(()); + } + 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)?; + } + Ok(()) + } + /// Prunes sections of the state that are no longer needed when calculating /// results (as determined by window frame boundaries and number of results generated). // For instance, if first `n` (not necessarily same with `n_out`) elements are no longer needed to @@ -1085,6 +1175,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)); @@ -1099,6 +1191,8 @@ impl BoundedWindowAggStream { baseline_metrics, search_mode, most_recent_row: None, + partition_idx, + finalized_state_observer, }) } @@ -1116,6 +1210,12 @@ impl BoundedWindowAggStream { )?; } + // Fire before `calculate_out_columns`: on causal frames every row + // already streamed out, so at EOS that call returns `None` and the + // prune path is skipped — the final partition would otherwise be + // dropped unobserved. + self.publish_finalized_states()?; + let schema = Arc::clone(&self.schema); let window_expr_out = self.search_mode.calculate_out_columns( &self.input_buffer, @@ -1955,6 +2055,386 @@ 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(()) + } + + #[tokio::test] + async fn test_finalized_state_observer_fires_on_causal_frame() -> Result<()> { + // Same setup as the non-causal test above, but with the causal frame + // `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` — the running-sum + // shape. Output can stream row-by-row without waiting for partition + // close, so this is the frame most likely to bypass the observer path + // if `is_end` were tied to buffered-output flushing. It isn't: + // `is_end` comes from PARTITION BY transition detection in + // `mark_partition_end` (mid-stream) and EOS (tail), independent of + // frame causality. This test locks that in. + 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(); + + 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(), + }]; + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ); + 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(()) + } + + /// Run one task's local BWAG for `SUM(sn) OVER (ORDER BY sn ROWS + /// UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, over + /// `input` sorted ascending. Returns the per-row output values and the + /// observed finalized state total (which the caller uses as a carry-in + /// for the next task). + async fn run_running_sum_task( + input: &[u64], + task_ctx: Arc, + ) -> Result<(Vec, u64)> { + use crate::windows::bounded_window_agg_exec::FinalizedWindowStateObserver; + use arrow::array::UInt64Array; + use datafusion_functions_aggregate::sum::sum_udaf; + use std::sync::Mutex; + + let schema = test_schema(); + let mut sn_b = UInt64Builder::with_capacity(input.len()); + let mut hash_b = Int64Builder::with_capacity(input.len()); + for &sn in input { + sn_b.append_value(sn); + hash_b.append_value(0); + } + 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("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(sum_udaf()); + let args = vec![col("sn", &schema)?]; + let partition_by: Vec> = vec![]; + let order_by = vec![PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }]; + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ); + let expr = create_window_expr( + &window_fn, + "running_sum".to_string(), + &args, + &partition_by, + &order_by, + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + let total_sink: Arc>> = Arc::new(Mutex::new(None)); + let sink = Arc::clone(&total_sink); + let observer: FinalizedWindowStateObserver = + Arc::new(move |_idx, key, states| { + assert!( + key.is_empty(), + "empty PartitionKey for no-PARTITION-BY plan" + ); + assert_eq!(states.len(), 1, "one window expression"); + let state = states[0].as_ref().expect("sum() → Some(state)"); + let total = match &state[0] { + ScalarValue::UInt64(Some(v)) => *v, + ScalarValue::Int64(Some(v)) => *v as u64, + other => panic!("unexpected sum state element: {other:?}"), + }; + let prev = sink.lock().unwrap().replace(total); + assert!(prev.is_none(), "observer must fire exactly once per task"); + Ok(()) + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_finalized_state_observer(observer); + let batches = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let mut out = Vec::with_capacity(input.len()); + for batch in &batches { + let col = batch + .column_by_name("running_sum") + .expect("running_sum column present"); + let arr = col + .as_any() + .downcast_ref::() + .expect("SUM(UInt64) → UInt64Array"); + for i in 0..arr.len() { + out.push(arr.value(i)); + } + } + let total = total_sink + .lock() + .unwrap() + .expect("observer must have fired at EOS"); + Ok((out, total)) + } + + #[tokio::test] + async fn test_prefix_scan_across_tasks_matches_single_bwag() -> Result<()> { + // Demonstrates the parallel-window shape reviewers asked about: + // range-shuffle `SUM(sn) OVER (ORDER BY sn UNBOUNDED PRECEDING TO + // CURRENT ROW)` across two tasks, then prefix-scan each task's + // finalized state (from the observer) to carry-in the next task's + // rows. Result must match a single BWAG over the concatenated input. + let task_ctx = Arc::new(TaskContext::default()); + + // Two tasks under range partition on sn: + let (task1_out, task1_total) = + run_running_sum_task(&[1, 1, 2, 2, 3, 3, 4, 4], Arc::clone(&task_ctx)) + .await?; + let (task2_out, task2_total) = + run_running_sum_task(&[5, 5, 6, 6, 7, 7, 8, 8], Arc::clone(&task_ctx)) + .await?; + + // Local (uncorrected) outputs and totals — first pass. + assert_eq!(task1_out, vec![1, 2, 4, 6, 9, 12, 16, 20]); + assert_eq!(task1_total, 20); + assert_eq!(task2_out, vec![5, 10, 16, 22, 29, 36, 44, 52]); + assert_eq!(task2_total, 52); + + // Prefix scan over per-task totals → carry-in for each task. Task 0's + // carry-in is 0; task N's carry-in is the sum of tasks [0, N). + let carry_ins = [0u64, task1_total]; + + // Second pass: shift each task's local values by its carry-in. + let task1_final: Vec = task1_out.iter().map(|v| v + carry_ins[0]).collect(); + let task2_final: Vec = task2_out.iter().map(|v| v + carry_ins[1]).collect(); + let parallel_result: Vec = task1_final + .iter() + .chain(task2_final.iter()) + .copied() + .collect(); + + // Oracle: single BWAG over the full concatenated input. + let (single_result, single_total) = run_running_sum_task( + &[1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8], + task_ctx, + ) + .await?; + + assert_eq!( + parallel_result, single_result, + "two-task prefix-scan must match single-BWAG oracle" + ); + // And matches the sequence in the design discussion. + assert_eq!( + single_result, + vec![1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49, 56, 64, 72] + ); + assert_eq!(single_total, 72); + Ok(()) + } + #[test] fn test_bounded_window_agg_cardinality_effect() -> Result<()> { let schema = test_schema();