From e80aa810c6d11bdf8e663a72317c24e020bea548 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 10:44:48 -0600 Subject: [PATCH 1/7] feat(core): scaffold PrefixMergeExec for cross-partition windowed-aggregate state merge Introduces PrefixMergeExec as the downstream half of the AQE range-shuffle prefix-scan pipeline: takes per-input-partition finalized window-aggregate state (from BoundedWindowAggExec::finalized_partition_state in apache/datafusion#24007) and, in follow-up work, merges it row-wise into the current partition's output so cross-partition running aggregates come out correct. This commit is scaffolding only. execute() forwards batches unchanged; the merge itself lands in a follow-up. The type signature and constructor are committed now so the surrounding plumbing (scheduler collection of per-task state via task-status transport, plan-rule placement, per-task state injection) can be built against a stable shape. FinalizedPartitionState is defined locally as a stand-in for datafusion::physical_plan::windows::FinalizedPartitionState until apache/datafusion#24007 lands; the doc comment marks the correspondence so the swap is unambiguous when the DF change is available. Co-Authored-By: Claude Opus 4.7 (1M context) --- ballista/core/src/execution_plans/mod.rs | 2 + .../core/src/execution_plans/prefix_merge.rs | 317 ++++++++++++++++++ 2 files changed, 319 insertions(+) create mode 100644 ballista/core/src/execution_plans/prefix_merge.rs diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 0b9e89104..dfa332799 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -24,6 +24,7 @@ mod distributed_explain_analyze; mod distributed_query; mod ordered_range_repartition; mod per_partition_filter; +mod prefix_merge; mod range_repartition_common; mod runtime_stats; mod shuffle_reader; @@ -42,6 +43,7 @@ pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use per_partition_filter::PerPartitionFilterExec; +pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec}; pub use runtime_stats::{ MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, collect_reports as collect_runtime_stats_reports, log_merged_runtime_stats, diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs new file mode 100644 index 000000000..c42dd2f59 --- /dev/null +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -0,0 +1,317 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cross-partition state merge for windowed aggregates in an AQE range-shuffle +//! pipeline. +//! +//! Range-shuffle produces `N` ordered disjoint partitions of a stream sorted +//! by the window's ORDER BY key. Each executor task then runs a windowed +//! aggregate over its slice, producing correct per-row values *within* the +//! slice but not across slices — task `k`'s running SUM starts at zero, not +//! at the sum of everything in partitions `[0..k)`. +//! +//! This operator closes that gap by merging the accumulated state from all +//! earlier partitions into partition `k`'s output. Each row's contribution +//! from prior partitions is folded in via [`Accumulator::merge_batch`]-shaped +//! composition — the same operation the group-by `Partial → Final` protocol +//! already uses to combine per-group state. Aggregate-agnostic by +//! construction: SUM adds, MIN/MAX take the extreme, sketches (KLL, TDigest, +//! HLL) merge as sketches. +//! +//! Input shape (per input partition `k`, provided in [`Self::try_new`]): +//! [`FinalizedPartitionState`] — a snapshot of every window aggregate's +//! finalized [`Accumulator::state`] for every PARTITION BY key seen while +//! draining partitions `[0..k)`. The scheduler bakes this in at plan time +//! after collecting per-task state from the upstream stage via task-status +//! transport. +//! +//! **This is scaffolding.** `execute` currently forwards batches unchanged. +//! The merge itself lands in a follow-up. The shape is committed to now so +//! the surrounding plumbing (scheduler collection, per-task state injection, +//! plan-rule placement) can be built against a stable signature. +//! +//! # Relation to DataFusion +//! +//! The per-input-partition input this operator expects is produced by +//! `BoundedWindowAggExec::finalized_partition_state`, added in +//! [apache/datafusion#24007]. Until that lands, callers can pass empty maps +//! (the operator forwards batches regardless) to exercise the plumbing. +//! +//! [`Accumulator::state`]: datafusion::logical_expr::Accumulator::state +//! [`Accumulator::merge_batch`]: datafusion::logical_expr::Accumulator::merge_batch +//! [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 + +use std::fmt::{self, Debug, Formatter}; +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, ScalarValue, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::window::PartitionKey; +use datafusion::physical_expr::{Distribution, OrderingRequirements}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use std::collections::HashMap; + +/// One input partition's finalized window-aggregate state, keyed by PARTITION +/// BY tuple. The inner `Vec` is indexed by window expression (same order the +/// upstream `BoundedWindowAggExec` reports in `window_expr()`); `None` at a +/// slot indicates a non-aggregate window function (`row_number`, `rank`, +/// `lead`/`lag`, ...) which contributes no state. +/// +/// Must mirror `datafusion::physical_plan::windows::FinalizedPartitionState` +/// once [apache/datafusion#24007] lands; the type alias here is a local +/// stand-in so this crate can compile against stable DataFusion 54. +/// +/// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 +pub type FinalizedPartitionState = HashMap>>>; + +/// Merge finalized window-aggregate state from earlier partitions into each +/// partition's output. See the [module-level docs][self] for the shape of +/// the input and the AQE pipeline this fits into. +/// +/// **Scaffolding.** `execute` forwards batches unchanged; the merge itself +/// lands in a follow-up. +pub struct PrefixMergeExec { + input: Arc, + /// `per_partition_state[k]` is the accumulated finalized state from every + /// input partition in `[0..k)`, ready to be merged into partition `k`'s + /// output. Length equals `input.output_partitioning().partition_count()`. + per_partition_state: Vec, + properties: Arc, +} + +impl PrefixMergeExec { + /// Wrap `input` with per-input-partition prefix state. + /// + /// Errors if `per_partition_state.len()` doesn't match the input's + /// partition count. + pub fn try_new( + input: Arc, + per_partition_state: Vec, + ) -> Result { + let partition_count = input.output_partitioning().partition_count(); + if per_partition_state.len() != partition_count { + return internal_err!( + "PrefixMergeExec: per_partition_state.len() {} does not match \ + input partition count {}", + per_partition_state.len(), + partition_count + ); + } + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + input, + per_partition_state, + properties, + }) + } + + /// The prefix state carried per input partition. + /// `per_partition_state()[k]` is what will be merged into partition `k`. + pub fn per_partition_state(&self) -> &[FinalizedPartitionState] { + &self.per_partition_state + } +} + +impl Debug for PrefixMergeExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("PrefixMergeExec") + .field("partition_count", &self.per_partition_state.len()) + .finish() + } +} + +impl DisplayAs for PrefixMergeExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "PrefixMergeExec: partitions={}", + self.per_partition_state.len() + ) + } + DisplayFormatType::TreeRender => { + write!(f, "PrefixMergeExec") + } + } + } +} + +impl ExecutionPlan for PrefixMergeExec { + fn name(&self) -> &str { + "PrefixMergeExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input] = children.as_slice() else { + return internal_err!( + "PrefixMergeExec expects exactly one child, got {}", + children.len() + ); + }; + Ok(Arc::new(PrefixMergeExec::try_new( + input.clone(), + self.per_partition_state.clone(), + )?)) + } + + /// Passthrough: no distribution requirement on the child. + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + /// Passthrough: no ordering requirement on the child. In practice the + /// upstream range-shuffle already delivers sorted-by-ORDER-BY input; the + /// merge doesn't reorder rows within a partition. + fn required_input_ordering(&self) -> Vec> { + vec![None] + } + + /// Each output row corresponds 1:1 to an input row; the merge only + /// rewrites the window-aggregate columns. + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input.partition_statistics(partition) + } + + /// Every input row is emitted exactly once. + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + /// **Scaffolding.** Forwards the input stream unchanged. The merge itself + /// — folding `self.per_partition_state[partition]` into each row's + /// window-aggregate columns — lands in a follow-up. + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + if partition >= self.per_partition_state.len() { + return internal_err!( + "PrefixMergeExec: partition {} out of bounds ({} slots)", + partition, + self.per_partition_state.len() + ); + } + self.input.execute(partition, ctx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + + fn one_col_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])) + } + + /// Two-partition memory source; partition `k` carries the single batch + /// `[k * rows_per .. (k + 1) * rows_per)`. + fn partitioned_source(partitions: usize, rows_per: usize) -> Arc { + let schema = one_col_schema(); + let mut per_partition: Vec> = Vec::with_capacity(partitions); + for k in 0..partitions { + let start = (k * rows_per) as i64; + let arr = Int64Array::from_iter_values(start..start + rows_per as i64); + per_partition.push(vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(), + ]); + } + let source = MemorySourceConfig::try_new(&per_partition, schema, None).unwrap(); + Arc::new(DataSourceExec::new(Arc::new(source))) + } + + /// Mismatched state length surfaces as an error rather than a panic at + /// runtime. + #[test] + fn try_new_rejects_state_length_mismatch() { + let input = partitioned_source(2, 3); + let err = PrefixMergeExec::try_new(input, vec![FinalizedPartitionState::new()]) + .expect_err("length mismatch must surface as an error"); + assert!( + err.to_string() + .contains("does not match input partition count"), + "unexpected error: {err}" + ); + } + + /// While the merge is unimplemented, `execute` behaves as a passthrough: + /// each partition's rows are forwarded verbatim. + #[tokio::test] + async fn scaffold_forwards_input_rows_unchanged() -> Result<()> { + let input = partitioned_source(2, 3); + let state: Vec = vec![ + FinalizedPartitionState::new(), + FinalizedPartitionState::new(), + ]; + let exec = Arc::new(PrefixMergeExec::try_new(input, state)?); + + let ctx = SessionContext::new().task_ctx(); + for k in 0..2 { + let batches: Vec = + exec.execute(k, ctx.clone())?.try_collect().await?; + assert_eq!(batches.len(), 1); + let arr = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("Int64 column"); + let expected: Vec = ((k * 3) as i64..((k + 1) * 3) as i64).collect(); + assert_eq!(arr.values(), &expected); + } + Ok(()) + } +} From fe8bb2a4edd9b3396d190601f5c629498640dbdb Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 10:51:09 -0600 Subject: [PATCH 2/7] docs(prefix-merge): clarify scheduler/executor division of labor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prefix-merge is inherently a scheduler-side computation (each partition's offset needs global visibility across tasks), so this operator receives *already-merged* state rather than raw per-partition states waiting to be combined. Documents that split explicitly, and notes that for scalar aggregates a plain ProjectionExec with the offset as a literal is equivalent — this operator earns its keep as the aggregate-agnostic escape hatch for cases where the apply isn't a per-row scalar expression. No functional change. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/prefix_merge.rs | 84 ++++++++++++------- 1 file changed, 56 insertions(+), 28 deletions(-) diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index c42dd2f59..5c1df95b6 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -24,29 +24,48 @@ //! slice but not across slices — task `k`'s running SUM starts at zero, not //! at the sum of everything in partitions `[0..k)`. //! -//! This operator closes that gap by merging the accumulated state from all -//! earlier partitions into partition `k`'s output. Each row's contribution -//! from prior partitions is folded in via [`Accumulator::merge_batch`]-shaped -//! composition — the same operation the group-by `Partial → Final` protocol -//! already uses to combine per-group state. Aggregate-agnostic by -//! construction: SUM adds, MIN/MAX take the extreme, sketches (KLL, TDigest, -//! HLL) merge as sketches. +//! # Division of labor +//! +//! The prefix-merge splits cleanly across the scheduler/executor boundary: +//! +//! - **Scheduler (global step):** collects each upstream task's finalized +//! [`Accumulator::state`] via task-status transport, then computes the +//! prefix-merge — for each partition `k`, combining the individual states +//! of partitions `[0..k)` into a single already-merged state per PARTITION +//! BY key and per window expression. This is the step that requires global +//! visibility across tasks; only the scheduler has it. +//! - **Executor (local step, this operator):** receives the *already-merged* +//! state for its partition in its constructor and folds it row-wise into +//! the window-aggregate columns. Aggregate-agnostic by construction — the +//! fold is the same [`Accumulator::merge_batch`]-shaped composition the +//! group-by `Partial → Final` protocol uses, so SUM adds, MIN/MAX take the +//! extreme, sketches (KLL, TDigest, HLL) merge as sketches, without this +//! operator knowing which aggregate is which. +//! +//! For scalar aggregates whose apply reduces to arithmetic (SUM, COUNT, +//! MIN/MAX, and AVG once decomposed to SUM + COUNT + division) the scheduler +//! can instead emit a plain `ProjectionExec` with the offset as a literal, +//! and this operator isn't needed. It exists for cases where the apply is +//! not expressible as a per-row scalar expression — the operator gives an +//! aggregate-agnostic escape hatch that avoids inventing a per-aggregate +//! projection form. //! //! Input shape (per input partition `k`, provided in [`Self::try_new`]): -//! [`FinalizedPartitionState`] — a snapshot of every window aggregate's -//! finalized [`Accumulator::state`] for every PARTITION BY key seen while -//! draining partitions `[0..k)`. The scheduler bakes this in at plan time -//! after collecting per-task state from the upstream stage via task-status -//! transport. +//! [`FinalizedPartitionState`] — the *pre-merged* state, one +//! [`Accumulator::state`] value per (PARTITION BY key, window expression), +//! ready to be folded directly into partition `k`'s rows. The scheduler +//! bakes this in when it constructs the downstream stage after the upstream +//! stage's tasks complete. //! //! **This is scaffolding.** `execute` currently forwards batches unchanged. -//! The merge itself lands in a follow-up. The shape is committed to now so +//! The row-wise fold lands in a follow-up. The shape is committed to now so //! the surrounding plumbing (scheduler collection, per-task state injection, //! plan-rule placement) can be built against a stable signature. //! //! # Relation to DataFusion //! -//! The per-input-partition input this operator expects is produced by +//! The upstream tasks' finalized state — which the scheduler prefix-merges +//! before handing the result to this operator — is produced by //! `BoundedWindowAggExec::finalized_partition_state`, added in //! [apache/datafusion#24007]. Until that lands, callers can pass empty maps //! (the operator forwards batches regardless) to exercise the plumbing. @@ -70,30 +89,38 @@ use datafusion::physical_plan::{ }; use std::collections::HashMap; -/// One input partition's finalized window-aggregate state, keyed by PARTITION +/// A single already-prefix-merged window-aggregate state, keyed by PARTITION /// BY tuple. The inner `Vec` is indexed by window expression (same order the /// upstream `BoundedWindowAggExec` reports in `window_expr()`); `None` at a /// slot indicates a non-aggregate window function (`row_number`, `rank`, /// `lead`/`lag`, ...) which contributes no state. /// -/// Must mirror `datafusion::physical_plan::windows::FinalizedPartitionState` -/// once [apache/datafusion#24007] lands; the type alias here is a local -/// stand-in so this crate can compile against stable DataFusion 54. +/// The scheduler produces one of these per input partition, having already +/// combined the individual states from every prior partition into one merged +/// value per (key, window expr) — see the [module-level docs][self] for the +/// division of labor. This operator applies it; it does not compute it. +/// +/// The type mirrors `datafusion::physical_plan::windows::FinalizedPartitionState` +/// from [apache/datafusion#24007]; the alias here is a local stand-in so +/// this crate compiles against stable DataFusion 54 until that PR lands. /// /// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 pub type FinalizedPartitionState = HashMap>>>; -/// Merge finalized window-aggregate state from earlier partitions into each -/// partition's output. See the [module-level docs][self] for the shape of -/// the input and the AQE pipeline this fits into. +/// Apply pre-merged window-aggregate state (computed by the scheduler) to +/// each row of the current partition's output. See the [module-level +/// docs][self] for the division of labor between scheduler and executor and +/// the AQE pipeline this fits into. /// -/// **Scaffolding.** `execute` forwards batches unchanged; the merge itself +/// **Scaffolding.** `execute` forwards batches unchanged; the row-wise fold /// lands in a follow-up. pub struct PrefixMergeExec { input: Arc, - /// `per_partition_state[k]` is the accumulated finalized state from every - /// input partition in `[0..k)`, ready to be merged into partition `k`'s - /// output. Length equals `input.output_partitioning().partition_count()`. + /// `per_partition_state[k]` is the scheduler-provided *already-merged* + /// state summarising every input partition in `[0..k)`. This operator + /// folds it into partition `k`'s rows; it does not perform the merge + /// itself (the scheduler already did). Length equals + /// `input.output_partitioning().partition_count()`. per_partition_state: Vec, properties: Arc, } @@ -225,9 +252,10 @@ impl ExecutionPlan for PrefixMergeExec { CardinalityEffect::Equal } - /// **Scaffolding.** Forwards the input stream unchanged. The merge itself - /// — folding `self.per_partition_state[partition]` into each row's - /// window-aggregate columns — lands in a follow-up. + /// **Scaffolding.** Forwards the input stream unchanged. The row-wise + /// fold — applying `self.per_partition_state[partition]` (already merged + /// upstream by the scheduler) to each row's window-aggregate columns — + /// lands in a follow-up. fn execute( &self, partition: usize, From 11f2129fecb4fc2db91f894c5dbc26ef9ad8fb7f Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 11:30:46 -0600 Subject: [PATCH 3/7] feat(prefix-merge): add WindowApply / ScalarOp descriptors for per-column apply MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend PrefixMergeExec's constructor with a Vec — one entry per output column that needs cross-partition correction, telling the operator *how* to rewrite that column. Two shapes: - Scalar: monoidal op (Add/Min/Max/Overwrite) between the row's existing value and a scheduler-provided per-partition scalar offset. Fits SUM, COUNT, MIN, MAX, row_number, first_value, last_value. - Aggregate: construct a fresh Accumulator seeded from the pre-merged state, feed args per row, overwrite the column with evaluate(). Fits AVG (without decomposition), sketch-backed windows (APPROX_DISTINCT, APPROX_QUANTILE), statistical aggregates (STDDEV/VAR/covar), and concat aggregates (ARRAY_AGG/STRING_AGG). Non-corrected functions (lead/lag/nth_value via halos in the shuffle layer; rank family via future segment-tree infrastructure) don't appear in the applies list and aren't PrefixMergeExec's concern. The scaffold still forwards batches unchanged — the row-wise fold lands in a follow-up. Two new tests cover Scalar-offset-length and output_column-range validation. Co-Authored-By: Claude Opus 4.7 (1M context) --- ballista/core/src/execution_plans/mod.rs | 2 +- .../core/src/execution_plans/prefix_merge.rs | 264 +++++++++++++++--- 2 files changed, 232 insertions(+), 34 deletions(-) diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index dfa332799..7ab5136f9 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -43,7 +43,7 @@ pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use per_partition_filter::PerPartitionFilterExec; -pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec}; +pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec, ScalarOp, WindowApply}; pub use runtime_stats::{ MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, collect_reports as collect_runtime_stats_reports, log_merged_runtime_stats, diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index 5c1df95b6..859ccfc6c 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -42,19 +42,33 @@ //! extreme, sketches (KLL, TDigest, HLL) merge as sketches, without this //! operator knowing which aggregate is which. //! -//! For scalar aggregates whose apply reduces to arithmetic (SUM, COUNT, -//! MIN/MAX, and AVG once decomposed to SUM + COUNT + division) the scheduler -//! can instead emit a plain `ProjectionExec` with the offset as a literal, -//! and this operator isn't needed. It exists for cases where the apply is -//! not expressible as a per-row scalar expression — the operator gives an -//! aggregate-agnostic escape hatch that avoids inventing a per-aggregate -//! projection form. +//! # Apply descriptors //! -//! Input shape (per input partition `k`, provided in [`Self::try_new`]): -//! [`FinalizedPartitionState`] — the *pre-merged* state, one -//! [`Accumulator::state`] value per (PARTITION BY key, window expression), -//! ready to be folded directly into partition `k`'s rows. The scheduler -//! bakes this in when it constructs the downstream stage after the upstream +//! `try_new` takes a `Vec<`[`WindowApply`]`>` — one entry per output column +//! that needs cross-partition correction, telling the operator *how* to +//! rewrite that column. Two shapes: +//! +//! - [`WindowApply::Scalar`] — fast path. Combines each row's existing value +//! with a scheduler-provided scalar via [`ScalarOp`] (`Add`/`Min`/`Max` for +//! SUM/COUNT/MIN/MAX and `row_number`; `Overwrite` for `first_value` / +//! `last_value`). No `Accumulator` constructed. +//! - [`WindowApply::Aggregate`] — fallback. Constructs a fresh `Accumulator` +//! seeded from the pre-merged state, feeds `args` per row, overwrites the +//! column with `evaluate()`. Fits AVG (without decomposition), sketch-backed +//! windows (APPROX_DISTINCT, APPROX_QUANTILE), and statistical aggregates. +//! +//! Non-corrected window functions don't appear in `applies` at all. `lead` / +//! `lag` / `nth_value` are solved by halo rows in the shuffle layer. +//! `rank` / `dense_rank` / `percent_rank` / `cume_dist` / `ntile` need a +//! separate segment-tree-plus-broadcast design and are out of scope here. +//! +//! # Prefix-state input +//! +//! [`FinalizedPartitionState`] — one map per input partition, holding +//! *pre-merged* [`Accumulator::state`] for every (PARTITION BY key, window +//! expression) pair. Only consumed by [`WindowApply::Aggregate`] entries; +//! [`WindowApply::Scalar`] carries its own offsets inline. The scheduler +//! bakes both when it constructs the downstream stage after the upstream //! stage's tasks complete. //! //! **This is scaffolding.** `execute` currently forwards batches unchanged. @@ -80,8 +94,9 @@ use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::{Result, ScalarValue, Statistics, internal_err}; use datafusion::execution::TaskContext; +use datafusion::logical_expr::AggregateUDF; use datafusion::physical_expr::window::PartitionKey; -use datafusion::physical_expr::{Distribution, OrderingRequirements}; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, @@ -107,6 +122,97 @@ use std::collections::HashMap; /// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 pub type FinalizedPartitionState = HashMap>>>; +/// How to combine each row's existing value in an output column with a +/// scheduler-provided scalar offset. The result overwrites the column. +/// +/// [`Overwrite`] ignores the row's existing value and just writes the offset; +/// it's the shape needed for `first_value` / `last_value`, where the scheduler +/// picks the correct global value once and every row gets a copy. +/// +/// [`Overwrite`]: ScalarOp::Overwrite +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ScalarOp { + /// `output := row_value + offset`. Fits SUM, COUNT, and ranking functions + /// like `row_number` (which are effectively `COUNT(*)`). + Add, + /// `output := min(row_value, offset)`. Fits MIN. + Min, + /// `output := max(row_value, offset)`. Fits MAX. + Max, + /// `output := offset`. Ignores `row_value`. Fits `first_value` / + /// `last_value`, where the scheduler picks a single global scalar and + /// every row gets the same corrected value. + Overwrite, +} + +/// How to correct one window-function output column at row-apply time. Each +/// entry describes exactly one column PrefixMergeExec should rewrite. +/// +/// Two shapes are covered by construction; anything else is out of scope +/// (`lead`/`lag`/`nth_value` are solved by halo rows in the shuffle layer, and +/// ranking-family functions like `rank`/`percent_rank`/`ntile` want a separate +/// segment-tree-plus-broadcast infrastructure). +#[derive(Debug, Clone)] +pub enum WindowApply { + /// Fast path: monoidal op between each row's existing value and a + /// scheduler-provided scalar. No `Accumulator` is constructed. + /// + /// Fits every aggregate whose per-row output is itself a valid partial + /// state that composes via a single scalar op — SUM, COUNT, MIN, MAX — + /// plus ranking functions like `row_number` (offset = prior row count) + /// and value-selection functions like `first_value` / `last_value` + /// (offset = scheduler-picked global value; op = [`ScalarOp::Overwrite`]). + Scalar { + /// Combining op between `row_value` and `offset`. + op: ScalarOp, + /// One scalar per input partition. `offset[k]` combines with every + /// row passing through partition `k`. Length must match the input's + /// output partition count. + offset: Vec, + /// Column overwritten with `op(row_value, offset[partition])`. + output_column: usize, + }, + /// Fallback path: fresh `Accumulator` per partition, seeded with the + /// merged offset state via [`Accumulator::merge_batch`], updated per row + /// with `args` evaluated against that row, then [`Accumulator::evaluate`] + /// overwrites `output_column`. + /// + /// Fits aggregates whose per-row output isn't a valid partial state: + /// AVG (without decomposition), sketch-backed windows like + /// APPROX_DISTINCT and APPROX_QUANTILE, and statistical aggregates like + /// STDDEV / VAR / correlation whose state is a tuple of running moments. + /// + /// [`Accumulator::merge_batch`]: datafusion::logical_expr::Accumulator::merge_batch + /// [`Accumulator::evaluate`]: datafusion::logical_expr::Accumulator::evaluate + Aggregate { + /// UDF used to construct a fresh `Accumulator` per partition. + udf: Arc, + /// Aggregate's argument expressions, evaluated against each input row + /// and fed to `Accumulator::update_batch`. For SUM/COUNT/MIN/MAX where + /// re-running the accumulator is redundant, prefer the [`Scalar`] + /// variant; this path is for cases where re-running is required. + /// + /// [`Scalar`]: WindowApply::Scalar + args: Vec>, + /// Column overwritten with the accumulator's `evaluate()` result. + output_column: usize, + /// Position in the upstream `BoundedWindowAggExec`'s `window_expr()` + /// list — the index into the inner `Vec` inside + /// [`FinalizedPartitionState`] where this aggregate's merged offset + /// state lives. + window_expr_index: usize, + }, +} + +impl WindowApply { + fn output_column(&self) -> usize { + match self { + WindowApply::Scalar { output_column, .. } + | WindowApply::Aggregate { output_column, .. } => *output_column, + } + } +} + /// Apply pre-merged window-aggregate state (computed by the scheduler) to /// each row of the current partition's output. See the [module-level /// docs][self] for the division of labor between scheduler and executor and @@ -116,22 +222,32 @@ pub type FinalizedPartitionState = HashMap, + /// One entry per window-function output column that needs cross-partition + /// correction. Non-corrected columns (e.g. `lead`/`lag` handled by halos, + /// or ranking functions left to segment-tree infrastructure) don't appear + /// here. + applies: Vec, /// `per_partition_state[k]` is the scheduler-provided *already-merged* - /// state summarising every input partition in `[0..k)`. This operator - /// folds it into partition `k`'s rows; it does not perform the merge - /// itself (the scheduler already did). Length equals + /// state summarising every input partition in `[0..k)`. Only consumed by + /// [`WindowApply::Aggregate`] entries — [`WindowApply::Scalar`] carries + /// its own offsets. Length equals /// `input.output_partitioning().partition_count()`. per_partition_state: Vec, properties: Arc, } impl PrefixMergeExec { - /// Wrap `input` with per-input-partition prefix state. + /// Wrap `input` with per-column apply descriptors and per-input-partition + /// prefix state. /// - /// Errors if `per_partition_state.len()` doesn't match the input's - /// partition count. + /// Errors on any of: + /// - `per_partition_state.len()` != input's partition count. + /// - Any [`WindowApply::Scalar`]'s `offset.len()` != input's partition + /// count. + /// - Any entry's `output_column` outside the input schema's field range. pub fn try_new( input: Arc, + applies: Vec, per_partition_state: Vec, ) -> Result { let partition_count = input.output_partitioning().partition_count(); @@ -143,6 +259,26 @@ impl PrefixMergeExec { partition_count ); } + let field_count = input.schema().fields().len(); + for (i, apply) in applies.iter().enumerate() { + let col = apply.output_column(); + if col >= field_count { + return internal_err!( + "PrefixMergeExec: applies[{i}] output_column {col} out of \ + range (schema has {field_count} fields)" + ); + } + if let WindowApply::Scalar { offset, .. } = apply + && offset.len() != partition_count + { + return internal_err!( + "PrefixMergeExec: applies[{i}] Scalar offset.len() {} does \ + not match input partition count {}", + offset.len(), + partition_count + ); + } + } let properties = Arc::new(PlanProperties::new( input.equivalence_properties().clone(), input.output_partitioning().clone(), @@ -151,13 +287,20 @@ impl PrefixMergeExec { )); Ok(Self { input, + applies, per_partition_state, properties, }) } - /// The prefix state carried per input partition. - /// `per_partition_state()[k]` is what will be merged into partition `k`. + /// Per-column apply descriptors. `applies()[i]` corresponds to one + /// output column that will be rewritten by the prefix-merge. + pub fn applies(&self) -> &[WindowApply] { + &self.applies + } + + /// The prefix state carried per input partition. Only consumed by + /// [`WindowApply::Aggregate`] entries. pub fn per_partition_state(&self) -> &[FinalizedPartitionState] { &self.per_partition_state } @@ -167,6 +310,7 @@ impl Debug for PrefixMergeExec { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("PrefixMergeExec") .field("partition_count", &self.per_partition_state.len()) + .field("applies", &self.applies.len()) .finish() } } @@ -177,8 +321,9 @@ impl DisplayAs for PrefixMergeExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!( f, - "PrefixMergeExec: partitions={}", - self.per_partition_state.len() + "PrefixMergeExec: partitions={}, applies={}", + self.per_partition_state.len(), + self.applies.len() ) } DisplayFormatType::TreeRender => { @@ -217,6 +362,7 @@ impl ExecutionPlan for PrefixMergeExec { }; Ok(Arc::new(PrefixMergeExec::try_new( input.clone(), + self.applies.clone(), self.per_partition_state.clone(), )?)) } @@ -302,13 +448,20 @@ mod tests { Arc::new(DataSourceExec::new(Arc::new(source))) } + fn empty_state(partitions: usize) -> Vec { + (0..partitions) + .map(|_| FinalizedPartitionState::new()) + .collect() + } + /// Mismatched state length surfaces as an error rather than a panic at /// runtime. #[test] fn try_new_rejects_state_length_mismatch() { let input = partitioned_source(2, 3); - let err = PrefixMergeExec::try_new(input, vec![FinalizedPartitionState::new()]) - .expect_err("length mismatch must surface as an error"); + let err = + PrefixMergeExec::try_new(input, vec![], vec![FinalizedPartitionState::new()]) + .expect_err("length mismatch must surface as an error"); assert!( err.to_string() .contains("does not match input partition count"), @@ -316,16 +469,57 @@ mod tests { ); } - /// While the merge is unimplemented, `execute` behaves as a passthrough: - /// each partition's rows are forwarded verbatim. + /// A [`WindowApply::Scalar`] whose `offset.len()` doesn't match the + /// input partition count is caught at construction, before any rows flow. + #[test] + fn try_new_rejects_scalar_offset_length_mismatch() { + let input = partitioned_source(2, 3); + let apply = WindowApply::Scalar { + op: ScalarOp::Add, + offset: vec![ScalarValue::Int64(Some(0))], // only 1, need 2 + output_column: 0, + }; + let err = PrefixMergeExec::try_new(input, vec![apply], empty_state(2)) + .expect_err("scalar offset length mismatch must surface as an error"); + assert!( + err.to_string().contains("Scalar offset.len()"), + "unexpected error: {err}" + ); + } + + /// An `output_column` past the input schema's field count errors. + #[test] + fn try_new_rejects_output_column_out_of_range() { + let input = partitioned_source(2, 3); + let apply = WindowApply::Scalar { + op: ScalarOp::Add, + offset: vec![ScalarValue::Int64(Some(0)); 2], + output_column: 5, // schema has 1 field + }; + let err = PrefixMergeExec::try_new(input, vec![apply], empty_state(2)) + .expect_err("out-of-range output_column must surface as an error"); + assert!( + err.to_string().contains("output_column 5 out of range"), + "unexpected error: {err}" + ); + } + + /// While the merge is unimplemented, `execute` behaves as a passthrough + /// even with `applies` populated: each partition's rows are forwarded + /// verbatim regardless of what the descriptors say. #[tokio::test] async fn scaffold_forwards_input_rows_unchanged() -> Result<()> { let input = partitioned_source(2, 3); - let state: Vec = vec![ - FinalizedPartitionState::new(), - FinalizedPartitionState::new(), - ]; - let exec = Arc::new(PrefixMergeExec::try_new(input, state)?); + let apply = WindowApply::Scalar { + op: ScalarOp::Add, + offset: vec![ScalarValue::Int64(Some(100)); 2], + output_column: 0, + }; + let exec = Arc::new(PrefixMergeExec::try_new( + input, + vec![apply], + empty_state(2), + )?); let ctx = SessionContext::new().task_ctx(); for k in 0..2 { @@ -338,7 +532,11 @@ mod tests { .downcast_ref::() .expect("Int64 column"); let expected: Vec = ((k * 3) as i64..((k + 1) * 3) as i64).collect(); - assert_eq!(arr.values(), &expected); + assert_eq!( + arr.values(), + &expected, + "scaffold must not apply the offset yet" + ); } Ok(()) } From 115e511527e1531bc123aae09eca7a7c85ff545a Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 11:45:43 -0600 Subject: [PATCH 4/7] =?UTF-8?q?feat(prefix-merge):=20implement=20WindowApp?= =?UTF-8?q?ly::Aggregate=20=E2=80=94=20real=20per-row=20seeded=20merge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit execute() now actually applies WindowApply::Aggregate entries: for each partition it builds a fresh Accumulator via AggregateExprBuilder, seeds it via merge_batch with the pre-merged offset state from per_partition_state, then per row evaluates the aggregate's args, feeds them to update_batch, and overwrites output_column with evaluate(). Assumes at most one PARTITION BY key per input partition — the AQE synthetic-PARTITION-BY case. If a state map carries more than one key, execute() errors rather than silently applying only one; multi-key support can grow when a workload needs it. The WindowApply::Scalar path is still a passthrough for now, flagged in the doc so it doesn't come as a surprise; batch-arithmetic apply via arrow kernels lands in a follow-up. New test approx_distinct_corrects_running_distinct_across_partitions runs the sketch case end-to-end: two mock BWAG outputs with local running distinct counts, partition 1's offset seeded from partition 0's terminal HLL state via approx_distinct_udaf's own Accumulator::state. Verifies partition 1's corrected running distinct matches what a single BWAG over the concatenated input would have produced (4, 5, 6 instead of the local 1, 2, 3). This is the "worst case" — an aggregate whose state can't be recovered from a running scalar output, so per-row scalar correction via ProjectionExec is provably insufficient and only the seeded-accumulator re-run gets the right answer. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/prefix_merge.rs | 349 +++++++++++++++++- 1 file changed, 332 insertions(+), 17 deletions(-) diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index 859ccfc6c..c90a29fb1 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -71,10 +71,16 @@ //! bakes both when it constructs the downstream stage after the upstream //! stage's tasks complete. //! -//! **This is scaffolding.** `execute` currently forwards batches unchanged. -//! The row-wise fold lands in a follow-up. The shape is committed to now so -//! the surrounding plumbing (scheduler collection, per-task state injection, -//! plan-rule placement) can be built against a stable signature. +//! **Status.** The [`WindowApply::Aggregate`] path is implemented: for each +//! partition it builds a fresh `Accumulator`, seeds it via `merge_batch` from +//! the offset state, and replays each row through `update_batch` + `evaluate` +//! to overwrite the output column. Assumes at most one PARTITION BY key per +//! input partition (matches the AQE synthetic-PARTITION-BY pattern); +//! multi-key handling can grow when a workload needs it. +//! +//! The [`WindowApply::Scalar`] path is still a passthrough — those output +//! columns are left untouched. Applying the scalar op via arrow kernels +//! lands in a follow-up. //! //! # Relation to DataFusion //! @@ -91,18 +97,23 @@ use std::fmt::{self, Debug, Formatter}; use std::sync::Arc; +use datafusion::arrow::array::{ArrayRef, RecordBatch}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::{Result, ScalarValue, Statistics, internal_err}; use datafusion::execution::TaskContext; -use datafusion::logical_expr::AggregateUDF; +use datafusion::logical_expr::{Accumulator, AggregateUDF}; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; use datafusion::physical_expr::window::PartitionKey; use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - SendableRecordBatchStream, + ColumnarValue, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; +use futures::{Stream, StreamExt, ready}; use std::collections::HashMap; +use std::pin::Pin; +use std::task::{Context, Poll}; /// A single already-prefix-merged window-aggregate state, keyed by PARTITION /// BY tuple. The inner `Vec` is indexed by window expression (same order the @@ -218,8 +229,9 @@ impl WindowApply { /// docs][self] for the division of labor between scheduler and executor and /// the AQE pipeline this fits into. /// -/// **Scaffolding.** `execute` forwards batches unchanged; the row-wise fold -/// lands in a follow-up. +/// [`WindowApply::Aggregate`] entries are applied per row via a seeded +/// `Accumulator`; [`WindowApply::Scalar`] entries are currently passthrough +/// (their apply lands in a follow-up). pub struct PrefixMergeExec { input: Arc, /// One entry per window-function output column that needs cross-partition @@ -398,10 +410,6 @@ impl ExecutionPlan for PrefixMergeExec { CardinalityEffect::Equal } - /// **Scaffolding.** Forwards the input stream unchanged. The row-wise - /// fold — applying `self.per_partition_state[partition]` (already merged - /// upstream by the scheduler) to each row's window-aggregate columns — - /// lands in a follow-up. fn execute( &self, partition: usize, @@ -414,7 +422,176 @@ impl ExecutionPlan for PrefixMergeExec { self.per_partition_state.len() ); } - self.input.execute(partition, ctx) + let input_schema = self.input.schema(); + let output_schema = self.schema(); + + // Pick this partition's pre-merged state. For the AQE synthetic- + // PARTITION-BY case each task carries exactly one key; multi-key + // handling can grow when a real workload needs it. + let state = &self.per_partition_state[partition]; + let key_state: Option<&Vec>>> = match state.len() { + 0 => None, + 1 => state.values().next(), + n => { + return internal_err!( + "PrefixMergeExec: {n} PARTITION BY keys in the state map \ + for partition {partition}; multi-key apply is not yet \ + implemented" + ); + } + }; + + let mut aggregate_appliers: Vec = Vec::new(); + for (i, apply) in self.applies.iter().enumerate() { + match apply { + WindowApply::Aggregate { + udf, + args, + output_column, + window_expr_index, + } => { + let offset_state: Option<&Vec> = key_state + .and_then(|per_expr| per_expr.get(*window_expr_index)) + .and_then(|slot| slot.as_ref()); + aggregate_appliers.push(AggregateApply::new( + i, + udf, + args, + *output_column, + &input_schema, + offset_state, + )?); + } + WindowApply::Scalar { .. } => { + // Scalar variant application lands in a follow-up. Its + // output column is left untouched for now (passthrough). + } + } + } + + let input = self.input.execute(partition, ctx)?; + let stream = ApplyStream { + input, + appliers: aggregate_appliers, + schema: Arc::clone(&output_schema), + }; + Ok(Box::pin(stream)) + } +} + +/// A single [`WindowApply::Aggregate`] prepared for a specific input +/// partition: the Accumulator has already been seeded from the offset state. +struct AggregateApply { + /// Position of the source `WindowApply` in the exec's `applies` list — + /// carried through so error messages can point at the offender. + apply_index: usize, + accumulator: Box, + args: Vec>, + output_column: usize, +} + +impl AggregateApply { + fn new( + apply_index: usize, + udf: &Arc, + args: &[Arc], + output_column: usize, + input_schema: &SchemaRef, + offset_state: Option<&Vec>, + ) -> Result { + let agg_expr = AggregateExprBuilder::new(Arc::clone(udf), args.to_vec()) + .schema(Arc::clone(input_schema)) + .alias(format!("prefix_merge_apply_{apply_index}")) + .build()?; + let mut accumulator = agg_expr.create_accumulator()?; + if let Some(state_scalars) = offset_state { + let offset_arrays: Vec = state_scalars + .iter() + .map(|s| s.to_array_of_size(1)) + .collect::>>()?; + accumulator.merge_batch(&offset_arrays)?; + } + Ok(Self { + apply_index, + accumulator, + args: args.to_vec(), + output_column, + }) + } + + /// Evaluate `args` against `batch`, replay them through `accumulator` + /// row by row, and overwrite `output_column` with the accumulator's + /// per-row `evaluate()` result. + fn apply(&mut self, batch: RecordBatch) -> Result { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(batch); + } + let arg_arrays: Vec = self + .args + .iter() + .map(|expr| match expr.evaluate(&batch)? { + ColumnarValue::Array(a) => Ok(a), + ColumnarValue::Scalar(s) => s.to_array_of_size(num_rows), + }) + .collect::>>()?; + + let mut new_values: Vec = Vec::with_capacity(num_rows); + for i in 0..num_rows { + let row_args: Vec = + arg_arrays.iter().map(|a| a.slice(i, 1)).collect(); + self.accumulator.update_batch(&row_args)?; + new_values.push(self.accumulator.evaluate()?); + } + let new_column = ScalarValue::iter_to_array(new_values)?; + + if self.output_column >= batch.num_columns() { + return internal_err!( + "PrefixMergeExec: applies[{}] output_column {} out of range \ + at execute time (batch has {} columns)", + self.apply_index, + self.output_column, + batch.num_columns() + ); + } + let mut columns = batch.columns().to_vec(); + columns[self.output_column] = new_column; + Ok(RecordBatch::try_new(batch.schema(), columns)?) + } +} + +struct ApplyStream { + input: SendableRecordBatchStream, + appliers: Vec, + schema: SchemaRef, +} + +impl Stream for ApplyStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(mut batch)) => { + for applier in &mut self.appliers { + batch = match applier.apply(batch) { + Ok(b) => b, + Err(e) => return Poll::Ready(Some(Err(e))), + }; + } + Poll::Ready(Some(Ok(batch))) + } + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => Poll::Ready(None), + } + } +} + +impl RecordBatchStream for ApplyStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) } } @@ -504,9 +681,147 @@ mod tests { ); } - /// While the merge is unimplemented, `execute` behaves as a passthrough - /// even with `applies` populated: each partition's rows are forwarded - /// verbatim regardless of what the descriptors say. + /// Runs the accumulator to completion over `values`, then returns the + /// state — the shape a real upstream `BoundedWindowAggExec` reports for + /// this window aggregate at partition close. + fn approx_distinct_state( + udf: &Arc, + input_schema: &SchemaRef, + values: &[i64], + ) -> Result> { + let column: Arc = + Arc::new(datafusion::physical_expr::expressions::Column::new("v", 0)); + let agg_expr = AggregateExprBuilder::new(Arc::clone(udf), vec![column]) + .schema(Arc::clone(input_schema)) + .alias("state_helper") + .build()?; + let mut acc = agg_expr.create_accumulator()?; + let arr: ArrayRef = + Arc::new(Int64Array::from_iter_values(values.iter().copied())); + acc.update_batch(&[arr])?; + acc.state() + } + + /// End-to-end demonstration on the sketch case that motivates this whole + /// design: cumulative APPROX_DISTINCT across a range-shuffled ordered + /// stream, corrected per row by re-running the HLL accumulator seeded + /// with the pre-merged state of every prior partition. + /// + /// Setup: + /// - Partition 0's mock upstream output: `[(1,1), (2,2), (3,3)]` — v and + /// the local running distinct count. + /// - Partition 1's mock upstream output: `[(4,1), (5,2), (6,3)]` — again + /// local, so wrong globally: the running distinct at the last row + /// should be 6, not 3. + /// - The scheduler collects partition 0's terminal Accumulator state + /// (HLL registers seeded with {1,2,3}) and hands it to partition 1 as + /// its offset. Partition 0 gets an empty offset (nothing before it). + /// + /// Expected corrected output: `[(1,1),(2,2),(3,3)]` on partition 0 and + /// `[(4,4),(5,5),(6,6)]` on partition 1 — matches what a single BWAG on + /// the concatenated `[1..6]` would produce. + /// + /// This is the case that a two-pass halo scheme can't handle without + /// emitting HLL registers on every output row: recovering state from + /// a running distinct-count scalar is not tractable. The side-channel + /// design routes state around the data stream and applies it per row + /// here. + #[tokio::test] + async fn approx_distinct_corrects_running_distinct_across_partitions() -> Result<()> { + use datafusion::functions_aggregate::approx_distinct::approx_distinct_udaf; + use datafusion::physical_expr::expressions::Column; + + // BWAG's output schema: original argument column + the running + // approx-distinct column that the aggregate produced. + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Int64, false), + Field::new("running_distinct", DataType::UInt64, false), + ])); + let batch = |vs: &[i64], rs: &[u64]| -> RecordBatch { + let v_col = Arc::new(Int64Array::from_iter_values(vs.iter().copied())); + let r_col = + Arc::new(datafusion::arrow::array::UInt64Array::from_iter_values( + rs.iter().copied(), + )); + RecordBatch::try_new(schema.clone(), vec![v_col, r_col]).unwrap() + }; + // Local (wrong-globally) BWAG output per partition. + let p0 = batch(&[1, 2, 3], &[1, 2, 3]); + let p1 = batch(&[4, 5, 6], &[1, 2, 3]); + let source = + MemorySourceConfig::try_new(&[vec![p0], vec![p1]], schema.clone(), None)?; + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + + let udf = approx_distinct_udaf(); + // Partition 1's offset = the HLL state after ingesting {1,2,3}. + let p0_state = approx_distinct_state( + &udf, + &Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])), + &[1, 2, 3], + )?; + // Assemble the per-partition state maps. PARTITION BY tuple is empty + // (global window); one aggregate at window_expr_index 0. + let empty_key: PartitionKey = vec![]; + let mut state_p1 = FinalizedPartitionState::new(); + state_p1.insert(empty_key.clone(), vec![Some(p0_state)]); + let per_partition_state = vec![FinalizedPartitionState::new(), state_p1]; + + let apply = WindowApply::Aggregate { + udf: Arc::clone(&udf), + // Feed the accumulator with the original argument column `v` + // from the batch (not the already-computed running_distinct). + args: vec![Arc::new(Column::new("v", 0))], + output_column: 1, + window_expr_index: 0, + }; + let exec = Arc::new(PrefixMergeExec::try_new( + input, + vec![apply], + per_partition_state, + )?); + + let ctx = SessionContext::new().task_ctx(); + + // Partition 0: offset is empty, running_distinct output unchanged. + let out_p0: Vec = + exec.execute(0, ctx.clone())?.try_collect().await?; + let p0_running = out_p0[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p0_running, + vec![1, 2, 3], + "partition 0 with empty offset should preserve local running distinct" + ); + + // Partition 1: seeded with partition 0's HLL, running_distinct + // corrected to the global cumulative count. + let out_p1: Vec = + exec.execute(1, ctx.clone())?.try_collect().await?; + let p1_running = out_p1[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p1_running, + vec![4, 5, 6], + "partition 1 seeded with partition-0 state should show \ + global-cumulative running distinct" + ); + Ok(()) + } + + /// The Scalar apply variant is still passthrough — its output column is + /// left untouched by execute. This test guards against silent behavior + /// change until the Scalar path is implemented. #[tokio::test] async fn scaffold_forwards_input_rows_unchanged() -> Result<()> { let input = partitioned_source(2, 3); From 201ebea611f90755d814c2500af6a6ad73344a0a Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 11:51:03 -0600 Subject: [PATCH 5/7] test(prefix-merge): add AVG cross-partition demo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the APPROX_DISTINCT test on the aggregate whose state has been the reference point for "why can't we just add offsets": AVG stores (sum, count) but emits sum/count per row, so recovering (sum, count) from a running mean is impossible. A ProjectionExec that added a scalar offset to running_avg would give the wrong answer at every row. The seeded-Accumulator re-run handles it uniformly — same code path as APPROX_DISTINCT, no decomposition needed. Partition 1's mock local running means `[40.0, 45.0, 50.0]` become the correct global running means `[25.0, 30.0, 35.0]` after seeding with partition 0's terminal AVG state (sum=60, count=3). Uses Float64 throughout because DataFusion's AvgAccumulator is only implemented for Float64/decimal/duration inputs; Int64 would rely on a planner-inserted cast that isn't needed to demonstrate the mechanism. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/prefix_merge.rs | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index c90a29fb1..ddb15f058 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -819,6 +819,127 @@ mod tests { Ok(()) } + /// AVG's per-row output (`sum / count`) also isn't a valid partial state: + /// you can't recover `(sum, count)` from a single mean, so a naive + /// `ProjectionExec` adding a scalar offset would give the wrong answer. + /// The seeded-`Accumulator` re-run handles it uniformly — same code path + /// that worked for APPROX_DISTINCT. + /// + /// Setup mirrors the APPROX_DISTINCT test: + /// - Partition 0's mock upstream output: `[(10, 10.0), (20, 15.0), + /// (30, 20.0)]` — v and the local running mean. + /// - Partition 1's mock upstream output: `[(40, 40.0), (50, 45.0), + /// (60, 50.0)]` — again local, so wrong globally. + /// - Partition 1's offset = AVG's terminal state after ingesting + /// `{10, 20, 30}` = `(sum=60, count=3)`. + /// + /// Expected corrected partition-1 running mean: + /// - `(60+40)/(3+1) = 25.0` + /// - `(100+50)/(4+1) = 30.0` + /// - `(150+60)/(5+1) = 35.0` + /// + /// Which is what a single BWAG on the concatenated `[10..60 step 10]` + /// would produce. + #[tokio::test] + async fn avg_corrects_running_mean_across_partitions() -> Result<()> { + use datafusion::arrow::array::Float64Array; + use datafusion::functions_aggregate::average::avg_udaf; + use datafusion::physical_expr::expressions::Column; + + // DataFusion's `AvgAccumulator` is only wired up for Float64 (and + // decimal / duration) inputs — the planner casts Int64 to Float64 + // before AVG in a real query. Use Float64 directly so the test + // hits the standard accumulator without stitching a cast in. + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Float64, false), + Field::new("running_avg", DataType::Float64, false), + ])); + let batch = |vs: &[f64], means: &[f64]| -> RecordBatch { + let v_col = Arc::new(Float64Array::from_iter_values(vs.iter().copied())); + let m_col = Arc::new(Float64Array::from_iter_values(means.iter().copied())); + RecordBatch::try_new(schema.clone(), vec![v_col, m_col]).unwrap() + }; + // Local (wrong-globally) BWAG output per partition. + let p0 = batch(&[10.0, 20.0, 30.0], &[10.0, 15.0, 20.0]); + let p1 = batch(&[40.0, 50.0, 60.0], &[40.0, 45.0, 50.0]); + let source = + MemorySourceConfig::try_new(&[vec![p0], vec![p1]], schema.clone(), None)?; + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + + let udf = avg_udaf(); + // Partition 1's offset = AVG's terminal state after ingesting + // partition 0's values. Constructed inline (not through the + // Int64-only helper the APPROX_DISTINCT test uses) because AVG needs + // Float64 args. + let single_col_schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let p0_state = { + let column: Arc = Arc::new(Column::new("v", 0)); + let agg_expr = AggregateExprBuilder::new(Arc::clone(&udf), vec![column]) + .schema(single_col_schema) + .alias("avg_state_helper") + .build()?; + let mut acc = agg_expr.create_accumulator()?; + let arr: ArrayRef = + Arc::new(Float64Array::from_iter_values([10.0, 20.0, 30.0])); + acc.update_batch(&[arr])?; + acc.state()? + }; + let empty_key: PartitionKey = vec![]; + let mut state_p1 = FinalizedPartitionState::new(); + state_p1.insert(empty_key.clone(), vec![Some(p0_state)]); + let per_partition_state = vec![FinalizedPartitionState::new(), state_p1]; + + let apply = WindowApply::Aggregate { + udf: Arc::clone(&udf), + args: vec![Arc::new(Column::new("v", 0))], + output_column: 1, + window_expr_index: 0, + }; + let exec = Arc::new(PrefixMergeExec::try_new( + input, + vec![apply], + per_partition_state, + )?); + + let ctx = SessionContext::new().task_ctx(); + + // Partition 0: no offset, output preserved. + let out_p0: Vec = + exec.execute(0, ctx.clone())?.try_collect().await?; + let p0_avg = out_p0[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p0_avg, + vec![10.0, 15.0, 20.0], + "partition 0 with empty offset should preserve local running mean" + ); + + // Partition 1: seeded with (sum=60, count=3), corrected per row. + let out_p1: Vec = + exec.execute(1, ctx.clone())?.try_collect().await?; + let p1_avg = out_p1[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p1_avg, + vec![25.0, 30.0, 35.0], + "partition 1 seeded with partition-0 AVG state should show \ + global-cumulative running mean" + ); + Ok(()) + } + /// The Scalar apply variant is still passthrough — its output column is /// left untouched by execute. This test guards against silent behavior /// change until the Scalar path is implemented. From 1ab1f2650448782eb5db60d5d5c64d9fb1610822 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 11:58:13 -0600 Subject: [PATCH 6/7] feat(prefix-merge): implement WindowApply::Scalar via arrow kernels + SUM demo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the Scalar fast path through arrow kernels: - ScalarOp::Add: arrow::compute::kernels::numeric::add - ScalarOp::Min: cmp::lt_eq + zip::zip (element-wise min) - ScalarOp::Max: cmp::gt_eq + zip::zip (element-wise max) - ScalarOp::Overwrite: constant-fill from the offset The stream now dispatches through a PreparedApply enum wrapping either a ScalarApply (batch-arithmetic, no Accumulator) or an AggregateApply (seeded Accumulator, per-row replay). One code path holds them both, each applied per batch in order. Replaces the scaffold-passthrough test with sum_corrects_running_sum_ across_partitions — the canonical fast-path demo. Partition 1's mock local running sums [4, 9, 15] become the correct global [10, 15, 21] after adding the scheduler-provided offset 6. Rounds out the trio: SUM via Scalar::Add, AVG via Aggregate re-run, APPROX_DISTINCT via Aggregate re-run. Same descriptor list; the operator picks the shape each entry needs. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/prefix_merge.rs | 203 ++++++++++++++---- 1 file changed, 161 insertions(+), 42 deletions(-) diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index ddb15f058..b29615091 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -71,16 +71,18 @@ //! bakes both when it constructs the downstream stage after the upstream //! stage's tasks complete. //! -//! **Status.** The [`WindowApply::Aggregate`] path is implemented: for each -//! partition it builds a fresh `Accumulator`, seeds it via `merge_batch` from -//! the offset state, and replays each row through `update_batch` + `evaluate` -//! to overwrite the output column. Assumes at most one PARTITION BY key per -//! input partition (matches the AQE synthetic-PARTITION-BY pattern); -//! multi-key handling can grow when a workload needs it. +//! **Status.** Both apply paths are implemented. //! -//! The [`WindowApply::Scalar`] path is still a passthrough — those output -//! columns are left untouched. Applying the scalar op via arrow kernels -//! lands in a follow-up. +//! - [`WindowApply::Aggregate`] builds a fresh `Accumulator` per partition, +//! seeds it via `merge_batch` from the offset state, and replays each row +//! through `update_batch` + `evaluate` to overwrite the output column. +//! - [`WindowApply::Scalar`] applies the [`ScalarOp`] batch-at-a-time via +//! arrow kernels — `numeric::add` for `Add`, `cmp::lt_eq`/`gt_eq` + `zip` +//! for `Min`/`Max`, and a constant-fill for `Overwrite`. +//! +//! Assumes at most one PARTITION BY key per input partition (matches the +//! AQE synthetic-PARTITION-BY pattern); multi-key handling can grow when a +//! workload needs it. //! //! # Relation to DataFusion //! @@ -230,8 +232,8 @@ impl WindowApply { /// the AQE pipeline this fits into. /// /// [`WindowApply::Aggregate`] entries are applied per row via a seeded -/// `Accumulator`; [`WindowApply::Scalar`] entries are currently passthrough -/// (their apply lands in a follow-up). +/// `Accumulator`; [`WindowApply::Scalar`] entries are applied per batch via +/// arrow kernels. pub struct PrefixMergeExec { input: Arc, /// One entry per window-function output column that needs cross-partition @@ -441,7 +443,7 @@ impl ExecutionPlan for PrefixMergeExec { } }; - let mut aggregate_appliers: Vec = Vec::new(); + let mut appliers: Vec = Vec::with_capacity(self.applies.len()); for (i, apply) in self.applies.iter().enumerate() { match apply { WindowApply::Aggregate { @@ -453,18 +455,26 @@ impl ExecutionPlan for PrefixMergeExec { let offset_state: Option<&Vec> = key_state .and_then(|per_expr| per_expr.get(*window_expr_index)) .and_then(|slot| slot.as_ref()); - aggregate_appliers.push(AggregateApply::new( + appliers.push(PreparedApply::Aggregate(AggregateApply::new( i, udf, args, *output_column, &input_schema, offset_state, - )?); + )?)); } - WindowApply::Scalar { .. } => { - // Scalar variant application lands in a follow-up. Its - // output column is left untouched for now (passthrough). + WindowApply::Scalar { + op, + offset, + output_column, + } => { + appliers.push(PreparedApply::Scalar(ScalarApply { + apply_index: i, + op: *op, + offset: offset[partition].clone(), + output_column: *output_column, + })); } } } @@ -472,13 +482,80 @@ impl ExecutionPlan for PrefixMergeExec { let input = self.input.execute(partition, ctx)?; let stream = ApplyStream { input, - appliers: aggregate_appliers, + appliers, schema: Arc::clone(&output_schema), }; Ok(Box::pin(stream)) } } +/// One [`WindowApply`] entry prepared for a specific input partition: the +/// scheduler's offset has been resolved down to a single value (or a seeded +/// `Accumulator`) that can be applied per batch. +enum PreparedApply { + Scalar(ScalarApply), + Aggregate(AggregateApply), +} + +impl PreparedApply { + fn apply(&mut self, batch: RecordBatch) -> Result { + match self { + PreparedApply::Scalar(s) => s.apply(batch), + PreparedApply::Aggregate(a) => a.apply(batch), + } + } +} + +/// A single [`WindowApply::Scalar`] prepared for one input partition: the +/// scheduler's per-partition offset has been narrowed down to a single +/// [`ScalarValue`] and the combining `op` is applied batch-at-a-time via +/// arrow kernels. +struct ScalarApply { + apply_index: usize, + op: ScalarOp, + offset: ScalarValue, + output_column: usize, +} + +impl ScalarApply { + fn apply(&self, batch: RecordBatch) -> Result { + use datafusion::arrow::compute::kernels::{cmp, numeric, zip}; + + if self.output_column >= batch.num_columns() { + return internal_err!( + "PrefixMergeExec: applies[{}] output_column {} out of range \ + at execute time (batch has {} columns)", + self.apply_index, + self.output_column, + batch.num_columns() + ); + } + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(batch); + } + let col = batch.column(self.output_column); + let offset_arr: ArrayRef = self.offset.to_array_of_size(num_rows)?; + let new_col: ArrayRef = match self.op { + ScalarOp::Add => numeric::add(col, &offset_arr)?, + ScalarOp::Min => { + // element-wise min: keep `col` where col ≤ offset, else offset + let mask = cmp::lt_eq(col, &offset_arr)?; + zip::zip(&mask, col, &offset_arr)? + } + ScalarOp::Max => { + // element-wise max: keep `col` where col ≥ offset, else offset + let mask = cmp::gt_eq(col, &offset_arr)?; + zip::zip(&mask, col, &offset_arr)? + } + ScalarOp::Overwrite => offset_arr, + }; + let mut columns = batch.columns().to_vec(); + columns[self.output_column] = new_col; + Ok(RecordBatch::try_new(batch.schema(), columns)?) + } +} + /// A single [`WindowApply::Aggregate`] prepared for a specific input /// partition: the Accumulator has already been seeded from the offset state. struct AggregateApply { @@ -562,7 +639,7 @@ impl AggregateApply { struct ApplyStream { input: SendableRecordBatchStream, - appliers: Vec, + appliers: Vec, schema: SchemaRef, } @@ -940,16 +1017,44 @@ mod tests { Ok(()) } - /// The Scalar apply variant is still passthrough — its output column is - /// left untouched by execute. This test guards against silent behavior - /// change until the Scalar path is implemented. + /// The scalar case — SUM, applied via the fast path (`WindowApply::Scalar` + /// with `ScalarOp::Add`, no `Accumulator` reconstructed). Rounds out the + /// demo trio (SUM/AVG/APPROX_DISTINCT), covers the operator's arrow-kernel + /// batch-arithmetic path, and shows the same shape of test works whether + /// the aggregate needs a seeded accumulator or a plain add. + /// + /// - Partition 0's mock upstream output: `[(1,1),(2,3),(3,6)]` — v and + /// the local running sum. + /// - Partition 1's mock upstream output: `[(4,4),(5,9),(6,15)]` — again + /// local, so wrong globally. + /// - Partition offsets: `[0, 6]` — the scheduler prefix-summed the + /// per-task terminal sums (partition 0 had none prior; partition 1 + /// inherits partition 0's total). + /// + /// Expected corrected running sums: partition 0 unchanged `[1,3,6]`, + /// partition 1 corrected `[10,15,21]`. #[tokio::test] - async fn scaffold_forwards_input_rows_unchanged() -> Result<()> { - let input = partitioned_source(2, 3); + async fn sum_corrects_running_sum_across_partitions() -> Result<()> { + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Int64, false), + Field::new("running_sum", DataType::Int64, false), + ])); + let batch = |vs: &[i64], sums: &[i64]| -> RecordBatch { + let v_col = Arc::new(Int64Array::from_iter_values(vs.iter().copied())); + let s_col = Arc::new(Int64Array::from_iter_values(sums.iter().copied())); + RecordBatch::try_new(schema.clone(), vec![v_col, s_col]).unwrap() + }; + let p0 = batch(&[1, 2, 3], &[1, 3, 6]); + let p1 = batch(&[4, 5, 6], &[4, 9, 15]); + let source = + MemorySourceConfig::try_new(&[vec![p0], vec![p1]], schema.clone(), None)?; + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + let apply = WindowApply::Scalar { op: ScalarOp::Add, - offset: vec![ScalarValue::Int64(Some(100)); 2], - output_column: 0, + offset: vec![ScalarValue::Int64(Some(0)), ScalarValue::Int64(Some(6))], + output_column: 1, }; let exec = Arc::new(PrefixMergeExec::try_new( input, @@ -958,22 +1063,36 @@ mod tests { )?); let ctx = SessionContext::new().task_ctx(); - for k in 0..2 { - let batches: Vec = - exec.execute(k, ctx.clone())?.try_collect().await?; - assert_eq!(batches.len(), 1); - let arr = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .expect("Int64 column"); - let expected: Vec = ((k * 3) as i64..((k + 1) * 3) as i64).collect(); - assert_eq!( - arr.values(), - &expected, - "scaffold must not apply the offset yet" - ); - } + + let out_p0: Vec = + exec.execute(0, ctx.clone())?.try_collect().await?; + let p0_sum = out_p0[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p0_sum, + vec![1, 3, 6], + "partition 0 with offset 0 should preserve local running sum" + ); + + let out_p1: Vec = + exec.execute(1, ctx.clone())?.try_collect().await?; + let p1_sum = out_p1[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + p1_sum, + vec![10, 15, 21], + "partition 1 with offset 6 should show global-cumulative running sum" + ); Ok(()) } } From 3955e1ae13f1a68637acac1e9cbdb6b53a898d0f Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Thu, 30 Jul 2026 13:06:54 -0600 Subject: [PATCH 7/7] refactor(prefix-merge): drop PARTITION BY keying to match DF-side simplification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follows apache/datafusion#24007's shape simplification: the DF getter now returns Vec>> directly rather than HashMap, since the cross-DF-partition prefix-scan use case is by construction scoped to at most one PARTITION BY group per DF partition. The DF side handles the scoping (Empty/Single/Multi slot internally, error on multi). Ballista's local FinalizedPartitionState alias matches — dropping the outer HashMap layer means execute() no longer has to project it away or error on multi-key (both handled upstream), and the demo tests construct state directly as Vec>> rather than via a HashMap keyed by an empty PartitionKey. No semantic change; the aggregate re-run tests still cover APPROX_DISTINCT and AVG, the scalar test still covers SUM, all 6 pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/src/execution_plans/prefix_merge.rs | 96 +++++++++---------- 1 file changed, 43 insertions(+), 53 deletions(-) diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index b29615091..d6e082091 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -31,9 +31,9 @@ //! - **Scheduler (global step):** collects each upstream task's finalized //! [`Accumulator::state`] via task-status transport, then computes the //! prefix-merge — for each partition `k`, combining the individual states -//! of partitions `[0..k)` into a single already-merged state per PARTITION -//! BY key and per window expression. This is the step that requires global -//! visibility across tasks; only the scheduler has it. +//! of partitions `[0..k)` into a single already-merged state per window +//! expression. This is the step that requires global visibility across +//! tasks; only the scheduler has it. //! - **Executor (local step, this operator):** receives the *already-merged* //! state for its partition in its constructor and folds it row-wise into //! the window-aggregate columns. Aggregate-agnostic by construction — the @@ -64,12 +64,15 @@ //! //! # Prefix-state input //! -//! [`FinalizedPartitionState`] — one map per input partition, holding -//! *pre-merged* [`Accumulator::state`] for every (PARTITION BY key, window -//! expression) pair. Only consumed by [`WindowApply::Aggregate`] entries; -//! [`WindowApply::Scalar`] carries its own offsets inline. The scheduler -//! bakes both when it constructs the downstream stage after the upstream -//! stage's tasks complete. +//! [`FinalizedPartitionState`] — one entry per input partition, holding +//! the *pre-merged* [`Accumulator::state`] for each aggregate window +//! expression (indexed by position in `BoundedWindowAggExec::window_expr()`). +//! Only consumed by [`WindowApply::Aggregate`] entries; +//! [`WindowApply::Scalar`] carries its own offsets inline. The DF-side getter +//! already commits to at-most-one PARTITION BY group per DF partition (see +//! `apache/datafusion#24007`), so no per-key dimension is exposed here. The +//! scheduler bakes the state when it constructs the downstream stage after +//! the upstream stage's tasks complete. //! //! **Status.** Both apply paths are implemented. //! @@ -80,9 +83,10 @@ //! arrow kernels — `numeric::add` for `Add`, `cmp::lt_eq`/`gt_eq` + `zip` //! for `Min`/`Max`, and a constant-fill for `Overwrite`. //! -//! Assumes at most one PARTITION BY key per input partition (matches the -//! AQE synthetic-PARTITION-BY pattern); multi-key handling can grow when a -//! workload needs it. +//! Inherits the DF-side getter's at-most-one-PARTITION-BY-group scoping +//! (matches the AQE synthetic-PARTITION-BY pattern) — multi-key queries +//! are handled by DataFusion's normal key-partitioned distribution and +//! don't route through this operator. //! //! # Relation to DataFusion //! @@ -105,7 +109,6 @@ use datafusion::common::{Result, ScalarValue, Statistics, internal_err}; use datafusion::execution::TaskContext; use datafusion::logical_expr::{Accumulator, AggregateUDF}; use datafusion::physical_expr::aggregate::AggregateExprBuilder; -use datafusion::physical_expr::window::PartitionKey; use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{ @@ -113,19 +116,20 @@ use datafusion::physical_plan::{ PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use futures::{Stream, StreamExt, ready}; -use std::collections::HashMap; use std::pin::Pin; use std::task::{Context, Poll}; -/// A single already-prefix-merged window-aggregate state, keyed by PARTITION -/// BY tuple. The inner `Vec` is indexed by window expression (same order the -/// upstream `BoundedWindowAggExec` reports in `window_expr()`); `None` at a -/// slot indicates a non-aggregate window function (`row_number`, `rank`, -/// `lead`/`lag`, ...) which contributes no state. +/// A single already-prefix-merged window-aggregate state. Indexed by window +/// expression (same order the upstream `BoundedWindowAggExec` reports in +/// `window_expr()`); `None` at a slot indicates a non-aggregate window +/// function (`row_number`, `rank`, `lead`/`lag`, ...) that contributes no +/// state. 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 scheduler produces one of these per input partition, having already /// combined the individual states from every prior partition into one merged -/// value per (key, window expr) — see the [module-level docs][self] for the +/// value per window expression — see the [module-level docs][self] for the /// division of labor. This operator applies it; it does not compute it. /// /// The type mirrors `datafusion::physical_plan::windows::FinalizedPartitionState` @@ -133,7 +137,7 @@ use std::task::{Context, Poll}; /// this crate compiles against stable DataFusion 54 until that PR lands. /// /// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 -pub type FinalizedPartitionState = HashMap>>>; +pub type FinalizedPartitionState = Vec>>; /// How to combine each row's existing value in an output column with a /// scheduler-provided scalar offset. The result overwrites the column. @@ -427,21 +431,11 @@ impl ExecutionPlan for PrefixMergeExec { let input_schema = self.input.schema(); let output_schema = self.schema(); - // Pick this partition's pre-merged state. For the AQE synthetic- - // PARTITION-BY case each task carries exactly one key; multi-key - // handling can grow when a real workload needs it. - let state = &self.per_partition_state[partition]; - let key_state: Option<&Vec>>> = match state.len() { - 0 => None, - 1 => state.values().next(), - n => { - return internal_err!( - "PrefixMergeExec: {n} PARTITION BY keys in the state map \ - for partition {partition}; multi-key apply is not yet \ - implemented" - ); - } - }; + // The DF-side getter (BoundedWindowAggExec::finalized_partition_state) + // already commits to at-most-one PARTITION BY group per DF partition, + // so this operator receives the state indexed only by window + // expression — no PARTITION BY dimension to project away. + let key_state = &self.per_partition_state[partition]; let mut appliers: Vec = Vec::with_capacity(self.applies.len()); for (i, apply) in self.applies.iter().enumerate() { @@ -453,7 +447,7 @@ impl ExecutionPlan for PrefixMergeExec { window_expr_index, } => { let offset_state: Option<&Vec> = key_state - .and_then(|per_expr| per_expr.get(*window_expr_index)) + .get(*window_expr_index) .and_then(|slot| slot.as_ref()); appliers.push(PreparedApply::Aggregate(AggregateApply::new( i, @@ -703,9 +697,7 @@ mod tests { } fn empty_state(partitions: usize) -> Vec { - (0..partitions) - .map(|_| FinalizedPartitionState::new()) - .collect() + (0..partitions).map(|_| Vec::new()).collect() } /// Mismatched state length surfaces as an error rather than a panic at @@ -713,9 +705,8 @@ mod tests { #[test] fn try_new_rejects_state_length_mismatch() { let input = partitioned_source(2, 3); - let err = - PrefixMergeExec::try_new(input, vec![], vec![FinalizedPartitionState::new()]) - .expect_err("length mismatch must surface as an error"); + let err = PrefixMergeExec::try_new(input, vec![], vec![Vec::new()]) + .expect_err("length mismatch must surface as an error"); assert!( err.to_string() .contains("does not match input partition count"), @@ -837,12 +828,11 @@ mod tests { &Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])), &[1, 2, 3], )?; - // Assemble the per-partition state maps. PARTITION BY tuple is empty - // (global window); one aggregate at window_expr_index 0. - let empty_key: PartitionKey = vec![]; - let mut state_p1 = FinalizedPartitionState::new(); - state_p1.insert(empty_key.clone(), vec![Some(p0_state)]); - let per_partition_state = vec![FinalizedPartitionState::new(), state_p1]; + // Partition 1 gets partition 0's HLL state at the one aggregate slot + // (window_expr_index 0). Partition 0 gets an empty state — nothing + // to merge in. + let per_partition_state: Vec = + vec![vec![], vec![Some(p0_state)]]; let apply = WindowApply::Aggregate { udf: Arc::clone(&udf), @@ -963,10 +953,10 @@ mod tests { acc.update_batch(&[arr])?; acc.state()? }; - let empty_key: PartitionKey = vec![]; - let mut state_p1 = FinalizedPartitionState::new(); - state_p1.insert(empty_key.clone(), vec![Some(p0_state)]); - let per_partition_state = vec![FinalizedPartitionState::new(), state_p1]; + // Partition 1 gets partition 0's (sum, count) state at the one + // aggregate slot; partition 0 has nothing to merge in. + let per_partition_state: Vec = + vec![vec![], vec![Some(p0_state)]]; let apply = WindowApply::Aggregate { udf: Arc::clone(&udf),