diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 0f0ec647a50ae..050e7c3bddb06 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -609,6 +609,17 @@ pub struct WindowState { pub state: WindowAggState, pub window_fn: WindowFn, } + +impl WindowState { + /// `Accumulator::state()` if this is an aggregate window; `None` otherwise. + pub fn aggregate_state(&mut self) -> Result>> { + match &mut self.window_fn { + WindowFn::Aggregate(acc) => acc.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 d5863080895f6..1ce26f6b32ce5 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, OnceLock}; use std::task::{Context, Poll}; use super::utils::create_schema; @@ -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,6 +75,11 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; +/// Per-output-partition slots for finalized `Accumulator::state()`, indexed by +/// window expression. Outer `Vec` is one entry per window expression; inner +/// `Option>` is `Some` for aggregates, `None` otherwise. +type FinalizedPartitionStates = Arc<[OnceLock>>>]>; + /// Window execution plan #[derive(Debug, Clone)] pub struct BoundedWindowAggExec { @@ -99,6 +104,9 @@ pub struct BoundedWindowAggExec { cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, + /// Populated by the stream at partition-close; read via + /// [`Self::finalized_partition_state`]. + finalized_state: FinalizedPartitionStates, } impl BoundedWindowAggExec { @@ -130,6 +138,10 @@ impl BoundedWindowAggExec { } }; let cache = Self::compute_properties(&input, &schema, &window_expr)?; + let finalized_state: FinalizedPartitionStates = + (0..cache.partitioning.partition_count()) + .map(|_| OnceLock::new()) + .collect(); Ok(Self { input, window_expr, @@ -139,6 +151,7 @@ impl BoundedWindowAggExec { ordered_partition_by_indices, cache: Arc::new(cache), can_repartition, + finalized_state, }) } @@ -152,6 +165,24 @@ impl BoundedWindowAggExec { &self.input } + /// Finalized `Accumulator::state()` per aggregate window expression on + /// output partition `partition`. `None` for the outer option means the + /// partition is out of range or the stream hasn't drained; inner `None` + /// slots are non-aggregate window functions. + /// + /// Intended to support prefix scans over ordered window aggregates. + /// Assumes at most one PARTITION BY group closes per output partition; + /// additional groups are silently dropped. + pub fn finalized_partition_state( + &self, + partition: usize, + ) -> Option<&[Option>]> { + self.finalized_state + .get(partition)? + .get() + .map(Vec::as_slice) + } + /// Return the output sort order of partition keys: For example /// OVER(PARTITION BY a, ORDER BY b) -> would give sorting of the column a // We are sure that partition by columns are always at the beginning of sort_keys @@ -377,6 +408,8 @@ impl ExecutionPlan for BoundedWindowAggExec { input, BaselineMetrics::new(&self.metrics, partition), search_mode, + Arc::clone(&self.finalized_state), + partition, )?); Ok(stream) } @@ -1010,6 +1043,11 @@ 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`]; write once, + /// first-close-wins. + finalized_state: FinalizedPartitionStates, + /// Which slot in [`Self::finalized_state`] this stream writes to. + partition: usize, } impl BoundedWindowAggStream { @@ -1021,6 +1059,9 @@ 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 any partition/window-expr pair + // whose entry is about to be dropped, before the retains below fire. + self.publish_finalized_state()?; // Prune `self.window_agg_states`: self.prune_out_columns(); // Prune `self.partition_batches`: @@ -1031,6 +1072,37 @@ impl BoundedWindowAggStream { self.search_mode.prune(n_out); Ok(()) } + + /// Bundle each window expression's `Accumulator::state()` for the closing + /// group into one `Vec>>` and try to publish it to + /// our slot. `OnceLock::set` is first-write-wins, so subsequent groups on + /// the same output partition are silently dropped. + fn publish_finalized_state(&mut self) -> Result<()> { + let Some(slot) = self.finalized_state.get(self.partition) else { + return Ok(()); + }; + if slot.get().is_some() { + return Ok(()); + } + let n_exprs = self.window_expr.len(); + let mut per_expr: Vec>> = vec![None; n_exprs]; + let mut any_end = false; + for (expr_idx, window_agg_state) in self.window_agg_states.iter_mut().enumerate() + { + for (_key, ws) in window_agg_state.iter_mut() { + if !ws.state.is_end { + continue; + } + per_expr[expr_idx] = ws.aggregate_state()?; + any_end = true; + break; + } + } + if any_end { + let _ = slot.set(per_expr); + } + Ok(()) + } } impl Stream for BoundedWindowAggStream { @@ -1053,6 +1125,8 @@ impl BoundedWindowAggStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, search_mode: Box, + finalized_state: FinalizedPartitionStates, + partition: usize, ) -> Result { let state = window_expr.iter().map(|_| IndexMap::new()).collect(); let empty_batch = RecordBatch::new_empty(Arc::clone(&schema)); @@ -1066,6 +1140,8 @@ impl BoundedWindowAggStream { window_expr, baseline_metrics, search_mode, + finalized_state, + partition, }) } @@ -1924,4 +2000,81 @@ mod tests { )); Ok(()) } + + async fn run_sum_over_pk(rows: &[(i64, i64)]) -> 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 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(pks)), + Arc::new(Int64Array::from(vs)), + ], + )?; + + 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?; + Ok(exec) + } + + #[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) + .expect("stream drained; slot should be populated"); + 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))); + + assert!(exec.finalized_partition_state(99).is_none()); + Ok(()) + } }