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..e6316c411ec57 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, internal_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,53 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; +/// 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 = 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)] pub struct BoundedWindowAggExec { @@ -99,6 +147,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 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 { @@ -130,6 +188,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(FinalStateSlot::default())) + .collect::>() + .into(); Ok(Self { input, window_expr, @@ -139,6 +202,7 @@ impl BoundedWindowAggExec { ordered_partition_by_indices, cache: Arc::new(cache), can_repartition, + finalized_state, }) } @@ -235,6 +299,68 @@ impl BoundedWindowAggExec { } } + /// 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 + /// + /// 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. + /// + /// 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( + &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 + ) + })?; + 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 { let win_cols = self.window_expr.len(); let input_cols = self.input.schema().fields().len(); @@ -371,12 +497,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 +1139,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 +1156,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 +1193,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 +1208,8 @@ impl BoundedWindowAggStream { window_expr, baseline_metrics, search_mode, + finalized_state, + partition, }) } @@ -1200,6 +1344,64 @@ 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 + ) + })?; + // 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(()) + } + /// 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 +2110,114 @@ mod tests { Ok(()) } + /// 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; + + 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) + } + + /// 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(), 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();