feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates

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<Vec<f64>> {
+        (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<f64> {
+        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<TempDir> {
+        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<Vec<(f64, f64)>> {
+        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<f64> = (1..=TOTAL_ROWS).map(|v| v as f64).collect();
+        let actual_v: Vec<f64> = rows.iter().map(|(v, _)| *v).collect();
+        assert_eq!(actual_v, expected_v, "input rows must survive the rewrite");
+
+        let actual_rs: Vec<f64> = 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<f64> = rows.iter().map(|(_, rs)| *rs).collect();
+        assert_eq!(actual_rs, expected_running_sums());
+        Ok(())
+    }
+}

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 `<data>.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.

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;