diff --git a/ballista/client/tests/prefix_window.rs b/ballista/client/tests/prefix_window.rs new file mode 100644 index 000000000..50642d8f1 --- /dev/null +++ b/ballista/client/tests/prefix_window.rs @@ -0,0 +1,210 @@ +// 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. + +//! End-to-end correctness for the prefix-scan window rewrite. +//! +//! `PrefixWindowRule` splits an `UNBOUNDED PRECEDING` window across K tasks, +//! each computing a partition-local running aggregate, and corrects them +//! afterwards with the merged state of all prior partitions. The failure +//! mode it has to be held against is the one that doesn't show up on a +//! single node: each partition's running total silently restarting at zero. +//! +//! These run the real distributed path (scheduler, shuffle, executor) via +//! the standalone in-process cluster, so a wrong stage boundary or a lost +//! accumulator state shows up as wrong numbers rather than a plan diff. + +mod common; + +#[cfg(test)] +#[cfg(feature = "standalone")] +mod prefix_window_tests { + use ballista::prelude::{SessionConfigExt, SessionContextExt}; + use ballista_core::config::{ + BALLISTA_ADAPTIVE_PLANNER_ENABLED, BALLISTA_PARALLEL_WINDOW_ENABLED, + BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK, + }; + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::common::Result; + use datafusion::execution::SessionStateBuilder; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; + use std::fs::File; + use std::sync::Arc; + use tempfile::TempDir; + + /// Rows per input partition. + const ROWS_PER_PARTITION: usize = 4; + /// Input partitions, and therefore the K the rewrite range-repartitions to. + const PARTITIONS: usize = 4; + const TOTAL_ROWS: usize = ROWS_PER_PARTITION * PARTITIONS; + + /// `1.0 ..= 16.0`, split across [`PARTITIONS`] partitions in ascending + /// runs. Every value and every running sum of them is exactly + /// representable in `f64`, so a mismatch is a semantic bug rather than + /// floating-point drift. + fn input_partitions() -> Vec> { + (0..PARTITIONS) + .map(|p| { + (0..ROWS_PER_PARTITION) + .map(|r| (p * ROWS_PER_PARTITION + r + 1) as f64) + .collect() + }) + .collect() + } + + /// Running sums of `1.0 ..= 16.0` in ascending order: 1, 3, 6, 10, ... + /// + /// Computed here rather than captured from a Ballista run, so the test + /// can't agree with a uniformly-wrong engine. + fn expected_running_sums() -> Vec { + let mut total = 0.0; + (1..=TOTAL_ROWS) + .map(|v| { + total += v as f64; + total + }) + .collect() + } + + /// `max_partitions_per_task`, so tasks carry a multi-partition slice + /// rather than one partition each. That exercises the task builder's + /// index remapping: `PrefixMergeExec`'s state is keyed by global + /// partition, but a task's `execute(k)` numbers its own slice from zero. + const MAX_PARTITIONS_PER_TASK: usize = 2; + + async fn context(parallel_window: bool) -> SessionContext { + let config = SessionConfig::new_with_ballista() + .with_target_partitions(PARTITIONS) + .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, true) + .set_str( + BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK, + &MAX_PARTITIONS_PER_TASK.to_string(), + ) + .set_bool(BALLISTA_PARALLEL_WINDOW_ENABLED, parallel_window); + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + SessionContext::standalone_with_state(state).await.unwrap() + } + + /// Register `t(v Float64)` as one parquet file per partition. + /// + /// Parquet rather than a `MemTable` because Ballista ships the *logical* + /// plan to the scheduler, and an in-memory provider has no + /// `LogicalExtensionCodec` — it fails at serialization before any of the + /// distributed path runs. One file per partition so the scan really is + /// [`PARTITIONS`]-wide rather than depending on how DataFusion chooses to + /// split a single small file. + /// + /// `v` is `Float64` because ORRE routes on a T-Digest, which is + /// Float64-only until the KLL migration. That restriction is what keeps + /// h2o Q7 (`ORDER BY id3`, `Int64`) from being rewritten today. + /// + /// The returned [`TempDir`] owns the files and must outlive the query. + async fn register_input(ctx: &SessionContext) -> Result { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let dir = TempDir::new().expect("temp dir"); + for (partition, values) in input_partitions().into_iter().enumerate() { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float64Array::from(values))], + )?; + let path = dir.path().join(format!("part-{partition}.parquet")); + let file = File::create(&path).expect("create parquet file"); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None)?; + writer.write(&batch)?; + writer.close()?; + } + ctx.register_parquet( + "t", + dir.path().to_str().expect("utf-8 temp path"), + ParquetReadOptions::default(), + ) + .await?; + Ok(dir) + } + + /// `(v, running_sum)` pairs ordered by `v`. + async fn running_sums(ctx: &SessionContext) -> Result> { + let batches = ctx + .sql( + "SELECT v, \ + sum(v) OVER (ORDER BY v \ + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS rs \ + FROM t \ + ORDER BY v", + ) + .await? + .collect() + .await?; + + let mut rows = Vec::with_capacity(TOTAL_ROWS); + for batch in &batches { + let v = datafusion::common::cast::as_float64_array(batch.column(0))?; + let rs = datafusion::common::cast::as_float64_array(batch.column(1))?; + for i in 0..batch.num_rows() { + rows.push((v.value(i), rs.value(i))); + } + } + Ok(rows) + } + + /// The rewrite must not change the answer. + /// + /// The failure this guards against is quiet: if prior-partition state + /// never reaches `PrefixMergeExec`, each task emits its own + /// partition-local running sum and every partition after the first is + /// short by exactly the sum of everything before it. Plausible-looking + /// numbers, wrong totals, no error. + #[tokio::test] + async fn prefix_scan_matches_serial_running_sum() -> Result<()> { + let ctx = context(true).await; + let _data = register_input(&ctx).await?; + let rows = running_sums(&ctx).await?; + + let expected_v: Vec = (1..=TOTAL_ROWS).map(|v| v as f64).collect(); + let actual_v: Vec = rows.iter().map(|(v, _)| *v).collect(); + assert_eq!(actual_v, expected_v, "input rows must survive the rewrite"); + + let actual_rs: Vec = rows.iter().map(|(_, rs)| *rs).collect(); + assert_eq!( + actual_rs, + expected_running_sums(), + "running sums must be global, not partition-local — a partition \ + whose total restarts near zero means prior-partition state never \ + reached PrefixMergeExec" + ); + Ok(()) + } + + /// Same query with the rewrite off, as a guard on the test itself: if + /// this ever fails, the harness is wrong rather than the rewrite. + #[tokio::test] + async fn serial_running_sum_is_correct_without_rewrite() -> Result<()> { + let ctx = context(false).await; + let _data = register_input(&ctx).await?; + let rows = running_sums(&ctx).await?; + + let actual_rs: Vec = rows.iter().map(|(_, rs)| *rs).collect(); + assert_eq!(actual_rs, expected_running_sums()); + Ok(()) + } +} diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index bd20ed65d..1544001aa 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -59,6 +59,7 @@ message BallistaPhysicalPlanNode { PartitionedBoundedWindowAggExecNode partitioned_bounded_window_agg = 11; RangeShuffleReaderExecNode range_shuffle_reader = 12; RangeFilterExecNode range_filter = 13; + PrefixMergeExecNode prefix_merge = 14; } } @@ -158,6 +159,72 @@ message RangeBound { datafusion_common.ScalarValue hi = 2; } +// Applies scheduler-computed prefix state to a window-aggregate column, so +// each range-disjoint task's local running aggregate becomes a global one. +message PrefixMergeExecNode { + // One entry per output column needing correction. Empty makes the operator + // a passthrough. + repeated WindowApplyNode applies = 1; + // Resolved prefix state, one entry per input partition, `[k]` summarising + // every partition before `k`. Encoding refuses while unresolved — an + // over-the-wire plan always ships with state bound, since an executor has + // no way to obtain it. + repeated FinalizedPartitionStateNode per_partition_state = 2; +} + +// How to correct one window-function output column. +message WindowApplyNode { + oneof apply { + ScalarWindowApplyNode scalar = 1; + AggregateWindowApplyNode aggregate = 2; + } +} + +// Fast path: combine each row's value with a per-partition scalar. +message ScalarWindowApplyNode { + ScalarOpNode op = 1; + // One scalar per input partition. + repeated datafusion_common.ScalarValue offset = 2; + uint32 output_column = 3; +} + +// How a scalar offset combines with a row's existing value. +enum ScalarOpNode { + SCALAR_OP_NODE_ADD = 0; + SCALAR_OP_NODE_MIN = 1; + SCALAR_OP_NODE_MAX = 2; + // Ignores the row's value entirely. + SCALAR_OP_NODE_OVERWRITE = 3; +} + +// Fallback path: seed a fresh accumulator from the prefix state and replay +// each row through it. Needed where a row's output is not itself a valid +// partial state — AVG, and sketch-backed aggregates like approx_distinct. +message AggregateWindowApplyNode { + // Resolved from the executor's function registry on decode. + string udf_name = 1; + repeated datafusion.PhysicalExprNode args = 2; + uint32 output_column = 3; + // Position in the upstream window operator's `window_expr()` list, which + // is what indexes into each FinalizedPartitionStateNode. + uint32 window_expr_index = 4; +} + +// Prefix state for one input partition: one slot per window expression. +message FinalizedPartitionStateNode { + repeated AggregateStateSlotNode slots = 1; +} + +message AggregateStateSlotNode { + // Unset when that window expression published no state — a non-aggregate + // window function. Distinct from a present-but-empty state. + AggregateStateNode state = 1; +} + +message AggregateStateNode { + repeated datafusion_common.ScalarValue values = 1; +} + // Wrapper for `BoundedWindowAggExec` that overrides // `required_input_distribution` to `Unspecified` — see the module doc on // `execution_plans::partitioned_bounded_window_agg` for what makes that safe. @@ -656,6 +723,48 @@ message SuccessfulTask { // executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag // to combine reports across tasks/executors. repeated RuntimeStatsReport runtime_stats = 3; + // Finalized window-aggregate state captured during this task, one entry + // per (output partition, window expression, PARTITION BY group) that + // closed. Empty unless the plan contains an ever-expanding-frame window + // (`UNBOUNDED PRECEDING`), which is the only shape DataFusion will + // publish accumulator state for. The scheduler prefix-merges these across + // tasks and bakes the result into a downstream `PrefixMergeExec`. + // + // TODO: watch the size of this. Task completion is a hot, frequent + // message, and sketch-backed aggregates make the payload unbounded in a + // way row counts and quantile sketches are not — an HLL or KLL state is + // kilobytes per window expression per partition, and a task covering a + // wide partition slice carries one of each. If it stops being small, + // write the state as a sidecar next to the shuffle files instead, the way + // sort-shuffle already writes `.arrow.index` beside its data + // (`sort_shuffle::get_index_path`), and send only a reference here. That + // keeps the completion message fixed-size regardless of aggregate. + repeated WindowStateReport window_state = 4; +} + +// One finalized window-aggregate state from a task's +// `BoundedWindowAggExec`. +message WindowStateReport { + // The stage's *global* output partition this state belongs to. + // + // DataFusion reports a task-local index, because a task's plan is + // restricted to a partition slice. The producing `ShuffleWriterExec` + // translates it through the `global_output_partition_ids` the scheduler + // stamped on it, so what crosses the wire is already global — the + // scheduler must not re-derive it. + uint32 global_partition_id = 1; + // Position in the window operator's `window_expr()` list. Indexes window + // expressions, not partitions. + uint32 window_expr_index = 2; + // The PARTITION BY tuple that closed. Empty for a window with no + // PARTITION BY, which is the only shape the prefix rewrite plants today. + repeated datafusion_common.ScalarValue partition_key = 3; + // `Accumulator::state` for the closed group: one element for SUM / COUNT / + // MIN / MAX, two for AVG's (sum, count), one opaque Binary for + // sketch-backed aggregates like approx_distinct. Carried as ScalarValue + // rather than a numeric field so non-decomposable aggregates work + // unchanged. + repeated datafusion_common.ScalarValue state = 4; } // One report per `RuntimeStatsExec` in the executed plan. diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 8c83050bd..97a223995 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -26,6 +26,7 @@ mod ordered_range_repartition; mod partitioned_bounded_window_agg; mod per_partition_filter; pub mod plan_algebra; +mod prefix_merge; mod range_filter; mod range_repartition_common; mod range_shuffle_reader; @@ -36,6 +37,7 @@ mod shuffle_writer_trait; pub mod sort_shuffle; mod unordered_range_repartition; mod unresolved_shuffle; +pub mod window_state; use std::path::{Path, PathBuf}; @@ -48,6 +50,7 @@ pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec; pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates}; pub use plan_algebra::{preserves_distribution, preserves_partitioning}; +pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec, ScalarOp, WindowApply}; pub use range_filter::{RangeBound, RangeFilterExec, WidenedBound}; pub use range_shuffle_reader::RangeShuffleReaderExec; pub use runtime_stats::{ @@ -65,6 +68,10 @@ pub use shuffle_writer_trait::ShuffleWriter; pub use sort_shuffle::SortShuffleWriterExec; pub use unordered_range_repartition::UnorderedRangeRepartitionExec; pub use unresolved_shuffle::UnresolvedShuffleExec; +pub use window_state::{ + ObservedWindowState, TaskWindowState, WindowStateCollector, + prefix_merge_window_state, window_state_from_proto, window_state_to_proto, +}; use crate::JobId; diff --git a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs index ccb4ab561..c00a74d8d 100644 --- a/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs +++ b/ballista/core/src/execution_plans/partitioned_bounded_window_agg.rs @@ -68,8 +68,11 @@ use datafusion::execution::TaskContext; use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{CardinalityEffect, InputOrderMode}; use datafusion::physical_plan::metrics::MetricsSet; -use datafusion::physical_plan::windows::BoundedWindowAggExec; -use datafusion::physical_plan::windows::WindowExpr; +use datafusion::physical_plan::windows::{ + BoundedWindowAggExec, WindowExpr, WindowStateObserver, +}; + +use crate::execution_plans::window_state::{ObservedWindowState, WindowStateCollector}; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, StatisticsArgs, statistics::ChildStats, @@ -93,6 +96,18 @@ pub struct PartitionedBoundedWindowAggExec { inner_bwag: Arc, /// Multi-partition input; same `Arc` `inner_bwag.input()` holds. input: Arc, + /// Installed on `inner_bwag` when every frame permits it; see + /// [`Self::try_new`]. `None` for the halo shape, whose sliding frames + /// DataFusion refuses to observe. + /// + /// TODO: move the install off this wrapper. Catching the state has + /// nothing to do with overriding a distribution declaration — they are + /// colocated only because this wrapper happens to own the BWAG today. + /// When DataFusion's BWAG can declare a non-single input distribution + /// itself and this type collapses, the install moves to whatever ends up + /// holding the BWAG (a bare BWAG, or a small node planted above it by + /// the prefix rule). `WindowStateCollector` itself is unaffected. + state_collector: Option>, } impl PartitionedBoundedWindowAggExec { @@ -104,13 +119,61 @@ impl PartitionedBoundedWindowAggExec { window_expr: Vec>, input: Arc, ) -> Result { - let inner_bwag = Arc::new(BoundedWindowAggExec::try_new( + let bwag = BoundedWindowAggExec::try_new( window_expr, input.clone(), BWAG_INPUT_ORDER_MODE, BWAG_CAN_REPARTITION, - )?); - Ok(Self { inner_bwag, input }) + )?; + // Install a collector exactly when DataFusion will accept one: every + // frame ever-expanding, i.e. `UNBOUNDED PRECEDING`. A sliding frame's + // accumulator retracts as the frame advances, so at partition close + // it holds the last frame rather than the partition aggregate, and + // `with_state_observer` refuses it. + // + // Deriving this from the frames rather than taking it as a flag keeps + // the wire format unchanged: the executor rebuilds the same decision + // from the window expressions it decodes. It also means the halo + // rewrite, which plants this wrapper over finite frames, is untouched. + let state_collector = bwag + .window_expr() + .iter() + .all(|expr| expr.get_window_frame().is_ever_expanding()) + .then(|| Arc::new(WindowStateCollector::new(bwag.window_expr().to_vec()))); + let bwag = match &state_collector { + Some(collector) => bwag.with_state_observer(Some( + Arc::clone(collector) as Arc + ))?, + None => bwag, + }; + Ok(Self { + inner_bwag: Arc::new(bwag), + input, + state_collector, + }) + } + + /// The installed collector, or `None` when the frames don't permit one. + /// + /// Drained after the task completes to recover this task's contribution + /// to the cross-task prefix scan. + pub fn state_collector(&self) -> Option<&Arc> { + self.state_collector.as_ref() + } + + /// Everything the collector captured, or empty when no collector was + /// installed. + /// + /// Each entry's `partition_idx` is **task-local** — it indexes this + /// task's partition slice, not the stage's global partitions. This + /// operator has no way to know otherwise, and shouldn't: the writer at + /// the stage root is handed the slice-to-global mapping by the scheduler + /// and does the translation. + pub fn observed_window_state(&self) -> Vec { + self.state_collector + .as_ref() + .map(|collector| collector.observed()) + .unwrap_or_default() } /// The wrapped `BoundedWindowAggExec` — for accessors that don't exist diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs index 245a8c377..1f033dcb6 100644 --- a/ballista/core/src/execution_plans/plan_algebra.rs +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -25,6 +25,9 @@ //! conservative is the safety net: unrecognized node → property assumed //! false → caller falls back to the safer path. +use std::sync::Arc; + +use datafusion::common::{DataFusionError, Result}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::projection::ProjectionExec; @@ -32,7 +35,8 @@ use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::execution_plans::{ - BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, + BufferExec, PrefixMergeExec, RangeFilterExec, RuntimeStatsExec, ShuffleWriterExec, + SortShuffleWriterExec, }; /// Whitelisted ops preserve the routing key's row set, values, and @@ -68,3 +72,80 @@ pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool { // Stats tap; no data mutation. || plan.downcast_ref::().is_some() } + +/// Operators carrying data indexed by *global* input partition, which the +/// scheduler must slice when it restricts a stage plan to one task's +/// partition subset. +/// +/// The executor still never learns its task's global identity — slicing is +/// how that stays true. After restriction an operator holds a vec covering +/// exactly the partitions its task will run, in local order, so `execute(k)` +/// indexes it directly. +/// +/// Implement this next to the fields being sliced. The alternative — the +/// scheduler's task builder knowing each operator's internals — means adding +/// a partition-indexed field silently leaves a stale slicer two crates away, +/// discovered at runtime if some length check happens to catch it. +/// +/// Like the whitelists above, this exists because `ExecutionPlan` has no +/// "restrict yourself to a subset of input partitions" method. A defaulted +/// upstream one would let each operator answer for itself and retire +/// [`as_partition_sliceable`]. +pub trait PartitionSliceable: ExecutionPlan { + /// Rebuild over `child` — already restricted to `partitions` — carrying + /// only the entries for `partitions`, in that order. + /// + /// # Arguments + /// + /// * `child` - this operator's single input, already restricted + /// * `partitions` - global input partition indices this task will run + fn slice_to_partitions( + &self, + child: Arc, + partitions: &[usize], + ) -> Result>; +} + +/// Hand-maintained whitelist of [`PartitionSliceable`] operators, for the +/// same reason as the property whitelists above: `dyn ExecutionPlan` can't +/// be downcast to an arbitrary trait. +pub fn as_partition_sliceable( + plan: &Arc, +) -> Option<&dyn PartitionSliceable> { + if let Some(op) = plan.downcast_ref::() { + return Some(op); + } + if let Some(op) = plan.downcast_ref::() { + return Some(op); + } + None +} + +/// Take `values[global]` for each global partition index, in task-local +/// order. Errors rather than dropping, so a slice that outruns its operator's +/// data surfaces as a named failure instead of a silently short vec. +/// +/// # Arguments +/// +/// * `values` - the operator's per-global-partition entries +/// * `partitions` - global indices to keep, in task-local order +/// * `owner` - operator name, for the error message +/// * `what` - what the entries are, for the error message +pub fn slice_by_global_partition( + values: &[T], + partitions: &[usize], + owner: &str, + what: &str, +) -> Result> { + partitions + .iter() + .map(|&global| { + values.get(global).cloned().ok_or_else(|| { + DataFusionError::Internal(format!( + "{owner}: partition index {global} out of bounds ({} {what})", + values.len() + )) + }) + }) + .collect() +} 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..307d3fc94 --- /dev/null +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -0,0 +1,1481 @@ +// 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)`. +//! +//! # 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 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. +//! +//! # Apply descriptors +//! +//! `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 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. No PARTITION BY +//! dimension is exposed: DataFusion publishes state per closed group, but the +//! rewrite that plants this operator only fires on windows without a +//! PARTITION BY, and the scheduler rejects any report carrying a key rather +//! than flattening two groups together. A window that does have a PARTITION +//! BY needs no help from this operator, since `BoundedWindowAggExec` asks for +//! `KeyPartitioned` input and each partition's window is already independent. +//! The scheduler bakes the state when it constructs the downstream stage +//! after the upstream stage's tasks complete. +//! +//! **Status.** Both apply paths are implemented. +//! +//! - [`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`. +//! +//! 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 +//! +//! The upstream tasks' finalized state — which the scheduler prefix-merges +//! before handing the result to this operator — comes from the accumulators +//! inside `BoundedWindowAggExec`, captured by a `WindowStateObserver` and +//! shipped to the scheduler on task completion. See +//! [`window_state`](super::window_state) for that half. +//! +//! An empty `applies` list makes this operator a passthrough. A non-empty one +//! always rewrites its output columns, including when a state slot is `None` — +//! that seeds the accumulator with nothing rather than skipping the column, so +//! the result is the partition-local aggregate, not the input value. +//! +//! [`Accumulator::state`]: datafusion::logical_expr::Accumulator::state +//! [`Accumulator::merge_batch`]: datafusion::logical_expr::Accumulator::merge_batch + +use std::fmt::{self, Debug, Formatter}; +use std::sync::Arc; + +use log::debug; +use parking_lot::Mutex; + +use datafusion::arrow::array::{ArrayRef, RecordBatch}; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +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::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::{ + ColumnarValue, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, StatisticsArgs, + apply_expression_roots, statistics::ChildStats, +}; +use futures::{Stream, StreamExt, ready}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use crate::execution_plans::plan_algebra::{ + PartitionSliceable, slice_by_global_partition, +}; + +/// 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 window expression — see the `prefix_merge` module docs for the +/// division of labor. This operator applies it; it does not compute it. +/// +/// A newtype rather than an alias for `Vec>>`. That +/// type appears in this operator's public signatures, where it is neither +/// readable nor searchable, and it spells "no state for this window +/// expression" two ways — a missing index and a `None` — which every caller +/// then has to handle. [`Self::slot`] collapses both into one answer. Any +/// later change to the representation also stays internal rather than +/// breaking whoever wrote the concrete type. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct FinalizedPartitionState { + /// Indexed by position in the upstream operator's `window_expr()` list. + per_window_expr: Vec>>, +} + +impl FinalizedPartitionState { + /// Build from one slot per window expression, in `window_expr()` order. + pub fn new(per_window_expr: Vec>>) -> Self { + Self { per_window_expr } + } + + /// Merged state for `window_expr_index`, or `None` when that expression + /// published none or the index is past the end. + pub fn slot(&self, window_expr_index: usize) -> Option<&Vec> { + self.per_window_expr + .get(window_expr_index) + .and_then(|slot| slot.as_ref()) + } + + /// Every slot in window-expression order, for wire encoding. + pub fn slots(&self) -> &[Option>] { + &self.per_window_expr + } + + /// Number of window expressions this state covers. + pub fn len(&self) -> usize { + self.per_window_expr.len() + } + + /// True when no window expression published state. + pub fn is_empty(&self) -> bool { + self.per_window_expr.is_empty() + } +} + +/// 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 +/// +/// `#[non_exhaustive]`: more ops will land (the ranking family, and whatever +/// a segment-tree correction needs), and each would otherwise be a breaking +/// change for anyone matching on this. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +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` over an + /// ever-expanding frame: every row's answer is the same global first + /// value, so the scheduler picks it once and every row gets a copy. + /// + /// Not `last_value`. Over an ever-expanding frame that is the current + /// row's own value and needs no correction at all, so overwriting every + /// row with one scalar would be wrong. + 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). +/// `#[non_exhaustive]`: the two shapes here cover what the prefix rewrite +/// plants today, and a third (a segment-tree broadcast for the ranking +/// family) is anticipated. +#[derive(Debug, Clone)] +#[non_exhaustive] +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 `prefix_merge` module +/// docs for the division of labor between scheduler and executor and the AQE +/// pipeline this fits into. +/// +/// [`WindowApply::Aggregate`] entries are applied per row via a seeded +/// `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 + /// 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 *already-merged* 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()`. + /// + /// Late-bound: `None` until [`PrefixMergeExec::resolve_state`]. The rule + /// plants this operator at plan time, but the state only exists once + /// every upstream task has closed its window and published accumulator + /// state. `execute` refuses while unresolved rather than treating an + /// absent carry-in as zero, which would silently emit partition-local + /// aggregates. + per_partition_state: Arc>>>, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl PrefixMergeExec { + /// Wrap `input` with per-column apply descriptors, state pending. + /// + /// The rewrite rule's path: the operator is planted at plan time and the + /// scheduler calls [`Self::resolve_state`] once the upstream stage's + /// tasks have reported. + pub fn try_new_pending( + input: Arc, + applies: Vec, + ) -> Result { + Self::try_new_inner(input, applies, None) + } + + /// Wrap `input` with state already known — wire decode, and + /// task-restriction, which slices state parallel to the input. + pub fn try_new_resolved( + input: Arc, + applies: Vec, + per_partition_state: Vec, + ) -> Result { + Self::try_new_inner(input, applies, Some(per_partition_state)) + } + + /// Errors on any of: + /// - `per_partition_state.len()` != input's partition count, when given. + /// - Any [`WindowApply::Scalar`]'s `offset.len()` != input's partition + /// count. + /// - Any entry's `output_column` outside the input schema's field range. + fn try_new_inner( + input: Arc, + applies: Vec, + per_partition_state: Option>, + ) -> Result { + let partition_count = input.output_partitioning().partition_count(); + if let Some(state) = per_partition_state.as_ref() + && state.len() != partition_count + { + return internal_err!( + "PrefixMergeExec: per_partition_state.len() {} does not match \ + input partition count {}", + state.len(), + 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(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + input, + applies, + per_partition_state: Arc::new(Mutex::new(per_partition_state)), + properties, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + /// Bind the prefix state. Called by the scheduler once the upstream + /// stage's tasks have all reported their finalized accumulator state and + /// it has been prefix-merged into one carry-in per partition. + /// + /// Idempotent overwrite, matching `RangeFilterExec::resolve_bounds`: AQE + /// re-plans, and a later pass may resolve the same operator again with + /// the same values. + pub fn resolve_state(&self, state: Vec) -> Result<()> { + let partition_count = self.input.output_partitioning().partition_count(); + if state.len() != partition_count { + return internal_err!( + "PrefixMergeExec: resolve_state got {} entries for {} input \ + partitions", + state.len(), + partition_count + ); + } + debug!( + "PrefixMergeExec: resolved prefix state for {} partitions: {:?}", + state.len(), + state + ); + self.per_partition_state.lock().replace(state); + Ok(()) + } + + /// 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 + } + + /// This operator's input. The scheduler descends from here to find the + /// state-sync boundary and the window operator whose expressions the + /// reported state is indexed against. + pub fn input(&self) -> &Arc { + &self.input + } + + /// The prefix state carried per input partition, or `None` before + /// [`Self::resolve_state`]. Only consumed by [`WindowApply::Aggregate`] + /// entries. + pub fn per_partition_state(&self) -> Option> { + self.per_partition_state.lock().clone() + } +} + +impl Debug for PrefixMergeExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("PrefixMergeExec") + .field( + "partition_count", + &self.per_partition_state.lock().as_ref().map(|s| s.len()), + ) + .field("applies", &self.applies.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={}, applies={}", + self.per_partition_state + .lock() + .as_ref() + .map_or(-1_i64, |s| s.len() as i64), + self.applies.len() + ) + } + DisplayFormatType::TreeRender => { + write!(f, "PrefixMergeExec") + } + } + } +} + +impl PartitionSliceable for PrefixMergeExec { + /// Both the per-partition state and every [`WindowApply::Scalar`]'s + /// offsets are indexed by global input partition, so both slice parallel + /// to the input. [`WindowApply::Aggregate`] carries `window_expr_index`, + /// which indexes window expressions rather than partitions, so it rides + /// over unchanged. + fn slice_to_partitions( + &self, + child: Arc, + partitions: &[usize], + ) -> Result> { + let applies = self + .applies + .iter() + .map(|apply| match apply { + WindowApply::Scalar { + op, + offset, + output_column, + } => Ok(WindowApply::Scalar { + op: *op, + offset: slice_by_global_partition( + offset, + partitions, + "PrefixMergeExec", + "scalar offsets", + )?, + output_column: *output_column, + }), + aggregate => Ok(aggregate.clone()), + }) + .collect::>>()?; + let state = self.per_partition_state.lock().clone().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "PrefixMergeExec: task-restriction before resolve_state()".into(), + ) + })?; + Ok(Arc::new(Self::try_new_resolved( + child, + applies, + slice_by_global_partition( + &state, + partitions, + "PrefixMergeExec", + "state slots", + )?, + )?)) + } +} + +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] + } + + /// Only [`WindowApply::Aggregate`] holds expressions. The scalar path's + /// offsets are already-evaluated [`ScalarValue`]s baked in by the + /// scheduler, so there is nothing there for a rewriter to visit. + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots( + self.applies.iter().flat_map(|apply| match apply { + WindowApply::Aggregate { args, .. } => args.as_slice(), + WindowApply::Scalar { .. } => [].as_slice(), + }), + f, + ) + } + + 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() + ); + }; + let rebuilt = match self.per_partition_state.lock().clone() { + Some(state) => PrefixMergeExec::try_new_resolved( + input.clone(), + self.applies.clone(), + state, + )?, + None => { + PrefixMergeExec::try_new_pending(input.clone(), self.applies.clone())? + } + }; + Ok(Arc::new(rebuilt)) + } + + /// 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] + } + + /// Row count and per-column stats pass through unchanged: the merge + /// rewrites values in the window-aggregate columns but adds no rows. + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + /// Every input row is emitted exactly once. + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + let resolved = self.per_partition_state.lock().clone(); + // Refuse rather than treat an absent carry-in as zero: that would + // emit partition-local aggregates that look plausible and are wrong. + let Some(resolved) = resolved else { + return internal_err!( + "PrefixMergeExec: execute() called before resolve_state()" + ); + }; + if partition >= resolved.len() { + return internal_err!( + "PrefixMergeExec: partition {} out of bounds ({} slots)", + partition, + resolved.len() + ); + } + let input_schema = self.input.schema(); + let output_schema = self.schema(); + + // Indexed by window expression only. The scheduler rejects any + // report carrying a PARTITION BY key, so at most one group per + // partition reaches here and there is no key dimension to project + // away. + let key_state = &resolved[partition]; + + let mut appliers: Vec = Vec::with_capacity(self.applies.len()); + for (i, apply) in self.applies.iter().enumerate() { + match apply { + WindowApply::Aggregate { + udf, + args, + output_column, + window_expr_index, + } => { + let offset_state = key_state.slot(*window_expr_index); + appliers.push(PreparedApply::Aggregate(AggregateApply::new( + i, + udf, + args, + *output_column, + &input_schema, + offset_state, + )?)); + } + WindowApply::Scalar { + op, + offset, + output_column, + } => { + appliers.push(PreparedApply::Scalar(ScalarApply { + apply_index: i, + op: *op, + offset: offset[partition].clone(), + output_column: *output_column, + })); + } + } + } + + let input = self.input.execute(partition, ctx)?; + let stream = ApplyStream { + input, + appliers, + schema: Arc::clone(&output_schema), + baseline: BaselineMetrics::new(&self.metrics, partition), + path_metrics: PathMetrics { + scalar_time: MetricBuilder::new(&self.metrics) + .subset_time("scalar_apply_time", partition), + aggregate_time: MetricBuilder::new(&self.metrics) + .subset_time("aggregate_apply_time", partition), + rows_corrected: MetricBuilder::new(&self.metrics) + .counter("rows_corrected", partition), + }, + }; + 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; + rebuild_batch( + batch.schema(), + columns, + self.apply_index, + self.output_column, + ) + } +} + +/// 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; + rebuild_batch( + batch.schema(), + columns, + self.apply_index, + self.output_column, + ) + } +} + +/// Time split by apply path, so the accumulator replay's cost is a number the +/// operator reports rather than something only a benchmark can see. +/// +/// The two paths differ by more than a constant: [`WindowApply::Scalar`] is an +/// arrow kernel over a whole batch, while [`WindowApply::Aggregate`] seeds an +/// accumulator and replays every row through it. A sketch-heavy query pays the +/// second and a SUM-heavy one need not, which a single total would hide. +struct PathMetrics { + scalar_time: Time, + aggregate_time: Time, + rows_corrected: Count, +} + +struct ApplyStream { + input: SendableRecordBatchStream, + appliers: Vec, + schema: SchemaRef, + baseline: BaselineMetrics, + path_metrics: PathMetrics, +} + +impl Stream for ApplyStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Split the borrow once: the timer holds `baseline` while the loop + // mutates `appliers` and reads `path_metrics`. + let this = self.get_mut(); + let poll = match ready!(this.input.poll_next_unpin(cx)) { + Some(Ok(mut batch)) => { + let rows = batch.num_rows(); + let _timer = this.baseline.elapsed_compute().timer(); + for applier in &mut this.appliers { + // Charge each batch to the path that handled it, so the + // two stay separable when one query mixes them. + let path = match applier { + PreparedApply::Scalar(_) => &this.path_metrics.scalar_time, + PreparedApply::Aggregate(_) => &this.path_metrics.aggregate_time, + }; + let timer = path.timer(); + let applied = applier.apply(batch); + timer.done(); + batch = match applied { + Ok(b) => b, + Err(e) => return Poll::Ready(Some(Err(e))), + }; + } + if !this.appliers.is_empty() { + this.path_metrics.rows_corrected.add(rows); + } + Poll::Ready(Some(Ok(batch))) + } + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => Poll::Ready(None), + }; + this.baseline.record_poll(poll) + } +} + +impl RecordBatchStream for ApplyStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Rebuild a batch with one column replaced, naming the apply responsible if +/// the new column's type no longer matches the schema. +/// +/// Arrow's own error reports a type mismatch at a column index and nothing +/// about which correction produced it. An apply can drift from its column by +/// promotion in a scalar kernel, or by an accumulator's `evaluate()` widening +/// relative to the column it overwrites — both of which are about the apply, +/// not the batch. +/// +/// # Arguments +/// +/// * `schema` - the batch's schema, unchanged by the rewrite +/// * `columns` - columns with `output_column` already replaced +/// * `apply_index` - position in the exec's `applies` list, for the message +/// * `output_column` - the column that was replaced, for the message +fn rebuild_batch( + schema: SchemaRef, + columns: Vec, + apply_index: usize, + output_column: usize, +) -> Result { + let replaced = columns[output_column].data_type().clone(); + RecordBatch::try_new(Arc::clone(&schema), columns).map_err(|e| { + datafusion::common::DataFusionError::Internal(format!( + "PrefixMergeExec: applies[{apply_index}] produced {replaced} for column \ + {output_column}, which the schema declares as {}: {e}", + schema.field(output_column).data_type() + )) + }) +} + +#[cfg(test)] +mod tests { + /// Metrics must separate the two apply paths. A single total would hide + /// that a sketch-heavy query pays the accumulator replay while a + /// SUM-heavy one need not, which is the thing worth knowing. + #[tokio::test] + async fn metrics_split_time_by_apply_path() -> Result<()> { + use datafusion::physical_plan::common::collect; + + let exec = Arc::new(PrefixMergeExec::try_new_resolved( + partitioned_source(1, 3), + vec![WindowApply::Scalar { + op: ScalarOp::Add, + offset: vec![ScalarValue::Int64(Some(10))], + output_column: 0, + }], + empty_state(1), + )?); + + let ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let _ = collect(exec.execute(0, ctx)?).await?; + + let metrics = exec.metrics().expect("operator must report metrics"); + let named = |name: &str| { + metrics + .iter() + .find(|m| m.value().name() == name) + .map(|m| m.value().as_usize()) + }; + assert_eq!( + named("rows_corrected"), + Some(3), + "every row passing a non-empty applies list is corrected" + ); + assert!( + named("scalar_apply_time").is_some(), + "the scalar path ran and must be timed separately: {metrics:?}" + ); + assert!( + named("aggregate_apply_time").is_some_and(|t| t == 0), + "the aggregate path did not run and must read zero, not be absent: \ + {metrics:?}" + ); + Ok(()) + } + + /// A drifted apply must name itself. Arrow reports a type mismatch at a + /// column index and nothing about which correction caused it, which is + /// the wrong half of the story when several applies rewrite one batch. + #[test] + fn rebuild_batch_names_the_drifted_apply() { + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("keep", DataType::Int64, false), + Field::new("corrected", DataType::Float64, false), + ])); + // applies[3] wrote an Int64 over a Float64 column. + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(Int64Array::from(vec![2])), + ]; + let err = rebuild_batch(schema, columns, 3, 1) + .expect_err("type drift must be rejected"); + let message = err.to_string(); + for expected in ["applies[3]", "column 1", "Int64", "Float64"] { + assert!( + message.contains(expected), + "error must mention {expected}: {message}" + ); + } + } + + 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))) + } + + fn empty_state(partitions: usize) -> Vec { + (0..partitions) + .map(|_| FinalizedPartitionState::default()) + .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_resolved( + input, + vec![], + vec![FinalizedPartitionState::default()], + ) + .expect_err("length mismatch must surface as an error"); + assert!( + err.to_string() + .contains("does not match input partition count"), + "unexpected error: {err}" + ); + } + + /// 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_resolved(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_resolved(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}" + ); + } + + /// 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], + )?; + // 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![ + FinalizedPartitionState::default(), + FinalizedPartitionState::new(vec![Some(p0_state)]), + ]; + + 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_resolved( + 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(()) + } + + /// 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()? + }; + // 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![ + FinalizedPartitionState::default(), + FinalizedPartitionState::new(vec![Some(p0_state)]), + ]; + + 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_resolved( + 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 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 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(0)), ScalarValue::Int64(Some(6))], + output_column: 1, + }; + let exec = Arc::new(PrefixMergeExec::try_new_resolved( + input, + vec![apply], + empty_state(2), + )?); + + let ctx = SessionContext::new().task_ctx(); + + 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(()) + } +} diff --git a/ballista/core/src/execution_plans/range_filter.rs b/ballista/core/src/execution_plans/range_filter.rs index 7df912cae..054c23571 100644 --- a/ballista/core/src/execution_plans/range_filter.rs +++ b/ballista/core/src/execution_plans/range_filter.rs @@ -85,6 +85,10 @@ use datafusion::scalar::ScalarValue; use futures::{Stream, StreamExt, ready}; use parking_lot::Mutex; +use crate::execution_plans::plan_algebra::{ + PartitionSliceable, slice_by_global_partition, +}; + /// Half-open `[lo, hi)` bound for one input partition. `None` on either side /// means unbounded (virtual ±∞). pub type RangeBound = (Option, Option); @@ -295,6 +299,35 @@ impl DisplayAs for RangeFilterExec { } } +impl PartitionSliceable for RangeFilterExec { + /// `raw_bounds` is indexed by input partition, so it slices parallel to + /// the input. Halos and routing carry over verbatim — the fresh operator + /// re-widens the unwidened bounds itself. + fn slice_to_partitions( + &self, + child: Arc, + partitions: &[usize], + ) -> Result> { + let raw_bounds = self.raw_bounds().ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "RangeFilterExec: task-restriction before resolve_bounds()".into(), + ) + })?; + Ok(Arc::new(Self::try_new_resolved( + child, + self.routing_expr().clone(), + self.halo_lo().clone(), + self.halo_hi().clone(), + slice_by_global_partition( + &raw_bounds, + partitions, + "RangeFilterExec", + "raw bounds", + )?, + )?)) + } +} + impl ExecutionPlan for RangeFilterExec { fn name(&self) -> &str { "RangeFilterExec" diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 1cabe6e14..6060e0c3a 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -29,8 +29,8 @@ use std::time::Instant; use crate::JobId; use crate::error::BallistaError; use crate::execution_plans::{ - OrderedRangeRepartitionExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, - create_shuffle_path, + ObservedWindowState, OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, + SortShuffleWriterExec, UnorderedRangeRepartitionExec, create_shuffle_path, }; use crate::extension::SessionConfigExt; use crate::utils; @@ -497,6 +497,46 @@ impl ShuffleWriterExec { &self.global_output_partition_ids } + /// Drain every window-state collector in this stage, translating each + /// capture's task-local partition index to its global one. + /// + /// The translation belongs here rather than at the operator that captured + /// the state: a task's plan is restricted to a partition slice, so an + /// operator mid-plan only ever sees local indices, while this writer is + /// handed `global_output_partition_ids` by the scheduler. Reassembling the + /// order downstream instead would mean the scheduler re-deriving a mapping + /// it already computed, and a prefix scan fed a permuted order is wrong + /// with no error to show for it. + /// + /// Errors rather than dropping a capture it cannot place. Unlike runtime + /// stats, which are an optimization input, this state is load-bearing: + /// the downstream stage's prefix merge is arithmetically wrong without + /// every partition's contribution, and wrong in a way no later check + /// catches. Failing the task surfaces it while it is still a failure + /// rather than a wrong answer. + pub fn collect_window_state(&self) -> Result> { + let mut found: Vec<&PartitionedBoundedWindowAggExec> = Vec::new(); + collect_window_state_operators(&self.plan, &mut found); + found + .into_iter() + .flat_map(|op| op.observed_window_state()) + .map(|observation| { + let global = self + .global_output_partition_ids + .get(observation.partition_idx) + .copied() + .ok_or_else(|| { + DataFusionError::Internal(format!( + "ShuffleWriterExec: window state for local partition {} \ + has no global id (slice covers {:?})", + observation.partition_idx, self.global_output_partition_ids + )) + })?; + Ok((global, observation)) + }) + .collect() + } + /// Get the Job ID for this query stage pub fn job_id(&self) -> &JobId { &self.job_id @@ -951,6 +991,28 @@ pub(crate) fn summaries_to_batch( MemoryStream::try_new(vec![batch], schema, None) } +/// Collect every [`PartitionedBoundedWindowAggExec`] reachable from `plan`. +/// +/// Walks the whole subtree rather than a partition-preserving spine: an +/// operator holding window state is worth draining wherever it sits, and the +/// writer translates indices against its own slice regardless of depth. +/// +/// TODO: retarget when `PartitionedBoundedWindowAggExec` collapses. Whatever +/// ends up holding the `BoundedWindowAggExec` — a bare BWAG once DataFusion +/// can declare a non-single input distribution, or a small node planted above +/// it — is what this should look for. +fn collect_window_state_operators<'a>( + plan: &'a Arc, + out: &mut Vec<&'a PartitionedBoundedWindowAggExec>, +) { + if let Some(op) = plan.downcast_ref::() { + out.push(op); + } + for child in plan.children() { + collect_window_state_operators(child, out); + } +} + #[cfg(test)] #[allow(dead_code, unused_imports)] // clippy false positive with local imports mod tests { diff --git a/ballista/core/src/execution_plans/window_state.rs b/ballista/core/src/execution_plans/window_state.rs new file mode 100644 index 000000000..a28bbe70f --- /dev/null +++ b/ballista/core/src/execution_plans/window_state.rs @@ -0,0 +1,520 @@ +// 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. + +//! Executor-side capture of finalized window-aggregate state. +//! +//! DataFusion's `BoundedWindowAggExec` fires +//! [`WindowStateObserver::finalize_window_aggregate`] once per (output +//! partition, window expression, PARTITION BY tuple) as each group closes. +//! [`WindowStateCollector`] catches those and holds them until the executor +//! drains them at task completion. One task's captures are its contribution +//! to a cross-task prefix scan: the scheduler merges the contributions of +//! all prior partitions and bakes the result into a downstream +//! `PrefixMergeExec`. +//! +//! # Why catching is required rather than polling +//! +//! `Accumulator::state` is a destructive read — several built-in aggregates +//! `std::mem::take` their internal buffers to build it — so DataFusion fires +//! the callback at most once per group and errors on a second call. There is +//! no getter to poll afterwards. Whatever the callback hands over is either +//! retained here or gone. +//! +//! # Lifetime relative to the operator that installs it +//! +//! Nothing here is specific to `PartitionedBoundedWindowAggExec`. That +//! wrapper only exists until DataFusion's BWAG can declare a non-single +//! input distribution itself, at which point the install site moves to a bare +//! BWAG or to a small dedicated node — and this module is unaffected. Kept +//! separate for that reason: the collector outlives the operator currently +//! installing it. + +use std::collections::HashMap; +use std::fmt::{self, Debug, Formatter}; +use std::sync::{Arc, Mutex}; + +use datafusion::common::{Result, ScalarValue, internal_datafusion_err, internal_err}; +use datafusion::physical_expr::window::{PartitionKey, PlainAggregateWindowExpr}; +use datafusion::physical_plan::windows::{WindowExpr, WindowStateObserver}; +use log::debug; + +use crate::execution_plans::prefix_merge::FinalizedPartitionState; +use crate::serde::protobuf::WindowStateReport; + +/// One finalized window-aggregate state, as DataFusion reported it. +/// +/// Every dimension the callback is keyed by is preserved. Callers that only +/// support a subset — a single window expression, or no PARTITION BY — +/// enforce that themselves rather than having it flattened away here, so the +/// wire format downstream doesn't inherit today's gates. +#[derive(Debug, Clone, PartialEq)] +pub struct ObservedWindowState { + /// Output partition index of the `BoundedWindowAggExec` stream that + /// fired. Task-local: the scheduler restricted this task to a partition + /// slice, so this indexes within the slice, not globally. + pub partition_idx: usize, + /// Position of the window expression in the exec's `window_expr()` list. + /// + /// DataFusion hands the callback an `&Arc`; the collector + /// resolves it back to this index against the list it was built from. + pub window_expr_index: usize, + /// The PARTITION BY tuple that closed. Empty when the window has no + /// PARTITION BY, which is the only shape the prefix-scan rule plants + /// today. + pub partition_key: PartitionKey, + /// `Accumulator::state` for the closed group: 1 element for SUM / COUNT / + /// MIN / MAX, 2 for AVG's `(sum, count)`, 1 opaque `Binary` for + /// sketch-backed aggregates like `approx_distinct`. + pub state: Vec, +} + +/// Captures every [`ObservedWindowState`] a `BoundedWindowAggExec` publishes. +/// +/// Shared by `Arc` between the plan node that installs it and whatever drains +/// it after the task completes. Interior mutability is required because the +/// observer callback takes `&self`. +pub struct WindowStateCollector { + /// The list the callback's `&Arc` is resolved against. + /// Must hold the same `Arc`s the observed exec does — the exec clones the + /// `Vec` rather than the expressions when building its stream, so pointer + /// identity survives. + window_expr: Vec>, + observed: Mutex>, +} + +impl WindowStateCollector { + /// Build a collector resolving callbacks against `window_expr`. + pub fn new(window_expr: Vec>) -> Self { + Self { + window_expr, + observed: Mutex::new(Vec::new()), + } + } + + /// Every state captured so far, in the order DataFusion published them. + /// + /// That is close order, not partition order: a group closing mid-stream + /// precedes one closing at end-of-stream. Callers needing a specific + /// order sort on [`ObservedWindowState`]'s fields. + pub fn observed(&self) -> Vec { + self.observed + .lock() + .expect("WindowStateCollector mutex poisoned") + .clone() + } +} + +impl Debug for WindowStateCollector { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + // The expression list is already rendered by the installing + // operator's DisplayAs; only the capture count adds anything. + f.debug_struct("WindowStateCollector") + .field("window_exprs", &self.window_expr.len()) + .field("observed", &self.observed.lock().map(|o| o.len()).ok()) + .finish() + } +} + +impl WindowStateObserver for WindowStateCollector { + fn finalize_window_aggregate( + &self, + partition_idx: usize, + window_expr: &Arc, + partition_key: &PartitionKey, + state: Vec, + ) -> Result<()> { + let window_expr_index = self + .window_expr + .iter() + .position(|candidate| Arc::ptr_eq(candidate, window_expr)) + .ok_or_else(|| { + internal_datafusion_err!( + "WindowStateCollector: callback for `{}`, which is not among \ + the {} expressions this collector was built from", + window_expr.name(), + self.window_expr.len() + ) + })?; + let observation = ObservedWindowState { + partition_idx, + window_expr_index, + partition_key: partition_key.clone(), + state, + }; + // Captures are otherwise invisible until they surface as corrected + // values two stages later, so this is the cheapest place to see what + // a task actually published. + debug!( + "WindowStateCollector: partition {} expr {} key {:?} state {:?}", + observation.partition_idx, + observation.window_expr_index, + observation.partition_key, + observation.state, + ); + self.observed + .lock() + .map_err(|_| internal_datafusion_err!("WindowStateCollector mutex poisoned"))? + .push(observation); + Ok(()) + } +} + +/// Encode one capture for the wire, stamped with the **global** partition it +/// belongs to. +/// +/// The global id is supplied by the caller rather than read off the +/// observation: [`ObservedWindowState::partition_idx`] is task-local, and only +/// the stage's writer holds the slice-to-global mapping. +/// +/// # Arguments +/// +/// * `global_partition_id` - stage-global output partition for this capture +/// * `observed` - the capture itself +pub fn window_state_to_proto( + global_partition_id: usize, + observed: &ObservedWindowState, +) -> Result { + Ok(WindowStateReport { + global_partition_id: global_partition_id as u32, + window_expr_index: observed.window_expr_index as u32, + partition_key: scalars_to_proto(&observed.partition_key, "partition key")?, + state: scalars_to_proto(&observed.state, "accumulator state")?, + }) +} + +/// Decode a wire report back into `(global partition id, expr index, key, +/// state)`. Reverses [`window_state_to_proto`]. +pub fn window_state_from_proto( + proto: &WindowStateReport, +) -> Result<(usize, usize, Vec, Vec)> { + Ok(( + proto.global_partition_id as usize, + proto.window_expr_index as usize, + scalars_from_proto(&proto.partition_key, "partition key")?, + scalars_from_proto(&proto.state, "accumulator state")?, + )) +} + +fn scalars_to_proto( + scalars: &[ScalarValue], + what: &str, +) -> Result> { + scalars + .iter() + .map(datafusion_proto_common::ScalarValue::try_from) + .collect::, _>>() + .map_err(|e| internal_datafusion_err!("failed to encode {what} to proto: {e:?}")) +} + +fn scalars_from_proto( + proto: &[datafusion_proto_common::ScalarValue], + what: &str, +) -> Result> { + proto + .iter() + .map(ScalarValue::try_from) + .collect::, _>>() + .map_err(|e| { + internal_datafusion_err!("failed to decode {what} from proto: {e:?}") + }) +} + +/// A window-state report tagged with the task that produced it. +/// +/// The tag is what makes the stage's accumulated reports purgeable. Reports +/// are append-only, and a task that is reset — retried, or lost with its +/// executor — re-runs its partition slice and reports the same global +/// partitions again under a fresh `task_id`. Without dropping the original +/// attempt's entries, the stage would hold two states for one partition and +/// the prefix merge would double-count them, which is a wrong running +/// aggregate rather than a degraded one. +#[derive(Debug, Clone, PartialEq)] +pub struct TaskWindowState { + /// Producer task's `task_id` at the time it emitted the report. + pub producer_task_id: usize, + /// The report itself, already addressed by stage-global partition. + pub report: WindowStateReport, +} + +/// Prefix-scan per-partition window state into one carry-in per partition. +/// +/// `out[k]` is the merge of every partition strictly before `k`, so `out[0]` +/// is empty (nothing precedes partition 0), `out[1]` is partition 0's state, +/// `out[2]` is partitions 0 and 1 merged, and so on. That is what a +/// downstream `PrefixMergeExec` adds to each partition's local running +/// aggregate to make it global. +/// +/// Merging goes through the aggregate's own `Accumulator::merge_batch` +/// rather than any arithmetic here, which is what lets non-decomposable +/// aggregates work: two `approx_distinct` HLL sketches combine correctly +/// while two distinct *counts* could not. +/// +/// Built from the running prefix rather than from scratch each time: +/// `out[k]` is `merge(out[k-1], state[k-1])`, which the monoid's +/// associativity makes equivalent to merging every prior partition. That is +/// two merges per partition — O(K) for K partitions — where merging all +/// priors independently would be O(K²). +/// +/// A fresh accumulator per partition is still required: `Accumulator::state` +/// is a destructive read for several built-in aggregates and must not be +/// called twice, so one accumulator cannot be advanced across the scan. +/// Seeding the fresh one from the previous prefix's state is the same +/// round trip a two-phase aggregation makes, and is what DataFusion's own +/// cross-task prefix test exercises for `approx_distinct`. +/// +/// # Arguments +/// +/// * `reports` - every report the stage accumulated, in any order +/// * `window_expr` - the upstream window operator's expressions, positionally +/// matching each report's `window_expr_index` +/// * `partition_count` - K, the stage's global output partition count +pub fn prefix_merge_window_state( + reports: &[TaskWindowState], + window_expr: &[Arc], + partition_count: usize, +) -> Result> { + // (global partition, window expression) -> that group's finalized state. + let mut states: HashMap<(usize, usize), Vec> = HashMap::new(); + for tagged in reports { + let (partition, expr_index, partition_key, state) = + window_state_from_proto(&tagged.report)?; + // `FinalizedPartitionState` carries no PARTITION BY dimension, so a + // second group in one partition has nowhere to go. The prefix rule + // only plants no-PARTITION-BY windows, and a window that has one + // needs no prefix scan anyway — `BoundedWindowAggExec` asks for + // `KeyPartitioned` input, so each partition's window is already + // independent. This assertion keeps that gate honest rather than + // silently merging two groups. + if !partition_key.is_empty() { + return internal_err!( + "prefix merge: window state for partition {partition} carries a \ + PARTITION BY key {partition_key:?}; only no-PARTITION-BY windows \ + are supported" + ); + } + if partition >= partition_count { + return internal_err!( + "prefix merge: window state for partition {partition} exceeds the \ + stage's {partition_count} partitions" + ); + } + if expr_index >= window_expr.len() { + return internal_err!( + "prefix merge: window state for expression {expr_index} exceeds the \ + operator's {} expressions", + window_expr.len() + ); + } + if states.insert((partition, expr_index), state).is_some() { + // Retries are purged by producer task, so a duplicate here means + // that purge failed. Merging both would double-count. + return internal_err!( + "prefix merge: duplicate window state for partition {partition}, \ + expression {expr_index}" + ); + } + } + + let mut prefixes: Vec = Vec::with_capacity(partition_count); + for partition in 0..partition_count { + let mut per_expr = Vec::with_capacity(window_expr.len()); + for (expr_index, expr) in window_expr.iter().enumerate() { + // Nothing precedes partition 0. Beyond that, the carry-in is the + // previous carry-in merged with the previous partition's own + // state; either may be absent when a non-aggregate window + // function publishes nothing. + let carried = partition + .checked_sub(1) + .and_then(|prior| prefixes[prior].slot(expr_index)); + let preceding = partition + .checked_sub(1) + .and_then(|prior| states.get(&(prior, expr_index))); + if carried.is_none() && preceding.is_none() { + per_expr.push(None); + continue; + } + let Some(plain) = expr.as_any().downcast_ref::() + else { + return internal_err!( + "prefix merge: expression {expr_index} published state but is not \ + a plain aggregate window expression; only ever-expanding frames \ + can be prefix-merged" + ); + }; + let mut accumulator = plain.get_aggregate_expr().create_accumulator()?; + for state in [carried, preceding].into_iter().flatten() { + let arrays = state + .iter() + .map(|scalar| scalar.to_array_of_size(1)) + .collect::>>()?; + accumulator.merge_batch(&arrays)?; + } + per_expr.push(Some(accumulator.state()?)); + } + prefixes.push(FinalizedPartitionState::new(per_expr)); + } + Ok(prefixes) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::logical_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, + }; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_expr::expressions::col; + use datafusion::physical_plan::windows::create_window_expr; + + /// `sum(v) OVER (ORDER BY v ROWS UNBOUNDED PRECEDING TO CURRENT ROW)` — + /// the ever-expanding shape the prefix rule plants. + fn running_sum_expr() -> Arc { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let v = col("v", schema.as_ref()).expect("column v"); + create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum(v)".to_string(), + &[Arc::clone(&v)], + &[], + &[PhysicalSortExpr { + expr: v, + options: SortOptions::default(), + }], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + ) + .expect("window expr") + } + + fn report(global_partition_id: u32, sum: f64) -> TaskWindowState { + TaskWindowState { + producer_task_id: global_partition_id as usize, + report: window_state_to_proto( + global_partition_id as usize, + &ObservedWindowState { + partition_idx: 0, + window_expr_index: 0, + partition_key: vec![], + state: vec![ScalarValue::Float64(Some(sum))], + }, + ) + .expect("encode"), + } + } + + fn sums(prefixes: &[FinalizedPartitionState]) -> Vec> { + prefixes + .iter() + .map(|per_expr| match per_expr.slot(0) { + None => None, + Some(state) => match &state[0] { + ScalarValue::Float64(v) => *v, + other => panic!("unexpected sum state {other:?}"), + }, + }) + .collect() + } + + /// The carry-ins for the client e2e's input: four partitions holding + /// 1..4, 5..8, 9..12, 13..16, so per-partition sums 10/26/42/58 and + /// carry-ins 0 / 10 / 36 / 78. + #[test] + fn prefix_scan_accumulates_prior_partitions() -> Result<()> { + let exprs = vec![running_sum_expr()]; + let reports = vec![ + report(0, 10.0), + report(1, 26.0), + report(2, 42.0), + report(3, 58.0), + ]; + let prefixes = prefix_merge_window_state(&reports, &exprs, 4)?; + assert_eq!( + sums(&prefixes), + vec![None, Some(10.0), Some(36.0), Some(78.0)], + "partition k must carry the merge of every partition before it" + ); + Ok(()) + } + + /// Report order is close order, not partition order — a task's + /// end-of-stream group closes after a mid-stream one, and tasks complete + /// in any order. The scan must key on the reported partition, not on + /// arrival. + #[test] + fn prefix_scan_is_independent_of_report_order() -> Result<()> { + let exprs = vec![running_sum_expr()]; + let forward = prefix_merge_window_state( + &[report(0, 10.0), report(1, 26.0), report(2, 42.0)], + &exprs, + 3, + )?; + let shuffled = prefix_merge_window_state( + &[report(2, 42.0), report(0, 10.0), report(1, 26.0)], + &exprs, + 3, + )?; + assert_eq!(sums(&forward), sums(&shuffled)); + assert_eq!(sums(&forward), vec![None, Some(10.0), Some(36.0)]); + Ok(()) + } + + /// A partition that received no rows closes no group, so it publishes no + /// state. The carry-in must step over the gap rather than reset: building + /// each prefix from the previous one makes that automatic, but it is the + /// case the incremental form could plausibly get wrong. + #[test] + fn prefix_scan_carries_across_a_partition_with_no_state() -> Result<()> { + let exprs = vec![running_sum_expr()]; + // Partition 1 is empty; 0, 2 and 3 report. + let prefixes = prefix_merge_window_state( + &[report(0, 10.0), report(2, 42.0), report(3, 58.0)], + &exprs, + 4, + )?; + assert_eq!( + sums(&prefixes), + vec![None, Some(10.0), Some(10.0), Some(52.0)], + "an empty partition must neither reset the carry-in nor contribute" + ); + Ok(()) + } + + /// A duplicate means the producer-task purge failed; merging both would + /// double-count silently. + #[test] + fn prefix_scan_rejects_duplicate_partition_state() { + let exprs = vec![running_sum_expr()]; + let err = + prefix_merge_window_state(&[report(0, 10.0), report(0, 10.0)], &exprs, 2) + .expect_err("duplicate must be rejected"); + assert!( + err.to_string().contains("duplicate window state"), + "unexpected error: {err}" + ); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index a868ff2b3..5261fc8c4 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -67,6 +67,8 @@ pub mod ballista_physical_plan_node { RangeShuffleReader(super::RangeShuffleReaderExecNode), #[prost(message, tag = "13")] RangeFilter(super::RangeFilterExecNode), + #[prost(message, tag = "14")] + PrefixMerge(super::PrefixMergeExecNode), } } /// Value-range router over N locally-sorted overlapping input partitions. @@ -182,6 +184,83 @@ pub struct RangeBound { #[prost(message, optional, tag = "2")] pub hi: ::core::option::Option<::datafusion_proto_common::ScalarValue>, } +/// Applies scheduler-computed prefix state to a window-aggregate column, so +/// each range-disjoint task's local running aggregate becomes a global one. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PrefixMergeExecNode { + /// One entry per output column needing correction. Empty makes the operator + /// a passthrough. + #[prost(message, repeated, tag = "1")] + pub applies: ::prost::alloc::vec::Vec, + /// Resolved prefix state, one entry per input partition, `\[k\]` summarising + /// every partition before `k`. Encoding refuses while unresolved — an + /// over-the-wire plan always ships with state bound, since an executor has + /// no way to obtain it. + #[prost(message, repeated, tag = "2")] + pub per_partition_state: ::prost::alloc::vec::Vec, +} +/// How to correct one window-function output column. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct WindowApplyNode { + #[prost(oneof = "window_apply_node::Apply", tags = "1, 2")] + pub apply: ::core::option::Option, +} +/// Nested message and enum types in `WindowApplyNode`. +pub mod window_apply_node { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Apply { + #[prost(message, tag = "1")] + Scalar(super::ScalarWindowApplyNode), + #[prost(message, tag = "2")] + Aggregate(super::AggregateWindowApplyNode), + } +} +/// Fast path: combine each row's value with a per-partition scalar. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ScalarWindowApplyNode { + #[prost(enumeration = "ScalarOpNode", tag = "1")] + pub op: i32, + /// One scalar per input partition. + #[prost(message, repeated, tag = "2")] + pub offset: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + #[prost(uint32, tag = "3")] + pub output_column: u32, +} +/// Fallback path: seed a fresh accumulator from the prefix state and replay +/// each row through it. Needed where a row's output is not itself a valid +/// partial state — AVG, and sketch-backed aggregates like approx_distinct. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AggregateWindowApplyNode { + /// Resolved from the executor's function registry on decode. + #[prost(string, tag = "1")] + pub udf_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub args: ::prost::alloc::vec::Vec<::datafusion_proto::protobuf::PhysicalExprNode>, + #[prost(uint32, tag = "3")] + pub output_column: u32, + /// Position in the upstream window operator's `window_expr()` list, which + /// is what indexes into each FinalizedPartitionStateNode. + #[prost(uint32, tag = "4")] + pub window_expr_index: u32, +} +/// Prefix state for one input partition: one slot per window expression. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct FinalizedPartitionStateNode { + #[prost(message, repeated, tag = "1")] + pub slots: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AggregateStateSlotNode { + /// Unset when that window expression published no state — a non-aggregate + /// window function. Distinct from a present-but-empty state. + #[prost(message, optional, tag = "1")] + pub state: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AggregateStateNode { + #[prost(message, repeated, tag = "1")] + pub values: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, +} /// Wrapper for `BoundedWindowAggExec` that overrides /// `required_input_distribution` to `Unspecified` — see the module doc on /// `execution_plans::partitioned_bounded_window_agg` for what makes that safe. @@ -960,6 +1039,53 @@ pub struct SuccessfulTask { /// to combine reports across tasks/executors. #[prost(message, repeated, tag = "3")] pub runtime_stats: ::prost::alloc::vec::Vec, + /// Finalized window-aggregate state captured during this task, one entry + /// per (output partition, window expression, PARTITION BY group) that + /// closed. Empty unless the plan contains an ever-expanding-frame window + /// (`UNBOUNDED PRECEDING`), which is the only shape DataFusion will + /// publish accumulator state for. The scheduler prefix-merges these across + /// tasks and bakes the result into a downstream `PrefixMergeExec`. + /// + /// TODO: watch the size of this. Task completion is a hot, frequent + /// message, and sketch-backed aggregates make the payload unbounded in a + /// way row counts and quantile sketches are not — an HLL or KLL state is + /// kilobytes per window expression per partition, and a task covering a + /// wide partition slice carries one of each. If it stops being small, + /// write the state as a sidecar next to the shuffle files instead, the way + /// sort-shuffle already writes `.arrow.index` beside its data + /// (`sort_shuffle::get_index_path`), and send only a reference here. That + /// keeps the completion message fixed-size regardless of aggregate. + #[prost(message, repeated, tag = "4")] + pub window_state: ::prost::alloc::vec::Vec, +} +/// One finalized window-aggregate state from a task's +/// `BoundedWindowAggExec`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct WindowStateReport { + /// The stage's *global* output partition this state belongs to. + /// + /// DataFusion reports a task-local index, because a task's plan is + /// restricted to a partition slice. The producing `ShuffleWriterExec` + /// translates it through the `global_output_partition_ids` the scheduler + /// stamped on it, so what crosses the wire is already global — the + /// scheduler must not re-derive it. + #[prost(uint32, tag = "1")] + pub global_partition_id: u32, + /// Position in the window operator's `window_expr()` list. Indexes window + /// expressions, not partitions. + #[prost(uint32, tag = "2")] + pub window_expr_index: u32, + /// The PARTITION BY tuple that closed. Empty for a window with no + /// PARTITION BY, which is the only shape the prefix rewrite plants today. + #[prost(message, repeated, tag = "3")] + pub partition_key: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + /// `Accumulator::state` for the closed group: one element for SUM / COUNT / + /// MIN / MAX, two for AVG's (sum, count), one opaque Binary for + /// sketch-backed aggregates like approx_distinct. Carried as ScalarValue + /// rather than a numeric field so non-decomposable aggregates work + /// unchanged. + #[prost(message, repeated, tag = "4")] + pub state: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, } /// One report per `RuntimeStatsExec` in the executed plan. #[derive(Clone, PartialEq, ::prost::Message)] @@ -1546,6 +1672,40 @@ impl BufferMode { } } } +/// How a scalar offset combines with a row's existing value. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum ScalarOpNode { + Add = 0, + Min = 1, + Max = 2, + /// Ignores the row's value entirely. + Overwrite = 3, +} +impl ScalarOpNode { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Add => "SCALAR_OP_NODE_ADD", + Self::Min => "SCALAR_OP_NODE_MIN", + Self::Max => "SCALAR_OP_NODE_MAX", + Self::Overwrite => "SCALAR_OP_NODE_OVERWRITE", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "SCALAR_OP_NODE_ADD" => Some(Self::Add), + "SCALAR_OP_NODE_MIN" => Some(Self::Min), + "SCALAR_OP_NODE_MAX" => Some(Self::Max), + "SCALAR_OP_NODE_OVERWRITE" => Some(Self::Overwrite), + _ => None, + } + } +} /// Generated client implementations. pub mod scheduler_grpc_client { #![allow( diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 642feec2c..d9048279e 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -56,17 +56,19 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ - BufferExec, BufferMode, ChaosExec, CoalescePlan, OrderedRangeRepartitionExec, - PartitionGroup, PartitionedBoundedWindowAggExec, PerPartitionFilterExec, - RangeFilterExec, RangeShuffleReaderExec, RuntimeStatsExec, ShuffleReaderExec, - ShuffleWriterExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, - UnresolvedShuffleExec, + BufferExec, BufferMode, ChaosExec, CoalescePlan, FinalizedPartitionState, + OrderedRangeRepartitionExec, PartitionGroup, PartitionedBoundedWindowAggExec, + PerPartitionFilterExec, PrefixMergeExec, RangeFilterExec, RangeShuffleReaderExec, + RuntimeStatsExec, ScalarOp, ShuffleReaderExec, ShuffleWriterExec, + SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec, + WindowApply, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, ballista_physical_plan_node::PhysicalPlanType, }; use crate::serde::scheduler::PartitionLocation; +use datafusion::execution::FunctionRegistry; pub use generated::ballista as protobuf; /// Generated protobuf code from Ballista protocol definitions. @@ -366,6 +368,191 @@ impl Default for BallistaPhysicalExtensionCodec { } } +/// Encode a `ScalarValue` list for the wire, naming what failed. +fn encode_scalars( + values: &[datafusion::scalar::ScalarValue], + what: &str, +) -> Result, DataFusionError> { + let mut out = Vec::with_capacity(values.len()); + for value in values { + out.push( + datafusion_proto_common::ScalarValue::try_from(value).map_err(|e| { + DataFusionError::Internal(format!("failed to encode {what}: {e:?}")) + })?, + ); + } + Ok(out) +} + +/// Reverse of [`encode_scalars`]. +fn decode_scalars( + proto: &[datafusion_proto_common::ScalarValue], + what: &str, +) -> Result, DataFusionError> { + let mut out = Vec::with_capacity(proto.len()); + for value in proto { + out.push( + datafusion::scalar::ScalarValue::try_from(value).map_err(|e| { + DataFusionError::Internal(format!("failed to decode {what}: {e:?}")) + })?, + ); + } + Ok(out) +} + +/// Encode `PrefixMergeExec`'s per-partition prefix state. +/// +/// An absent slot stays absent: a window expression that published no state +/// is not the same as one that published an empty state. +fn encode_prefix_state( + state: &[FinalizedPartitionState], +) -> Result, DataFusionError> { + let mut partitions = Vec::with_capacity(state.len()); + for partition in state { + let mut slots = Vec::with_capacity(partition.len()); + for slot in partition.slots() { + let state = match slot { + Some(values) => Some(protobuf::AggregateStateNode { + values: encode_scalars(values, "prefix state")?, + }), + None => None, + }; + slots.push(protobuf::AggregateStateSlotNode { state }); + } + partitions.push(protobuf::FinalizedPartitionStateNode { slots }); + } + Ok(partitions) +} + +/// Reverse of [`encode_prefix_state`]. +fn decode_prefix_state( + proto: &[protobuf::FinalizedPartitionStateNode], +) -> Result, DataFusionError> { + let mut partitions = Vec::with_capacity(proto.len()); + for partition in proto { + let mut slots = Vec::with_capacity(partition.slots.len()); + for slot in &partition.slots { + let state = match &slot.state { + Some(state) => Some(decode_scalars(&state.values, "prefix state")?), + None => None, + }; + slots.push(state); + } + partitions.push(FinalizedPartitionState::new(slots)); + } + Ok(partitions) +} + +/// Encode `PrefixMergeExec`'s per-column apply descriptors. +/// +/// The aggregate arm carries its UDAF by name, resolved from the executor's +/// function registry on decode — the same way DataFusion moves UDFs +/// generally. +fn encode_window_applies( + applies: &[WindowApply], + codec: &dyn PhysicalExtensionCodec, +) -> Result, DataFusionError> { + let mut out = Vec::with_capacity(applies.len()); + for apply in applies { + let apply = match apply { + WindowApply::Scalar { + op, + offset, + output_column, + } => protobuf::window_apply_node::Apply::Scalar( + protobuf::ScalarWindowApplyNode { + op: encode_scalar_op(*op) as i32, + offset: encode_scalars(offset, "scalar apply offset")?, + output_column: *output_column as u32, + }, + ), + WindowApply::Aggregate { + udf, + args, + output_column, + window_expr_index, + } => { + let mut encoded_args = Vec::with_capacity(args.len()); + for arg in args { + encoded_args.push( + datafusion_proto::physical_plan::to_proto::serialize_physical_expr( + arg, codec, + )?, + ); + } + protobuf::window_apply_node::Apply::Aggregate( + protobuf::AggregateWindowApplyNode { + udf_name: udf.name().to_string(), + args: encoded_args, + output_column: *output_column as u32, + window_expr_index: *window_expr_index as u32, + }, + ) + } + }; + out.push(protobuf::WindowApplyNode { apply: Some(apply) }); + } + Ok(out) +} + +/// Reverse of [`encode_window_applies`]. `schema` is the operator's input +/// schema, which the aggregate arm's argument expressions resolve against. +fn decode_window_applies( + proto: &[protobuf::WindowApplyNode], + ctx: &TaskContext, + schema: &datafusion::arrow::datatypes::Schema, + codec: &dyn PhysicalExtensionCodec, +) -> Result, DataFusionError> { + let mut out = Vec::with_capacity(proto.len()); + for node in proto { + let apply = node.apply.as_ref().ok_or_else(|| { + DataFusionError::Internal("WindowApplyNode missing apply".into()) + })?; + out.push(match apply { + protobuf::window_apply_node::Apply::Scalar(scalar) => WindowApply::Scalar { + op: decode_scalar_op(scalar.op()), + offset: decode_scalars(&scalar.offset, "scalar apply offset")?, + output_column: scalar.output_column as usize, + }, + protobuf::window_apply_node::Apply::Aggregate(aggregate) => { + let mut args = Vec::with_capacity(aggregate.args.len()); + for arg in &aggregate.args { + args.push( + datafusion_proto::physical_plan::from_proto::parse_physical_expr( + arg, ctx, schema, codec, + )?, + ); + } + WindowApply::Aggregate { + udf: ctx.udaf(&aggregate.udf_name)?, + args, + output_column: aggregate.output_column as usize, + window_expr_index: aggregate.window_expr_index as usize, + } + } + }); + } + Ok(out) +} + +fn encode_scalar_op(op: ScalarOp) -> protobuf::ScalarOpNode { + match op { + ScalarOp::Add => protobuf::ScalarOpNode::Add, + ScalarOp::Min => protobuf::ScalarOpNode::Min, + ScalarOp::Max => protobuf::ScalarOpNode::Max, + ScalarOp::Overwrite => protobuf::ScalarOpNode::Overwrite, + } +} + +fn decode_scalar_op(op: protobuf::ScalarOpNode) -> ScalarOp { + match op { + protobuf::ScalarOpNode::Add => ScalarOp::Add, + protobuf::ScalarOpNode::Min => ScalarOp::Min, + protobuf::ScalarOpNode::Max => ScalarOp::Max, + protobuf::ScalarOpNode::Overwrite => ScalarOp::Overwrite, + } +} + impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { fn try_decode( &self, @@ -704,6 +891,20 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { predicates, )?)) } + PhysicalPlanType::PrefixMerge(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "PrefixMergeExec expects exactly 1 input, got {}", + inputs.len() + ))); + }; + let schema = input.schema(); + Ok(Arc::new(PrefixMergeExec::try_new_resolved( + input.clone(), + decode_window_applies(&node.applies, ctx, schema.as_ref(), self)?, + decode_prefix_state(&node.per_partition_state)?, + )?)) + } PhysicalPlanType::RangeFilter(node) => { let [input] = inputs else { return Err(DataFusionError::Internal(format!( @@ -1086,6 +1287,32 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + // An executor has no route to prefix state, so a plan reaching + // the wire unresolved could only emit partition-local aggregates. + // Refuse, the way RangeFilterExec refuses unresolved bounds. + let state = exec.per_partition_state().ok_or_else(|| { + DataFusionError::Internal( + "PrefixMergeExec: cannot serialize before resolve_state()".into(), + ) + })?; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::PrefixMerge( + protobuf::PrefixMergeExecNode { + applies: encode_window_applies( + exec.applies(), + self.default_codec.as_ref(), + )?, + per_partition_state: encode_prefix_state(&state)?, + }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode PrefixMergeExecNode: {e:?}" + )) + })?; + Ok(()) } else if let Some(exec) = node.downcast_ref::() { let raw_bounds = exec.raw_bounds().ok_or_else(|| { DataFusionError::Internal( @@ -1409,6 +1636,92 @@ mod test { /// `RangeShuffleReaderExec` carries its merge ordering across the wire — /// the reader's k-way merge machinery is inert without it, and the /// scheduler side plants the reader with the child's declared ordering. + /// Both `WindowApply` shapes and the prefix state must survive the wire. + /// The aggregate arm is the interesting one: its UDAF crosses by name and + /// is resolved from the executor's registry, and its state is an opaque + /// `Vec` rather than a number. + #[tokio::test] + async fn test_prefix_merge_exec_roundtrip() { + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::scalar::ScalarValue; + + let schema = create_test_schema(); + let input = Arc::new(EmptyExec::new(schema.clone())); + let applies = vec![ + WindowApply::Scalar { + op: ScalarOp::Overwrite, + offset: vec![ScalarValue::Float64(Some(1.5))], + output_column: 0, + }, + WindowApply::Aggregate { + udf: sum_udaf(), + args: vec![col("id", schema.as_ref()).unwrap()], + output_column: 1, + window_expr_index: 3, + }, + ]; + // One partition; slot 0 carries state, slot 1 is a window function + // that published none. `None` and `Some(vec![])` must stay distinct. + let state = vec![FinalizedPartitionState::new(vec![ + Some(vec![ScalarValue::Float64(Some(42.0))]), + None, + ])]; + let original = + PrefixMergeExec::try_new_resolved(input, applies, state.clone()).unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec + .try_encode( + Arc::new(original), + &mut buf, + &DefaultPhysicalProtoConverter {}, + ) + .unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec + .try_decode( + &buf, + &[Arc::new(EmptyExec::new(schema.clone()))], + &ctx, + &DefaultPhysicalProtoConverter {}, + ) + .unwrap(); + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected PrefixMergeExec"); + + assert_eq!(decoded.per_partition_state(), Some(state)); + match &decoded.applies()[0] { + WindowApply::Scalar { + op, + offset, + output_column, + } => { + assert!(matches!(op, ScalarOp::Overwrite)); + assert_eq!(offset, &[ScalarValue::Float64(Some(1.5))]); + assert_eq!(*output_column, 0); + } + other => panic!("expected Scalar apply, got {other:?}"), + } + match &decoded.applies()[1] { + WindowApply::Aggregate { + udf, + args, + output_column, + window_expr_index, + } => { + assert_eq!(udf.name(), "sum"); + assert_eq!(args.len(), 1); + assert_eq!(*output_column, 1); + assert_eq!(*window_expr_index, 3); + } + other => panic!("expected Aggregate apply, got {other:?}"), + } + } + #[tokio::test] async fn test_range_shuffle_reader_exec_roundtrip() { use datafusion::arrow::compute::SortOptions; diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index e673a51c8..fe2a9d278 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -100,6 +100,26 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display { ) -> Vec { Vec::new() } + + /// Drain finalized window-aggregate state captured during this task, + /// already stamped with the global partition each capture belongs to. + /// Called at task completion, like + /// [`Self::collect_runtime_stats_reports`], and rides the same + /// `SuccessfulTask` message. + /// + /// Empty for every plan without an ever-expanding-frame window, which is + /// nearly all of them. Default returns empty — implementers with real + /// plans override to drain their writer. + /// + /// Errors fail the task. This state is load-bearing for the downstream + /// stage's prefix merge, so losing a report yields a wrong answer rather + /// than a degraded one — unlike + /// [`Self::collect_runtime_stats_reports`], which is telemetry. + fn collect_window_state_reports( + &self, + ) -> Result> { + Ok(Vec::new()) + } } /// Default execution engine using DataFusion's ShuffleWriterExec. @@ -340,6 +360,27 @@ impl QueryStageExecutor for DefaultQueryStageExec { } } } + + fn collect_window_state_reports( + &self, + ) -> Result> { + // Only the passthrough writer can sit over a window: the sort writer + // is Hash-partitioned by construction, which the prefix rewrite never + // plants. + let captured = match &self.shuffle_writer { + ShuffleWriterVariant::Passthrough(writer) => writer.collect_window_state()?, + ShuffleWriterVariant::Sort(_) => return Ok(Vec::new()), + }; + captured + .iter() + .map(|(global_partition, observed)| { + ballista_core::execution_plans::window_state_to_proto( + *global_partition, + observed, + ) + }) + .collect() + } } /// Spawn K parallel `plan.execute(N, ctx)` calls against a shuffle writer, diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index 2390364c1..a439dff03 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -419,6 +419,18 @@ async fn run_received_task, BallistaError>>() .ok(); let runtime_stats = query_stage_exec.collect_runtime_stats_reports(); + // Collect only when the task otherwise succeeded: a failed task's + // partial state is meaningless, and its own error is the useful one. + // A collection failure fails the task — these are load-bearing for the + // downstream stage's prefix merge, so continuing without them would + // ship a wrong answer that nothing later detects. + let (execution_result, window_state) = match execution_result { + Ok(partitions) => match query_stage_exec.collect_window_state_reports() { + Ok(reports) => (Ok(partitions), reports), + Err(e) => (Err(e.into()), Vec::new()), + }, + Err(e) => (Err(e), Vec::new()), + }; let end_exec_time = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -440,6 +452,7 @@ async fn run_received_task ExecutorServer, BallistaError>>() .ok(); let runtime_stats = exec.collect_runtime_stats_reports(); + // Collect only when the task otherwise succeeded: a failed task's + // partial state is meaningless, and its own error is the useful one. + // A collection failure fails the task — these are load-bearing for the + // downstream stage's prefix merge, so continuing without them would + // ship a wrong answer that nothing later detects. + let (execution_result, window_state) = match execution_result { + Ok(partitions) => match exec.collect_window_state_reports() { + Ok(reports) => (Ok(partitions), reports), + Err(e) => (Err(e.into()), Vec::new()), + }, + Err(e) => (Err(e), Vec::new()), + }; let executor_id = &self.executor.metadata.id; let end_exec_time = SystemTime::now() @@ -481,6 +493,7 @@ impl ExecutorServer>, /// Runtime-stats reports harvested from `RuntimeStatsExec` taps in the plan. pub runtime_stats: Vec, + /// Finalized window-aggregate state captured by an ever-expanding-frame + /// window, already stamped with the global partition each entry belongs + /// to by the stage's `ShuffleWriterExec`. + pub window_state: Vec, } /// Converts a task execution result into a [`TaskStatus`] protobuf message. @@ -129,6 +133,7 @@ pub fn as_task_status( let TaskCompletionExtras { operator_metrics, runtime_stats, + window_state, } = extras; let metrics = operator_metrics.unwrap_or_default(); let task_id = key.task_id; @@ -136,9 +141,10 @@ pub fn as_task_status( Ok(partitions) => { debug!( "Task {task_id} finished with operator_metrics array size {} \ - and {} runtime-stats report(s)", + and {} runtime-stats report(s), {} window-state report(s)", metrics.len(), runtime_stats.len(), + window_state.len(), ); TaskStatus { task_id: task_id as u32, @@ -153,6 +159,7 @@ pub fn as_task_status( executor_id, partitions, runtime_stats, + window_state, })), } } diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 5ee9d5fbd..85ae0bcc2 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -686,6 +686,7 @@ mod test { executor_id: "executor-1".to_owned(), partitions, runtime_stats: vec![], + window_state: vec![], })), }; diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index e11463924..17b0d11e9 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -867,6 +867,7 @@ mod tests { }) .collect(), runtime_stats: vec![], + window_state: vec![], })), }) .collect(); diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index d06b19d3e..58e7db120 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -29,8 +29,8 @@ use crate::state::task_manager::UpdatedStages; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{ - RangeFilterExec, cut_partitions, merge_runtime_stats_reports, - repartition_routing_expr, + PartitionedBoundedWindowAggExec, PrefixMergeExec, RangeFilterExec, cut_partitions, + merge_runtime_stats_reports, prefix_merge_window_state, repartition_routing_expr, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; @@ -39,10 +39,12 @@ use ballista_core::serde::protobuf::{ job_status, task_status, }; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; +use datafusion::common::DataFusionError; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::windows::WindowExpr; use datafusion::prelude::SessionConfig; use datafusion::scalar::ScalarValue; use log::{debug, error, info, warn}; @@ -355,6 +357,65 @@ impl AdaptiveExecutionGraph { })) } + /// Prefix-merge the window state a completed stage reported, and bind the + /// result to the `PrefixMergeExec` waiting on it downstream. + /// + /// No-op for the overwhelming majority of stages, which report no window + /// state at all. When a stage did report, failing to find its consumer is + /// an error rather than a skip: the state exists precisely because a + /// `PrefixMergeExec` downstream cannot produce correct values without it. + fn resolve_prefix_merge_state( + &self, + stage_id: usize, + ) -> ballista_core::error::Result<()> { + let Some(ExecutionStage::Running(stage)) = self.stages.get(&stage_id) else { + return Ok(()); + }; + if stage.window_state_reports.is_empty() { + return Ok(()); + } + let reports = stage.window_state_reports.clone(); + let mut resolved = 0usize; + self.planner + .plan + .apply(|node| { + let Some(prefix_merge) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + // Each PrefixMergeExec consumes exactly one upstream stage — + // the one behind the state-sync boundary directly below it. + if descend_to_boundary_stage_id(prefix_merge.input()) != Some(stage_id) { + return Ok(TreeNodeRecursion::Continue); + } + let window_expr = descend_to_window_expr(prefix_merge.input()) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "PrefixMergeExec over stage {stage_id} has no window \ + operator below it to interpret its state against" + )) + })?; + let partition_count = prefix_merge + .input() + .properties() + .output_partitioning() + .partition_count(); + let prefixes = + prefix_merge_window_state(&reports, &window_expr, partition_count)?; + prefix_merge.resolve_state(prefixes)?; + resolved += 1; + Ok(TreeNodeRecursion::Continue) + }) + .map_err(|e| BallistaError::General(e.to_string()))?; + if resolved == 0 { + return Err(BallistaError::General(format!( + "stage {stage_id} reported {} window-state entries but no \ + PrefixMergeExec consumes them; the correction would be dropped", + reports.len() + ))); + } + Ok(()) + } + /// Return a Vec of stages to cancel fn update_stage_progress( &mut self, @@ -366,6 +427,7 @@ impl AdaptiveExecutionGraph { .update_exchange_locations(stage_id, locations)?; if is_completed { + self.resolve_prefix_merge_state(stage_id)?; let partitions = self.planner.take_stage_output_partitions(stage_id)?; // Range-repartition stages need overlap-based remap of their partitions @@ -900,6 +962,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph { let ballista_core::serde::protobuf::SuccessfulTask { partitions, runtime_stats, + window_state, .. } = successful_task; debug!( @@ -911,6 +974,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { ); running_stage .append_runtime_stats_reports(task_id, runtime_stats); + running_stage + .append_window_state_reports(task_id, window_state); locations.append( &mut crate::state::execution_graph::partition_to_location( @@ -1536,3 +1601,48 @@ macro_rules! assert_plan { insta::assert_snapshot!(actual_lines, @ $EXPECTED_LINES); }; } + +/// Descend the single-child spine below a `PrefixMergeExec` to its +/// state-sync `ExchangeExec` and return that boundary's stage id. +/// +/// `None` when the boundary has no stage id yet (the upstream stage has not +/// been assigned one), or when the spine forks before reaching it — a +/// prefix merge only pairs with a single upstream stage. +fn descend_to_boundary_stage_id(start: &Arc) -> Option { + let mut node = Arc::clone(start); + loop { + if let Some(exchange) = node.downcast_ref::() { + return exchange.stage_id(); + } + let children = node.children(); + let [child] = children.as_slice() else { + return None; + }; + node = Arc::clone(*child); + } +} + +/// Find the window operator below a `PrefixMergeExec` and return its window +/// expressions, which the reported state is indexed against. +/// +/// Walks the whole subtree rather than the spine: the operator sits below the +/// state-sync boundary, whose input subtree the exchange retains even once +/// resolved. +/// +/// TODO: retarget when `PartitionedBoundedWindowAggExec` collapses, alongside +/// the collector's own walk in `shuffle_writer`. +fn descend_to_window_expr( + start: &Arc, +) -> Option>> { + let mut found = None; + start + .apply(|node| { + if let Some(op) = node.downcast_ref::() { + found = Some(op.window_expr().to_vec()); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .ok()?; + found +} diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs index 392c4bc37..631e82c0c 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs @@ -20,10 +20,12 @@ pub mod coalesce_partitions; pub mod distributed_exchange; pub mod join_selection; pub mod parallel_window; +pub mod prefix_window; pub mod propagate_empty; pub use coalesce_partitions::*; pub use distributed_exchange::*; pub use join_selection::*; pub use parallel_window::*; +pub use prefix_window::*; pub use propagate_empty::*; diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/prefix_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/prefix_window.rs new file mode 100644 index 000000000..0d826c12f --- /dev/null +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/prefix_window.rs @@ -0,0 +1,628 @@ +// 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. + +//! Rewrite ever-expanding-frame windows (`UNBOUNDED PRECEDING`) into a +//! parallel prefix scan, so `BoundedWindowAggExec`'s single-partition +//! constraint isn't a serial bottleneck. +//! +//! Sibling of the halo rewrite in [`super::parallel_window`], and mutually +//! exclusive with it: that rule gates on `is_finite(&frame.start_bound)`, +//! this one on `start_bound.is_unbounded()`. No plan can match both, so +//! their order in the optimizer chain doesn't matter. +//! +//! # Status +//! +//! Correct end to end for the shape it gates on: no PARTITION BY, a single +//! ascending `Float64` ORDER BY column, and an ever-expanding frame. +//! `prefix_window.rs` in `ballista/client/tests` holds it to the serial +//! answer through a real cluster. +//! +//! Not yet reached by h2o Q7, which orders by an `Int64` column while ORRE +//! routes on a Float64-only T-Digest — a sketch restriction that lifts with +//! the KLL migration, not anything in this rule. See `declines_int64_routing`. +//! +//! # Why halo can't cover this +//! +//! The halo rewrite widens each range partition by the frame's reach so +//! every frame's rows are local, then drops the halo rows afterwards. That +//! needs a *finite* reach. An `UNBOUNDED PRECEDING` frame reaches back to +//! the first row of the dataset, so no halo width suffices — every +//! partition would need every prior partition. The prefix scan instead lets +//! each task compute a partition-local running aggregate, then corrects it +//! with the merged state of all prior partitions. +//! +//! # Why it's worth doing +//! +//! h2o `window.sql` Q7 at scale 1e7, 8 partitions, 2 executors × 4 vcores: +//! +//! ```sql +//! SELECT id1, id2, id3, v2, +//! sum(v2) OVER (ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING +//! AND CURRENT ROW) AS my_rolling_sum +//! FROM large; +//! ``` +//! +//! Stage 0 sorts 8 partitions in parallel and costs `elapsed_compute` 1.94s. +//! Stage 1 merges them to one partition and runs the whole window on a +//! single core: 9.75s, 5x stage 0 for the same 10M rows. +//! +//! # Actual rule input +//! +//! Captured by logging `optimize`'s argument on the Q7 run above. The rule +//! sits after DataFusion's optimizer chain and before +//! [`DistributedExchangeRule`](super::DistributedExchangeRule), so on the +//! first pass there is no exchange or shuffle reader in the tree yet — just +//! the DataFusion plan with `EnforceSorting`'s `SortExec` placement already +//! materialized: +//! +//! ```text +//! ProjectionExec: expr=[id1@0 as id1, id2@1 as id2, id3@2 as id3, v2@3 as v2, +//! sum(large.v2) ORDER BY [large.id3 ASC NULLS LAST] +//! ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 +//! as my_rolling_sum] +//! BoundedWindowAggExec: wdw=[sum(large.v2) ORDER BY [large.id3 ASC NULLS LAST] +//! ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: +//! Field { nullable Float64 }], +//! frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, +//! mode=[Sorted] +//! SortPreservingMergeExec: [id3@2 ASC NULLS LAST] +//! SortExec: expr=[id3@2 ASC NULLS LAST], preserve_partitioning=[true] +//! DataSourceExec: file_groups={8 groups}, projection=[id1, id2, id3, v2], +//! file_type=parquet, +//! sort_order_for_reorder=[id3@2 ASC NULLS LAST] +//! ``` +//! +//! Note the frame is `ROWS`, not `RANGE`. The halo rule gates on +//! `WindowFrameUnits::Range`; this one must accept both, since for an +//! unbounded start the two differ only in tie handling at the frame edge. +//! +//! ## The rule fires more than once +//! +//! AQE re-plans as stages resolve, and `optimize` is called on each pass — +//! three times for this query. Passes 2 and 3 receive the plan wrapped in +//! `AdaptiveDatafusionExec` with the source subtree already behind a +//! resolved exchange: +//! +//! ```text +//! AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending +//! ProjectionExec: ... +//! BoundedWindowAggExec: ... +//! SortPreservingMergeExec: [id3@2 ASC NULLS LAST] +//! ExchangeExec: partitioning=None, plan_id=0, stage_id=0, stage_resolved=true +//! SortExec: expr=[id3@2 ASC NULLS LAST], preserve_partitioning=[true] +//! DataSourceExec: ... +//! ``` +//! +//! So the rewrite needs the same idempotency guard the halo rule uses +//! (`subtree_contains_our_rewrite`): without it, pass 2 would wrap the +//! output of pass 1 again. +//! +//! # Target shape +//! +//! The input partitions are **not** range-disjoint. Stage 0 is a bare +//! `SortExec` over 8 file groups, so each partition is locally sorted but +//! spans the whole value range — which is exactly why the SPM is needed for +//! correctness today. A prefix scan needs "all prior partitions" to be well +//! defined, so the rewrite has to introduce the disjointness itself, with +//! the same `RSE#1 → SortExec → ORRE → RSE#2` preamble the halo rule builds: +//! +//! ```text +//! PrefixMergeExec [per-partition state baked in by the scheduler] +//! ExchangeExec (partitioning: None) <- boundary 2, planted here +//! PartitionedBoundedWindowAggExec [wraps BWAG; UnspecifiedDistribution, +//! WindowStateCollector installed] +//! RangeFilterExec (halo_lo=0, halo_hi=0, cuts=pending) +//! [ExchangeExec] <- boundary 1, inserted by DER +//! RuntimeStatsExec #2 +//! OrderedRangeRepartitionExec [K range-disjoint outputs] +//! SortExec (preserve_partitioning=true) +//! RuntimeStatsExec #1 [local sketch → cuts] +//! +//! ``` +//! +//! ## Where the two stage boundaries come from +//! +//! **Boundary 1 is free.** [`DistributedExchangeRule`](super::DistributedExchangeRule) +//! walks bottom-up and, for any single-child node, tests whether that child +//! is partition-preserving *and* sits directly on an ORRE/URRE. Our +//! `RangeFilterExec`'s child (`RuntimeStatsExec #2` over the ORRE) satisfies +//! both, so DER wraps it in an `ExchangeExec`. The result is then a +//! recognized `is_stage_boundary` shape — RFE over exchange — so nothing +//! inserts a second one underneath. This is the same path the halo rule +//! rides; we plant the ORRE and DER does the rest. +//! +//! **Boundary 2 is not.** That test is two levels deep, not a walk: above +//! PBWAG there is no ORRE left (boundary 1 consumed it), and PBWAG isn't in +//! the `preserves_partitioning` whitelist anyway. So the rule plants this +//! `ExchangeExec` itself. `partitioning: None` is the passthrough encoding — +//! `ExchangeExec::new_with_details` maps it to `input.output_partitioning()` +//! and clones the input's `eq_properties`, so partition count and each +//! partition's ordering both survive. Every exchange DER creates today is +//! already a `None` one, so this is the well-trodden shape rather than a new +//! kind of boundary. +//! +//! What *is* new is the boundary's purpose. Every other boundary in Ballista +//! exists because data has to move — a fan-in, a repartition. This one moves +//! each partition to itself, and exists only so the scheduler has a +//! synchronization point at which every stage-1 task has published its +//! accumulator state. +//! +//! ## Differences from the halo rewrite +//! +//! - One `RangeFilterExec` with zero halo, not a wide/narrow pair. The trim +//! above the shuffle reader is needed regardless — the reader delivers a +//! superset and RFE narrows to the partition's own cut range. Halo is the +//! *widening* on top of that trim, and an unbounded start has no finite +//! reach to widen by; the correction rides accumulator state instead of +//! neighbouring rows. +//! - The SPM above BWAG is dropped rather than kept — that collapse is the +//! bottleneck being removed. +//! - Three stages rather than two, for the state round trip. +//! +//! ## Why it can't be one stage +//! +//! `PrefixMergeExec::try_new` takes its per-partition state by value, and +//! that state only exists once every upstream task has closed its window and +//! published accumulator state. So the collector rides stage 1's tasks, the +//! scheduler prefix-merges the reports as they arrive, and `PrefixMergeExec` +//! is constructed for stage 2 with the merged result baked in. +//! +//! The cost of that round trip is a full materialization of the window +//! output — for Q7, 10M rows — written and read back across boundary 2. The +//! halo path never pays it. Worth measuring against the 9.75s serial +//! baseline before assuming the parallel window is a net win. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use crate::state::aqe::execution_plan::ExchangeExec; +use ballista_core::config::BallistaConfig; +use ballista_core::execution_plans::{ + OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, PrefixMergeExec, + RangeFilterExec, RuntimeStatsExec, WindowApply, +}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::window::PlainAggregateWindowExpr; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::physical_plan::windows::BoundedWindowAggExec; +use datafusion::scalar::ScalarValue; +use log::debug; + +/// Physical optimizer pass: rewrite `UNBOUNDED PRECEDING` window frames into +/// a parallel prefix scan. See the [module docs][self] for the measured +/// starting point and the target shape. +/// +/// Shares `ballista.planner.parallel_window.enabled` with the halo rule for +/// now. The two rewrites are mutually exclusive on frame shape, so one key +/// selects both without ambiguity; whether prefix earns its own key is still +/// open. +#[derive(Debug, Clone, Default)] +pub struct PrefixWindowRule { + /// Shared with every other rule that plants an `ExchangeExec`, so + /// boundary 2's `plan_id` can't collide with one DER hands out. + plan_id_generator: Arc, +} + +impl PrefixWindowRule { + pub(crate) fn new(plan_id_generator: Arc) -> Self { + Self { plan_id_generator } + } +} + +impl PhysicalOptimizerRule for PrefixWindowRule { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> datafusion::common::Result> { + let bc = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + if !bc.parallel_window_enabled() { + return Ok(plan); + } + // The module docs' "Actual rule input" section is a capture of this. + // Re-run with `RUST_LOG=ballista_scheduler=debug` to refresh it after + // anything upstream in the optimizer chain changes shape. + debug!( + "PrefixWindowRule input:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); + // Same K as the halo rule: at rule-fire time the source is still one + // file_group, so the config knob is the only honest source of the + // eventual width. + let output_partitions = config.execution.target_partitions.max(2); + plan.transform_up(|node| { + match maybe_rewrite_bwag(&node, output_partitions, &self.plan_id_generator)? { + Some(rewritten) => Ok(Transformed::yes(rewritten)), + None => Ok(Transformed::no(node)), + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "PrefixWindow" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// True if any descendant plants what this rule plants. Idempotency guard: +/// AQE re-plans and calls `optimize` on every pass (three times for Q7), so +/// without this the second pass would wrap the first pass's output. +fn subtree_contains_our_rewrite(nodes: &[&Arc]) -> bool { + for node in nodes { + if node.is::() + || node.is::() + || node.is::() + { + return true; + } + if subtree_contains_our_rewrite(node.children().as_slice()) { + return true; + } + } + false +} + +/// ORRE requires `nulls_first == false` (T-Digest has no NULL slot). BWAG's +/// `NULLS LAST` expressions already arrive that way; sanitize anyway so the +/// invariant is visible. +fn normalize_sort_expr(expr: &PhysicalSortExpr) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: expr.expr.clone(), + options: SortOptions { + descending: expr.options.descending, + nulls_first: false, + }, + } +} + +/// Match the prefix-scan shape rooted at `node` and splice the target from +/// the [module docs][self] in place of the DF-planted +/// `BWAG → SPM → SortExec → ` subtree. +/// +/// - `Ok(None)`: shape gate missed. Silent — this runs on every node of +/// every plan. +/// - `Ok(Some(_))`: rewrote. +/// - `Err(_)`: a gate-guaranteed invariant didn't hold, or a constructor +/// failed. +fn maybe_rewrite_bwag( + node: &Arc, + output_partitions: usize, + plan_id_generator: &Arc, +) -> datafusion::common::Result>> { + let Some(window) = node.downcast_ref::() else { + return Ok(None); + }; + let [expr] = window.window_expr() else { + return Ok(None); + }; + let [] = expr.partition_by() else { + return Ok(None); + }; + let [order] = expr.order_by() else { + return Ok(None); + }; + let Some(column) = order.expr.downcast_ref::() else { + return Ok(None); + }; + // DESC support lands with the halo rule's — both need the mirrored + // bound handling and `RangeFilterExec::sorted_on_key` refuses DESC. + if order.options.descending { + return Ok(None); + } + let frame = expr.get_window_frame(); + // The gate that splits this rule from the halo one. Unlike that rule we + // accept both ROWS and RANGE units: with an unbounded start the two + // differ only in tie handling at the frame edge, which the prefix + // correction doesn't observe. + if !frame.start_bound.is_unbounded() { + return Ok(None); + } + // An unbounded *end* means the window sees rows it hasn't reached, so + // `uses_bounded_memory()` is false and DataFusion plans a `WindowAggExec` + // instead — which carries no observer and never reaches us. Guard anyway. + if frame.end_bound.is_unbounded() { + return Ok(None); + } + if subtree_contains_our_rewrite(window.children().as_slice()) { + return Ok(None); + } + + let node_children = node.children(); + let [immediate] = node_children.as_slice() else { + return datafusion::common::internal_err!( + "PrefixWindowRule: BWAG must have exactly 1 child" + ); + }; + // Peel whatever SPM/Sort combination EnforceSorting materialized. A + // source claiming its order natively via `sort_order_for_reorder` may + // have no SortExec at all. + let mut base_source: Arc = (*immediate).clone(); + while base_source.is::() || base_source.is::() { + let children = base_source.children(); + let [inner] = children.as_slice() else { + return datafusion::common::internal_err!( + "PrefixWindowRule: SPM/SortExec must have exactly 1 child" + ); + }; + base_source = (*inner).clone(); + } + let source_schema = base_source.schema(); + + // ORRE routes on the ORDER BY column, Float64-only today (T-Digest). + let routing_type = order.expr.data_type(&source_schema)?; + if !matches!(routing_type, DataType::Float64) { + return Ok(None); + } + + let sort_expr = normalize_sort_expr(order); + + // RSE#1 below the pipeline-breaking Sort so its sketch fully ingests and + // reports while Sort buffers — ORRE then routes against final cuts. + let rse1: Arc = Arc::new(RuntimeStatsExec::try_new( + base_source, + Some(vec![sort_expr.clone()]), + )?); + let sort_lex = LexOrdering::new(vec![sort_expr.clone()]).ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "PrefixWindowRule: could not build LexOrdering from ORDER BY".into(), + ) + })?; + let sorted_over_rse1: Arc = + Arc::new(SortExec::new(sort_lex, rse1).with_preserve_partitioning(true)); + let orre: Arc = Arc::new(OrderedRangeRepartitionExec::try_new( + sorted_over_rse1, + vec![sort_expr.clone()], + output_partitions, + )?); + let rse2: Arc = Arc::new(RuntimeStatsExec::try_new( + orre, + Some(vec![sort_expr.clone()]), + )?); + // Zero halo: the shuffle reader delivers a superset and this trims each + // task to its own cut range. There is no halo to widen by — the prefix + // correction rides accumulator state, not neighbouring rows. + let trim: Arc = Arc::new(RangeFilterExec::try_new_pending( + rse2, + sort_expr.expr.clone(), + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(0.0)), + )?); + // DER inserts boundary 1 under this trim, because its child (RSE#2) is + // partition-preserving and sits directly on the ORRE. + + let partitioned_bwag: Arc = Arc::new( + PartitionedBoundedWindowAggExec::try_new(window.window_expr().to_vec(), trim)?, + ); + + // Boundary 2. `None` partitioning is the passthrough encoding: partition + // count and per-partition ordering both carry across. Exists so the + // scheduler gets a synchronization point where every task on the stage + // below has published its accumulator state, not to move data. + let state_boundary: Arc = Arc::new(ExchangeExec::new( + partitioned_bwag, + None, + plan_id_generator.fetch_add(1, Ordering::Relaxed), + )); + + // BWAG appends its window columns after the input's, so expression `i` + // lands at `input_field_count + i` and the input columns keep their + // indices — which is why the aggregate's own argument expressions carry + // over unchanged despite being resolved against the input schema. + // + // SUM goes through the Aggregate path even though the cheaper Scalar path + // covers it: seeding an accumulator and replaying rows is the shape + // non-decomposable aggregates need, and exercising it where the answer is + // independently checkable beats the arrow-kernel shortcut. Choosing Scalar + // where it applies is a later optimization, worth measuring. + let input_field_count = window.input().schema().fields().len(); + let mut applies = Vec::new(); + for (expr_index, expr) in window.window_expr().iter().enumerate() { + // Non-aggregate window functions publish no state to merge, so they + // get no apply — `lead`/`lag` are handled by halos and the ranking + // family needs separate infrastructure. + let Some(plain) = expr.as_any().downcast_ref::() else { + continue; + }; + let aggregate = plain.get_aggregate_expr(); + applies.push(WindowApply::Aggregate { + udf: Arc::new(aggregate.fun().clone()), + args: aggregate.expressions(), + output_column: input_field_count + expr_index, + window_expr_index: expr_index, + }); + } + + // State is pending: it only exists once every task in the stage below + // has closed its window and reported. The scheduler resolves it via + // `PrefixMergeExec::resolve_state` when that stage completes. + let prefix_merge: Arc = + Arc::new(PrefixMergeExec::try_new_pending(state_boundary, applies)?); + + debug!( + "PrefixWindowRule: rewrote BWAG on `{}` ({} UNBOUNDED PRECEDING - {:?})", + column.name(), + frame.units, + frame.end_bound, + ); + Ok(Some(prefix_merge)) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::config::ExtensionOptions; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::displayable; + use datafusion::prelude::{SessionConfig, SessionContext}; + + /// Plan `sql` over an 8-partition source, so the physical plan matches + /// the shape captured from the real Q7 run: a per-partition `SortExec` + /// with `preserve_partitioning=true` under a `SortPreservingMergeExec` + /// that collapses to one partition for BWAG. A single-partition source + /// produces neither, which would make the "SPM is gone" assertion pass + /// against an input that never had one. + async fn plan(sql: &str) -> datafusion::common::Result> { + const PARTITIONS: usize = 8; + let schema = Arc::new(Schema::new(vec![ + Field::new("id1", DataType::Int64, false), + Field::new("id2", DataType::Int64, false), + Field::new("id3", DataType::Int64, false), + Field::new("v2", DataType::Float64, false), + ])); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(PARTITIONS), + ); + let table = MemTable::try_new(Arc::clone(&schema), vec![Vec::new(); PARTITIONS])?; + ctx.register_table("large", Arc::new(table))?; + ctx.sql(sql).await?.create_physical_plan().await + } + + /// Runs the rule with the shared parallel-window flag on, so + /// `ignores_finite_frame_shape` exercises the frame gate rather than the + /// config gate. + fn optimize( + plan: Arc, + ) -> datafusion::common::Result> { + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 8; + let mut bc = BallistaConfig::default(); + bc.set("planner.parallel_window.enabled", "true")?; + config.extensions.insert(bc); + PrefixWindowRule::default().optimize(plan, &config) + } + + /// `UNBOUNDED PRECEDING` frame, no PARTITION BY, single ascending + /// ORDER BY on a `Float64` column. Asserts the target from the module + /// docs. + /// + /// Orders by `v2` rather than h2o Q7's `id3` because the routing gate is + /// `Float64`-only today — see `declines_int64_routing`, which pins Q7's + /// actual shape and why it doesn't rewrite. + /// + /// Only this rule runs here, so boundary 1's `ExchangeExec` is absent — + /// `DistributedExchangeRule` inserts that one later in the chain. What + /// the rule itself owns is the ORRE preamble, the zero-halo trim, the + /// PBWAG swap, boundary 2, and the correction on top. + #[tokio::test] + async fn rewrites_unbounded_preceding_shape() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + + assert!( + rendered.contains("OrderedRangeRepartitionExec"), + "prefix scan needs range-disjoint partitions for \"all prior \ + partitions\" to be well defined:\n{rendered}" + ); + assert!( + rendered.contains("RuntimeStatsExec"), + "ORRE routes against sketched cuts, which RSE produces:\n{rendered}" + ); + assert!( + rendered.contains("RangeFilterExec"), + "the shuffle reader delivers a superset; each task must be \ + trimmed to its own cut range:\n{rendered}" + ); + assert!( + rendered.contains("PartitionedBoundedWindowAggExec"), + "BWAG must be wrapped so it runs per-partition:\n{rendered}" + ); + assert!( + rendered.contains("PrefixMergeExec"), + "per-partition results need a downstream prefix correction:\n{rendered}" + ); + assert!( + rendered.contains("ExchangeExec"), + "boundary 2 carries accumulator state back to the scheduler and \ + is planted by this rule, not by DER:\n{rendered}" + ); + assert!( + !rendered.contains("SortPreservingMergeExec"), + "SPM collapses K partitions to 1, which is the bottleneck being \ + removed:\n{rendered}" + ); + Ok(()) + } + + /// h2o Q7 verbatim. It has the right frame and the right partitioning, + /// but orders by `id3`, which is `Int64` in the h2o schema — and both + /// window rules gate routing on `Float64`, a T-Digest restriction that + /// lifts with the KLL migration. + /// + /// So Q7 is blocked on the sketch, not on anything in this rule. When + /// KLL lands this test starts failing, which is the signal to widen the + /// gate and delete it. + #[tokio::test] + async fn declines_int64_routing() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY id3 \ + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("PrefixMergeExec"), + "ORRE cannot route on a non-Float64 key until the sketch \ + widens:\n{rendered}" + ); + Ok(()) + } + + /// The halo rule's shape — finite `PRECEDING` start — must fall through + /// untouched. The two rules partition the frame space between them, so a + /// hit here would mean both fire on the same plan. + #[tokio::test] + async fn ignores_finite_frame_shape() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY v2 \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + !rendered.contains("PrefixMergeExec"), + "finite-frame windows belong to the halo rewrite:\n{rendered}" + ); + Ok(()) + } +} diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 984fc6554..5f4b75f48 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -22,7 +22,7 @@ use crate::state::aqe::execution_plan::{ use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; use crate::state::aqe::optimizer_rule::{ CoalescePartitionsRule, DelayJoinSelectionRule, DistributedExchangeRule, - ParallelWindowRule, PropagateEmptyExecRule, SelectJoinRule, + ParallelWindowRule, PrefixWindowRule, PropagateEmptyExecRule, SelectJoinRule, }; use crate::state::distributed_explain::handle_explain_plan; use crate::state::execution_stage::StageOutput; @@ -560,6 +560,13 @@ impl AdaptivePlanner { // picks up as the shuffle-boundary K-space source. physical_optimizers.push(Arc::new(ParallelWindowRule)); + // Sibling of ParallelWindowRule for UNBOUNDED PRECEDING frames. Same + // chain position: the two gate on complementary frame shapes, so + // their relative order is irrelevant. Shares the plan-id generator + // because it plants its own state-sync ExchangeExec. + physical_optimizers + .push(Arc::new(PrefixWindowRule::new(plan_id_generator.clone()))); + // `DistributedExchangeRule` should be the last plan mutator rule in the chain physical_optimizers .push(Arc::new(DistributedExchangeRule::new(plan_id_generator))); diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 279ebff09..277634c48 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -956,10 +956,13 @@ impl ExecutionGraph for StaticExecutionGraph { let SuccessfulTask { partitions, runtime_stats, + window_state, .. } = successful_task; running_stage .append_runtime_stats_reports(task_id, runtime_stats); + running_stage + .append_window_state_reports(task_id, window_state); locations.append(&mut partition_to_location( &job_id, task_id, stage_id, executor, partitions, diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 4643e3efd..43bfc3d9f 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -33,12 +33,12 @@ use log::{debug, warn}; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ - ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, + ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, TaskWindowState, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::{ FailedTask, OperatorMetricsSet, ResultLost, RuntimeStatsReport, SuccessfulTask, - TaskKilled, TaskStatus, + TaskKilled, TaskStatus, WindowStateReport, }; use ballista_core::serde::protobuf::{RunningTask, task_status}; use ballista_core::serde::scheduler::PartitionLocation; @@ -228,6 +228,18 @@ pub struct RunningStage { /// `partition_id` field is producer-local. Merged and logged once the /// stage finalizes; dropped when the stage transitions to Successful. pub runtime_stats_reports: Vec, + /// Finalized window-aggregate state reported by tasks in this stage + /// attempt. + /// + /// Each report already carries the stage-global partition it belongs to, + /// translated by the producing `ShuffleWriterExec` from the task-local + /// index DataFusion reports — so unlike `runtime_stats_reports` the tag + /// is not needed to address a producer file. It is still needed to purge + /// on reset: a retried task reports the same global partitions again, and + /// two states for one partition would make the prefix merge double-count. + /// Prefix-merged in global partition order to seed a downstream + /// `PrefixMergeExec`; dropped when the stage transitions to Successful. + pub window_state_reports: Vec, } /// If a stage finishes successfully, its task statuses and metrics will be finalized @@ -641,6 +653,7 @@ impl RunningStage { stage_metrics: None, session_config, runtime_stats_reports: Vec::new(), + window_state_reports: Vec::new(), } } @@ -858,6 +871,33 @@ impl RunningStage { })); } + /// Accumulate window-state reports as tasks in this stage attempt + /// complete. No producer tag: each report is already addressed by its + /// stage-global partition id. + pub fn append_window_state_reports( + &mut self, + producer_task_id: usize, + reports: Vec, + ) { + // TODO: prefix-merge these in global partition order and seed the + // downstream PrefixMergeExec. Logged for now so arrival is visible + // while the consumer is still being built. + for report in &reports { + log::debug!( + "stage {} window state arrived: global partition {} expr {} state {:?}", + self.stage_id, + report.global_partition_id, + report.window_expr_index, + report.state, + ); + } + self.window_state_reports + .extend(reports.into_iter().map(|report| TaskWindowState { + producer_task_id, + report, + })); + } + /// update and upsert the task metrics to the stage metrics pub fn update_task_metrics( &mut self, @@ -1040,6 +1080,8 @@ impl RunningStage { self.pending.reschedule(partitions); self.runtime_stats_reports .retain(|s| s.producer_task_id != task_id); + self.window_state_reports + .retain(|s| s.producer_task_id != task_id); } /// Reset the running and completed tasks on a given executor by @@ -1079,6 +1121,8 @@ impl RunningStage { self.pending.reschedule(to_reschedule); self.runtime_stats_reports .retain(|s| !reset_task_ids.contains(&s.producer_task_id)); + self.window_state_reports + .retain(|s| !reset_task_ids.contains(&s.producer_task_id)); reset } @@ -1185,6 +1229,7 @@ impl SuccessfulStage { // Fresh attempt: previous attempt's stats are irrelevant. // Merged-cut logging fires per-attempt on final success. runtime_stats_reports: Vec::new(), + window_state_reports: Vec::new(), } } @@ -1405,6 +1450,7 @@ mod tests { executor_id: "executor-1".to_string(), partitions: vec![], runtime_stats: vec![], + window_state: vec![], })), metrics: vec![], } @@ -1782,6 +1828,41 @@ mod tests { ); } + /// A reset task's window-state reports must go with it. The retry + /// re-runs the same partition slice and reports the same global + /// partitions again, so leaving the original attempt's entries behind + /// would give the prefix merge two states for one partition and + /// double-count them — a wrong running aggregate, not a degraded one. + #[test] + fn test_reset_task_info_purges_window_state_reports() { + let mut stage = make_running_stage(2); + append_running_task(&mut stage, 0, "executor-1", vec![0]); + append_running_task(&mut stage, 1, "executor-1", vec![1]); + + stage.append_window_state_reports(0, vec![make_window_state(0)]); + stage.append_window_state_reports(1, vec![make_window_state(1)]); + assert_eq!(stage.window_state_reports.len(), 2); + + stage.reset_task_info(0); + + assert_eq!(stage.window_state_reports.len(), 1); + assert_eq!(stage.window_state_reports[0].producer_task_id, 1); + assert_eq!( + stage.window_state_reports[0].report.global_partition_id, 1, + "the surviving task's partition must be the one left" + ); + } + + /// One window-state report for `global_partition_id`. + fn make_window_state(global_partition_id: u32) -> WindowStateReport { + WindowStateReport { + global_partition_id, + window_expr_index: 0, + partition_key: vec![], + state: vec![], + } + } + /// Executor loss resets every task the executor was hosting; the /// runtime-stats reports those (previously-Successful) producers had /// already contributed must be purged along with the task status. @@ -1804,6 +1885,7 @@ mod tests { executor_id: executor.to_string(), partitions: vec![], runtime_stats: vec![], + window_state: vec![], }); stage.append_runtime_stats_reports( task_id, diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index f1ed152a2..9fe99cf17 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,9 +36,8 @@ //! flows from parent to descendants via function arguments, so sibling //! subtrees never share state and there's no traversal-order dependency. -use ballista_core::execution_plans::{ - RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, -}; +use ballista_core::execution_plans::plan_algebra::as_partition_sliceable; +use ballista_core::execution_plans::{RangeShuffleReaderExec, ShuffleReaderExec}; use datafusion::common::internal_err; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ @@ -87,41 +86,20 @@ fn restrict( return Ok(rewritten); } - // RangeFilterExec: raw_bounds is indexed by input partition; restriction - // slices bounds parallel to the input's partition subset. Halos + routing - // are carried over verbatim; RFE re-widens on the fresh operator. - if !under_collect && let Some(rf) = plan.downcast_ref::() { + // Operators carrying data indexed by global input partition slice it + // parallel to the input restriction. Each one implements the slicing + // beside its own fields; this walker only supplies the restricted child. + if !under_collect && let Some(sliceable) = as_partition_sliceable(&plan) { let children = plan.children(); let [child] = children.as_slice() else { return internal_err!( - "RangeFilterExec must have exactly 1 child, got {}", + "{} is PartitionSliceable but has {} children, expected 1", + plan.name(), children.len() ); }; let new_child = restrict((*child).clone(), partitions, false)?; - let raw_bounds = rf.raw_bounds().ok_or_else(|| { - datafusion::common::DataFusionError::Internal( - "RangeFilterExec: task-restriction before resolve_bounds()".into(), - ) - })?; - let sliced_bounds: Vec<_> = partitions - .iter() - .map(|&global| { - raw_bounds.get(global).cloned().ok_or_else(|| { - datafusion::common::DataFusionError::Internal(format!( - "RangeFilterExec: partition index {global} out of bounds ({} raw bounds)", - raw_bounds.len() - )) - }) - }) - .collect::>()?; - return Ok(Arc::new(RangeFilterExec::try_new_resolved( - new_child, - rf.routing_expr().clone(), - rf.halo_lo().clone(), - rf.halo_hi().clone(), - sliced_bounds, - )?)); + return sliceable.slice_to_partitions(new_child, partitions); } // UnionExec: parent partition `p` maps to exactly one child's local diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 08bc7b9f7..0bd6e20b7 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -309,6 +309,7 @@ pub fn default_task_runner() -> impl TaskRunner { executor_id: executor_id.clone(), partitions: partitions.clone(), runtime_stats: vec![], + window_state: vec![], })), }); } @@ -1265,6 +1266,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat executor_id: executor_id.to_owned(), partitions, runtime_stats: vec![], + window_state: vec![], })), } }