You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by av...@apache.org on 2026/08/25 06:57:31 UTC
(datafusion-ballista) branch main updated: feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates (#2211)
This is an automated email from the ASF dual-hosted git repository.
avantgardnerio pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 21a7d32a5 feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates (#2211)
21a7d32a5 is described below
commit 21a7d32a549c6b308d01b710c32219c786c68d27
Author: Brent Gardner <bg...@squarelabs.net>
AuthorDate: Tue Aug 25 00:57:27 2026 -0600
feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates (#2211)
* feat(core): add PrefixMergeExec for cross-partition windowed-aggregate state merge
Introduces PrefixMergeExec as the downstream half of the AQE range-shuffle
prefix-scan pipeline: it takes per-input-partition window-aggregate state that
the scheduler has already prefix-merged and applies it row-wise to the current
partition's output, so cross-partition running aggregates come out correct.
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.
Purely additive: nothing in-tree constructs a PrefixMergeExec. The remaining
work is the state source — collecting each upstream task's finalized
accumulator state out of BoundedWindowAggExec and transporting it to the
scheduler — which lands separately.
FinalizedPartitionState is defined locally, indexed by window-expression
position, as the shape this operator consumes.
wip(core,scheduler): prefix-window rewrite plants the shape, collector captures state
Follows the data flow end to end for the AQE prefix-scan pipeline. Stages 0
and 1 run on a real cluster; stage 2 is blocked on PrefixMergeExec serde.
PrefixWindowRule: sibling of ParallelWindowRule for UNBOUNDED PRECEDING
frames, gating on start_bound.is_unbounded() where that rule gates on
is_finite() — complementary, so no plan matches both. Plants the ORRE
preamble, a zero-halo RangeFilterExec trim, PBWAG, a passthrough ExchangeExec
for the state round trip, and PrefixMergeExec. Accepts ROWS as well as RANGE
units, which for an unbounded start differ only in tie handling.
Module docs record the rule's actual captured input rather than an assumed
one, including that AQE re-plans and calls optimize three times — hence the
idempotency guard.
WindowStateCollector: implements DataFusion 55's WindowStateObserver and
retains each finalized accumulator state. Retention rather than polling
because Accumulator::state is a destructive read fired at most once per
group. PBWAG installs one exactly when every frame is ever-expanding, the
same condition with_state_observer enforces, so the halo shape is untouched
and the wire format needs no new field.
PartitionSliceable: operators carrying data indexed by global input partition
now implement their own slicing next to the fields being sliced, replacing
two bespoke arms in the scheduler's task builder. RangeFilterExec's bounds and
PrefixMergeExec's state/offsets both slice when a task is restricted to a
partition subset — without which PrefixMergeExec attaches each partition's
offsets to the wrong rows.
Tests: a client-side e2e asserting the running sum against a computed oracle
through the real distributed path. Red until the transport lands, since
PrefixMergeExec is a passthrough with no state. The rule's unit tests pin why
h2o Q7 does not rewrite today — it orders by an Int64 column and ORRE routes
on a Float64-only T-Digest until KLL.
Known gaps: no serde for PrefixMergeExec; no transport from collector to
scheduler; observed partition_idx is task-local and needs pairing with the
task's global partition ids.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
wip(core,executor): ShuffleWriter translates window state to global partition ids
Continues following the data: BWAG accumulator -> collector -> PBWAG getter ->
ShuffleWriter -> executor. Logged at task completion rather than transported,
so the path is exercised end to end before anything is built on it.
The local-to-global translation lives on the writer, not on 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; the writer is the node the
scheduler hands global_output_partition_ids to. Reassembling downstream instead
would have the scheduler re-derive a mapping it already computed, and a prefix
scan fed a permuted order is wrong with nothing to show for it.
Verified on the client e2e: two tasks each covering two partitions previously
both reported local 0 and 1; they now report globals 0/1 and 2/3, with states
10/26/42/58 over input 1..16.
collect_window_state joins collect_plan_metrics and collect_runtime_stats_reports
as a task-completion peer on QueryStageExecutor, wired into both the pull
(execution_loop) and push (executor_server) task paths.
PBWAG grows observed_window_state() and keeps the TODO that the install site
moves when the wrapper collapses. Neither the collector nor the writer-side
walk depends on the wrapper beyond one downcast.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
wip(core,scheduler,executor): transport window state to the scheduler
Adds WindowStateReport to SuccessfulTask, completing the path from a task's
BWAG accumulator to the scheduler: collector -> PBWAG -> ShuffleWriter ->
wire -> RunningStage::window_state_reports. Verified on the client e2e, where
four range-disjoint partitions over input 1..16 arrive as globals 0..3 with
states 10/26/42/58.
State and partition key cross as datafusion_common.ScalarValue rather than a
numeric field, so sketch-backed aggregates (approx_distinct's HLL blob) work
unchanged. The proto carries a TODO on payload size: if it stops being small,
write the state as a sidecar beside the shuffle files the way sort-shuffle
already writes <data>.arrow.index, and send only a reference.
Failures fail the task rather than dropping a report. Unlike runtime stats,
which are an optimization input, this state is load-bearing: the downstream
prefix merge is arithmetically wrong without every partition's contribution,
and wrong in a way nothing later detects. Collection is skipped entirely when
execution already failed.
Reports are tagged with their producer task and purged on reset, in both
reset_task_info and reset_tasks. A retried task re-runs its slice and reports
the same global partitions again; without the purge the stage would hold two
states for one partition and the prefix merge would double-count them. The
file-addressing reason RuntimeStats needs its tag does not transfer — the
writer already stamped stage-global ids — but the purge reason does.
Both scheduler task-status paths (classic execution_graph and AQE) and both
executor task paths (pull execution_loop and push executor_server) are wired;
each pair are peer implementations that need every completion hook.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
wip(core): prefix-scan the accumulated window state on the scheduler
prefix_merge_window_state turns per-partition finalized states into one
carry-in per partition: out[0] empty, out[k] the merge of every partition
before k. 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
arithmetic here, which is what lets non-decomposable aggregates work — two
approx_distinct HLL sketches combine correctly where two distinct counts
could not. The accumulator comes from PlainAggregateWindowExpr, the type an
ever-expanding frame always produces; a sliding expression reaching this is an
error rather than a silently wrong answer.
Built incrementally, out[k] = merge(out[k-1], state[k-1]), so two merges per
partition rather than merging every prior from scratch. A fresh accumulator
per partition is still required because Accumulator::state is a destructive
read and must not be called twice; seeding it from the previous carry-in is
the same round trip two-phase aggregation makes.
Enforces here what stopped being DataFusion's guarantee when the API turned
out to be push-shaped: a report carrying a PARTITION BY key is rejected,
because FinalizedPartitionState has no key dimension and a second group in one
partition would have nowhere to go.
Tests cover the carry-in arithmetic, independence from report arrival order
(reports arrive in close order and tasks complete in any order), rejection of
a duplicate partition state (only reachable if the producer-task purge failed,
and would double-count), and carrying across an empty partition that closes no
group and so publishes nothing.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
wip(core,scheduler): bake the prefix-merged state into PrefixMergeExec
Closes the loop from reports to operator, all in scheduler memory. When a
stage completes, its accumulated window-state reports are prefix-merged and
bound to the PrefixMergeExec waiting on it downstream. Verified on the client
e2e, where four partitions reporting 10/26/42/58 resolve to carry-ins
[None, 10, 36, 78].
State is late-bound, mirroring RangeFilterExec's cuts: try_new_pending for the
rule's plant-time path, try_new_resolved for wire decode and task restriction,
resolve_state as the setter. Both execute() and slice_to_partitions refuse
while unresolved rather than treating an absent carry-in as zero — that would
emit partition-local aggregates, which look plausible and are wrong.
The scheduler hooks update_stage_progress on completion: walk the plan for the
PrefixMergeExec whose state-sync boundary carries this stage id, recover the
window expressions from the operator below that boundary (the exchange retains
its input subtree even once resolved), prefix-merge, resolve. A stage that
reported state with no consumer to bind it to is an error rather than a skip;
the state exists because something downstream cannot be correct without it.
Still passthrough: `applies` is empty, so the operator carries the state
without applying it. The descriptors that turn state into corrected columns,
and the serde that lets the operator reach an executor, are next.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
wip(core): serde for PrefixMergeExec
Stage 2 now reaches an executor and runs. The e2e's remaining failure is the
expected one: every partition's running sum is off by exactly its carry-in
(5+10=15, 9+36=45, 13+78=91), so the state that crossed the wire is provably
correct and nothing is applying it yet.
Both WindowApply shapes cross. The aggregate arm carries its UDAF by name,
resolved from the executor's function registry on decode, with args as
PhysicalExprNodes. State crosses as ScalarValue so sketch-backed aggregates
work unchanged, and an absent slot stays distinct from a present-but-empty one
— a non-aggregate window function publishes no state, which is not the same as
publishing nothing.
Encoding refuses while state is unresolved, matching RangeFilterExec's refusal
on unresolved bounds. An executor has no way to obtain prefix state, so a plan
reaching the wire without it could only produce partition-local aggregates.
Round-trip test covers both arms, the UDAF-by-name resolution, and the
None-vs-empty distinction in the state slots.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
feat(scheduler,core): apply descriptors — parallel prefix scan is correct
The client e2e passes: a running sum computed across four range-disjoint
partitions matches the serial answer through the real distributed engine,
scheduler and shuffle and executor included.
One WindowApply per aggregate window expression. 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. Non-aggregate window functions get no apply; they publish no
state to merge. This is the last of the stubs each earlier step stood on: with
`applies` empty the operator was a passthrough, which is why the state was
provably correct and the output was still partition-local.
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.
Also flattens the codec arms added in the previous commit. The prefix-state
encoder was three nested maps around a transpose; it is now six named helpers
built from plain loops, taking the decode arm from ~85 lines to 13 and the
encode arm from ~95 to 24. Option handling is an explicit match rather than
`.map(..).transpose()?`, which puts the absent-versus-empty distinction where
a reader can see it, and each helper names what it was converting so failures
read as "failed to encode prefix state" rather than an anonymous try_from.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
docs: correct status claims that went stale when the pipeline went green
Five places still described a half-built system. Each was accurate when
written and became a lie at a different commit:
- the client e2e's "Fails today" doc, now stating the quiet failure it
guards against rather than predicting one
- the rule's "Status: shape only", which claimed the rewrite corrected
nothing; it now records what it is correct for, and that h2o Q7 is blocked
on the sketch's Float64 restriction rather than on this rule
- prefix_merge's "Nothing in-tree collects it yet", which now points at the
collector that does
- the collector's and the stage's scaffolding-log comments, both explaining
themselves as stand-ins for consumers that now exist
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
fix(core): review follow-ups — Overwrite doc, non_exhaustive, named type drift
Three small ones from review.
ScalarOp::Overwrite claimed to fit last_value. Over an ever-expanding frame
last_value is the current row's own value and needs no correction, so
overwriting every row with one scalar would be wrong. The doc now says
first_value only, and says why last_value is excluded.
ScalarOp and WindowApply are #[non_exhaustive]. Both are expected to grow —
the ranking family needs a segment-tree broadcast shape — and each addition
would otherwise be a breaking change for anyone matching on them. Construction
is unaffected, so the rule still builds an Aggregate apply.
Type drift now names the apply responsible. Arrow reports a mismatch at a
column index and nothing about which correction produced it, which is the
wrong half when several applies rewrite one batch. rebuild_batch reports
"applies[3] produced Int64 for column 1, which the schema declares as
Float64", using the apply_index AggregateApply already carried for exactly
this and did not use.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
feat(core): metrics on PrefixMergeExec, split by apply path
BaselineMetrics for elapsed_compute and output_rows, per partition because
execute(partition) builds a stream each. Plus a rows_corrected counter and a
separate timer per apply path.
The split is the point. WindowApply::Scalar is an arrow kernel over a whole
batch; 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. On the client e2e the operator now reports
aggregate_apply_time=234.34us against elapsed_compute=238.60us, so the replay
is 98% of its time on a trivial SUM — a shape rather than a magnitude at 16
rows, but it makes the replay cost something the operator reports in
production rather than something only a benchmark can see.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
refactor(core): FinalizedPartitionState as a newtype
Vec<Option<Vec<ScalarValue>>> appeared in this operator's public signatures,
where it is neither readable nor searchable, and it spelled "no state for this
window expression" two ways: a missing index, and a None at a present index.
Every caller handled both. slot(window_expr_index) collapses them into one
answer, and a later change to the representation now stays internal.
Also removes two comments claiming DataFusion guarantees at most one PARTITION
BY group per partition. It does not — apache/datafusion#24035 shipped a
callback keyed by group, so that invariant is ours, and the scheduler enforces
it by rejecting any report carrying a key. A window that does have a PARTITION
BY needs nothing from this operator anyway: BoundedWindowAggExec asks for
KeyPartitioned input, so each partition's window is already independent.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
Conflicts:
ballista/core/src/execution_plans/mod.rs
ballista/core/src/serde/mod.rs
ballista/scheduler/src/state/aqe/mod.rs
ballista/scheduler/src/state/aqe/planner.rs
ballista/scheduler/src/state/task_builder.rs
* reviewing PrefixMerge
* reviewed PrefixMerge
* reviewed ShuffleWriter
* reviewed window_state.rs
* reviewing prefix rule
* bug fix
* reviewed aqe mod
* reviewed execution_stage.rs
* test(core): pin the wire size of a sketch window state
`approx_distinct`'s HLL is a dense register array, so one state encodes to
16393 bytes whatever the row count. Asserts the size is independent of
cardinality, with a loose ceiling as a tripwire for a representation change.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
* fix(chaos): reserve a cluster's ports without releasing them in between
`free_port` bound :0, read the port, then dropped the listener, so two
consecutive calls could return the same number. `spawn_executor` asked
twice and handed its child one port for both the flight and gRPC servers:
the second bind failed AddrInUse, the executor exited during startup, and
the harness timed out waiting for two registrations.
Seen in CI as `port: 38571, grpc_port: 38571` on one executor's
registration, with distinct ports on the other. Not reproducible locally
(0 collisions in 1000 attempts at the old pattern), so `free_ports` fixes
it by construction instead: hold every listener until all N ports have
been read, which the OS cannot satisfy with a duplicate.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---------
Co-authored-by: Claude Opus 5 (1M context) <no...@anthropic.com>
---
ballista/client/tests/prefix_window.rs | 342 +++++
ballista/core/proto/ballista.proto | 103 ++
ballista/core/src/execution_plans/mod.rs | 7 +
.../partitioned_bounded_window_agg.rs | 60 +-
ballista/core/src/execution_plans/plan_algebra.rs | 83 +-
ballista/core/src/execution_plans/prefix_merge.rs | 1380 ++++++++++++++++++++
ballista/core/src/execution_plans/range_filter.rs | 37 +-
ballista/core/src/execution_plans/runtime_stats.rs | 162 ++-
.../core/src/execution_plans/shuffle_writer.rs | 52 +-
ballista/core/src/execution_plans/window_state.rs | 554 ++++++++
ballista/core/src/serde/generated/ballista.rs | 156 ++-
ballista/core/src/serde/mod.rs | 321 ++++-
ballista/executor/src/execution_engine.rs | 40 +
ballista/executor/src/execution_loop.rs | 13 +
ballista/executor/src/executor_server.rs | 13 +
ballista/executor/src/lib.rs | 11 +-
ballista/scheduler/src/scheduler_server/mod.rs | 1 +
.../src/scheduler_server/query_stage_scheduler.rs | 1 +
ballista/scheduler/src/state/aqe/mod.rs | 111 +-
.../scheduler/src/state/aqe/optimizer_rule/mod.rs | 2 +
.../src/state/aqe/optimizer_rule/prefix_window.rs | 547 ++++++++
ballista/scheduler/src/state/aqe/planner.rs | 9 +-
ballista/scheduler/src/state/execution_graph.rs | 3 +
ballista/scheduler/src/state/execution_stage.rs | 83 +-
ballista/scheduler/src/state/task_builder.rs | 41 +-
ballista/scheduler/src/test_utils.rs | 2 +
chaos-testing/src/cluster.rs | 48 +-
27 files changed, 4104 insertions(+), 78 deletions(-)
diff --git a/ballista/client/tests/prefix_window.rs b/ballista/client/tests/prefix_window.rs
new file mode 100644
index 000000000..f95bf6f88
--- /dev/null
+++ b/ballista/client/tests/prefix_window.rs
@@ -0,0 +1,342 @@
+// 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, Int64Array};
+ 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()
+ }
+
+ /// Non-NULL keys per partition of the nullable-key table. The remaining
+ /// row of each partition carries a NULL key.
+ const KEYS_PER_PARTITION: usize = ROWS_PER_PARTITION - 1;
+ /// `v` on every NULL-key row.
+ ///
+ /// Equal across the whole run, which is what makes the expectation well
+ /// defined: NULLs tie under `ORDER BY k`, so their order among themselves
+ /// is arbitrary, and equal values give the same running sums whichever
+ /// order the engine picks.
+ const NULL_KEY_VALUE: f64 = 100.0;
+
+ /// `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.
+ ///
+ /// 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)
+ }
+
+ /// `(k, v)` rows per partition, where `k` is a nullable `Int64` routing
+ /// key and `v` is the value summed over it.
+ ///
+ /// One NULL per input partition rather than one partition of NULLs, so
+ /// ORRE has to gather the run out of every input instead of passing a
+ /// single partition through.
+ fn nullable_key_partitions() -> Vec<Vec<(Option<i64>, f64)>> {
+ (0..PARTITIONS)
+ .map(|partition| {
+ let mut rows: Vec<(Option<i64>, f64)> = (0..KEYS_PER_PARTITION)
+ .map(|row| {
+ let key = (partition * KEYS_PER_PARTITION + row + 1) as i64;
+ (Some(key), key as f64)
+ })
+ .collect();
+ rows.push((None, NULL_KEY_VALUE));
+ rows
+ })
+ .collect()
+ }
+
+ /// Running sums of [`nullable_key_partitions`] in `k ASC NULLS LAST`
+ /// order: the non-NULL keys ascending, then the NULL run.
+ fn expected_null_key_running_sums() -> Vec<f64> {
+ let mut total = 0.0;
+ let mut sums: Vec<f64> = (1..=PARTITIONS * KEYS_PER_PARTITION)
+ .map(|key| {
+ total += key as f64;
+ total
+ })
+ .collect();
+ sums.extend((0..PARTITIONS).map(|_| {
+ total += NULL_KEY_VALUE;
+ total
+ }));
+ sums
+ }
+
+ /// Register [`nullable_key_partitions`] as `tk(k Int64 NULL, v Float64)`,
+ /// one parquet file per partition. Same reasons as [`register_input`].
+ async fn register_nullable_key_input(ctx: &SessionContext) -> Result<TempDir> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int64, true),
+ Field::new("v", DataType::Float64, false),
+ ]));
+ let dir = TempDir::new().expect("temp dir");
+ for (partition, rows) in nullable_key_partitions().into_iter().enumerate() {
+ let keys: Int64Array = rows.iter().map(|(k, _)| *k).collect();
+ let values: Float64Array = rows.iter().map(|(_, v)| *v).collect();
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(keys), Arc::new(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(
+ "tk",
+ dir.path().to_str().expect("utf-8 temp path"),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+ Ok(dir)
+ }
+
+ /// Running sums over `tk`, ordered by the nullable key.
+ async fn running_sums_by_null_key(ctx: &SessionContext) -> Result<Vec<f64>> {
+ let batches = ctx
+ .sql(
+ "SELECT k, \
+ sum(v) OVER (ORDER BY k \
+ ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS rs \
+ FROM tk \
+ ORDER BY k",
+ )
+ .await?
+ .collect()
+ .await?;
+
+ let mut sums = Vec::with_capacity(TOTAL_ROWS);
+ for batch in &batches {
+ let rs = datafusion::common::cast::as_float64_array(batch.column(1))?;
+ for row in 0..batch.num_rows() {
+ sums.push(rs.value(row));
+ }
+ }
+ Ok(sums)
+ }
+
+ /// `(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, on an `Int64` key that carries NULLs — the two things the rule's
+ /// type gate used to refuse.
+ ///
+ /// A NULL run is the case where "all prior partitions" could quietly stop
+ /// meaning anything: the run has to sort wholly to one end and land in one
+ /// partition, or the rows around it get a prefix that skips it.
+ #[tokio::test]
+ async fn prefix_scan_matches_serial_running_sum_with_null_keys() -> Result<()> {
+ let ctx = context(true).await;
+ let _data = register_nullable_key_input(&ctx).await?;
+ let sums = running_sums_by_null_key(&ctx).await?;
+
+ assert_eq!(
+ sums,
+ expected_null_key_running_sums(),
+ "running sums over a nullable key must be global — the NULL run \
+ sorts last and must still be carried into the prefix"
+ );
+ 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(())
+ }
+
+ /// The nullable-key query with the rewrite off, guarding
+ /// `expected_null_key_running_sums` the same way.
+ #[tokio::test]
+ async fn serial_null_key_running_sum_is_correct_without_rewrite() -> Result<()> {
+ let ctx = context(false).await;
+ let _data = register_nullable_key_input(&ctx).await?;
+ let sums = running_sums_by_null_key(&ctx).await?;
+
+ assert_eq!(sums, expected_null_key_running_sums());
+ Ok(())
+ }
+}
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
index 4273fde81..7142d1246 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;
}
}
@@ -213,6 +214,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.
@@ -713,6 +780,42 @@ 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.
+ 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 2c5f0bee5..d8fc21701 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::{InputOrder, 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..5d5ba72c0 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,10 @@ pub struct PartitionedBoundedWindowAggExec {
inner_bwag: Arc<BoundedWindowAggExec>,
/// Multi-partition input; same `Arc` `inner_bwag.input()` holds.
input: Arc<dyn ExecutionPlan>,
+ /// 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.
+ state_collector: Option<Arc<WindowStateCollector>>,
}
impl PartitionedBoundedWindowAggExec {
@@ -104,13 +111,56 @@ impl PartitionedBoundedWindowAggExec {
window_expr: Vec<Arc<dyn WindowExpr>>,
input: Arc<dyn ExecutionPlan>,
) -> Result<Self> {
- 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.
+ 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<dyn WindowStateObserver>
+ ))?,
+ 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<WindowStateCollector>> {
+ 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<ObservedWindowState> {
+ 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::<RuntimeStatsExec>().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<dyn ExecutionPlan>,
+ partitions: &[usize],
+ ) -> Result<Arc<dyn ExecutionPlan>>;
+}
+
+/// 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<dyn ExecutionPlan>,
+) -> Option<&dyn PartitionSliceable> {
+ if let Some(op) = plan.downcast_ref::<RangeFilterExec>() {
+ return Some(op);
+ }
+ if let Some(op) = plan.downcast_ref::<PrefixMergeExec>() {
+ 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<T: Clone>(
+ values: &[T],
+ partitions: &[usize],
+ owner: &str,
+ what: &str,
+) -> Result<Vec<T>> {
+ 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..b17c1a1f8
--- /dev/null
+++ b/ballista/core/src/execution_plans/prefix_merge.rs
@@ -0,0 +1,1380 @@
+// 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.
+//!
+//! 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.
+///
+/// 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. This operator applies it; it does not compute it.
+#[derive(Debug, Clone, Default, PartialEq)]
+pub struct FinalizedPartitionState {
+ /// Indexed by position in the upstream operator's `window_expr()` list.
+ per_window_expr: Vec<Option<Vec<ScalarValue>>>,
+}
+
+impl FinalizedPartitionState {
+ /// Build from one slot per window expression, in `window_expr()` order.
+ pub fn new(per_window_expr: Vec<Option<Vec<ScalarValue>>>) -> 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<ScalarValue>> {
+ 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<Vec<ScalarValue>>] {
+ &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.
+#[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.
+#[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`.
+ offset: Vec<ScalarValue>,
+ /// 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.
+ Aggregate {
+ /// UDF used to construct a fresh `Accumulator` per partition.
+ udf: Arc<AggregateUDF>,
+ /// 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<Arc<dyn PhysicalExpr>>,
+ /// 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,
+ },
+ // TODO: segment trees, deconstructable ops (AVG, stddev)
+}
+
+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.
+pub struct PrefixMergeExec {
+ input: Arc<dyn ExecutionPlan>,
+ /// One entry per window-function output column that needs cross-partition
+ /// correction. Non-corrected columns here.
+ applies: Vec<WindowApply>,
+ /// `per_partition_state[k]` is the *already-merged* state summarizing
+ /// every input partition in `[0..k)`.
+ /// 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<Mutex<Option<Vec<FinalizedPartitionState>>>>,
+ properties: Arc<PlanProperties>,
+ 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<dyn ExecutionPlan>,
+ applies: Vec<WindowApply>,
+ ) -> Result<Self> {
+ 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<dyn ExecutionPlan>,
+ applies: Vec<WindowApply>,
+ per_partition_state: Vec<FinalizedPartitionState>,
+ ) -> Result<Self> {
+ Self::try_new_inner(input, applies, Some(per_partition_state))
+ }
+
+ fn try_new_inner(
+ input: Arc<dyn ExecutionPlan>,
+ applies: Vec<WindowApply>,
+ per_partition_state: Option<Vec<FinalizedPartitionState>>,
+ ) -> Result<Self> {
+ 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.
+ pub fn resolve_state(&self, state: Vec<FinalizedPartitionState>) -> 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.
+ pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
+ &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<Vec<FinalizedPartitionState>> {
+ 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.
+ fn slice_to_partitions(
+ &self,
+ child: Arc<dyn ExecutionPlan>,
+ partitions: &[usize],
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ 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,
+ }),
+ // [`WindowApply::Aggregate`] carries `window_expr_index`,
+ // which indexes window expressions rather than partitions, so it rides
+ // over unchanged.
+ aggregate => Ok(aggregate.clone()),
+ })
+ .collect::<Result<Vec<_>>>()?;
+ 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<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ 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<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ 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<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ 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<Distribution> {
+ 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<Option<OrderingRequirements>> {
+ 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<bool> {
+ vec![true]
+ }
+
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ 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<Statistics>],
+ _args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ Ok(Arc::clone(&input_stats[0]))
+ }
+
+ fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+ vec![ChildStats::At(partition)]
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ 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<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ 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<PreparedApply> = 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<RecordBatch> {
+ 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<RecordBatch> {
+ 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<dyn Accumulator>,
+ /// arguments for the accumulator, ie. the x in AVG(x)
+ args: Vec<Arc<dyn PhysicalExpr>>,
+ output_column: usize,
+}
+
+impl AggregateApply {
+ fn new(
+ apply_index: usize,
+ udf: &Arc<AggregateUDF>,
+ args: &[Arc<dyn PhysicalExpr>],
+ output_column: usize,
+ input_schema: &SchemaRef,
+ offset_state: Option<&Vec<ScalarValue>>,
+ ) -> Result<Self> {
+ 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<ArrayRef> = state_scalars
+ .iter()
+ .map(|s| s.to_array_of_size(1))
+ .collect::<Result<Vec<_>>>()?;
+ 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<RecordBatch> {
+ let num_rows = batch.num_rows();
+ if num_rows == 0 {
+ return Ok(batch);
+ }
+ let arg_arrays: Vec<ArrayRef> = 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::<Result<Vec<_>>>()?;
+
+ let mut new_values: Vec<ScalarValue> = Vec::with_capacity(num_rows);
+ for row_idx in 0..num_rows {
+ let row_args: Vec<ArrayRef> =
+ arg_arrays.iter().map(|a| a.slice(row_idx, 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<PreparedApply>,
+ schema: SchemaRef,
+ baseline: BaselineMetrics,
+ path_metrics: PathMetrics,
+}
+
+impl Stream for ApplyStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ // 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.
+///
+/// # 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<ArrayRef>,
+ apply_index: usize,
+ output_column: usize,
+) -> Result<RecordBatch> {
+ 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<ArrayRef> = 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<dyn ExecutionPlan> {
+ let schema = one_col_schema();
+ let mut per_partition: Vec<Vec<RecordBatch>> = 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<FinalizedPartitionState> {
+ (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<AggregateUDF>,
+ input_schema: &SchemaRef,
+ values: &[i64],
+ ) -> Result<Vec<ScalarValue>> {
+ let column: Arc<dyn PhysicalExpr> =
+ 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<dyn ExecutionPlan> =
+ 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<FinalizedPartitionState> = 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<RecordBatch> =
+ exec.execute(0, ctx.clone())?.try_collect().await?;
+ let p0_running = out_p0[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<datafusion::arrow::array::UInt64Array>()
+ .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<RecordBatch> =
+ exec.execute(1, ctx.clone())?.try_collect().await?;
+ let p1_running = out_p1[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<datafusion::arrow::array::UInt64Array>()
+ .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<dyn ExecutionPlan> =
+ 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<dyn PhysicalExpr> = 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<FinalizedPartitionState> = 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<RecordBatch> =
+ exec.execute(0, ctx.clone())?.try_collect().await?;
+ let p0_avg = out_p0[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .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<RecordBatch> =
+ exec.execute(1, ctx.clone())?.try_collect().await?;
+ let p1_avg = out_p1[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .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<dyn ExecutionPlan> =
+ 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<RecordBatch> =
+ exec.execute(0, ctx.clone())?.try_collect().await?;
+ let p0_sum = out_p0[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .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<RecordBatch> =
+ exec.execute(1, ctx.clone())?.try_collect().await?;
+ let p1_sum = out_p1[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .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 db7145c16..0062dd2cf 100644
--- a/ballista/core/src/execution_plans/range_filter.rs
+++ b/ballista/core/src/execution_plans/range_filter.rs
@@ -81,6 +81,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<ScalarValue>, Option<ScalarValue>);
@@ -154,7 +158,8 @@ impl RangeFilterExec {
/// `raw_bounds.len()`.
/// * `filter_expr` - numeric physical expression each row is compared by.
/// * `halo_lo`, `halo_hi` - non-negative widening amounts applied by
- /// [`RangeFilterExec::resolve_bounds`]. Both must be finite Float64 today.
+ /// [`RangeFilterExec::resolve_bounds`], in `filter_expr`'s own type. A
+ /// float halo must also be finite.
/// * `input_order` - the order in which rows will arrive
pub fn try_new_pending(
input: Arc<dyn ExecutionPlan>,
@@ -327,6 +332,36 @@ 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<dyn ExecutionPlan>,
+ partitions: &[usize],
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ 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.filter_expr().clone(),
+ self.halo_lo().clone(),
+ self.halo_hi().clone(),
+ self.input_order(),
+ 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/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs
index 60e40bc56..e2b33272a 100644
--- a/ballista/core/src/execution_plans/runtime_stats.rs
+++ b/ballista/core/src/execution_plans/runtime_stats.rs
@@ -955,29 +955,35 @@ pub fn cut_partitions(
);
};
let sub_part_id = file.partition_id.partition_id as u32;
- // A file of NULLs has a null count but no value range, and it
- // belongs wholly to the partition holding the run — which the
- // `nulls_first` end names, not an overlap check.
let entry = ranges.get(&(task_id as usize, sub_part_id));
- let null_only = entry
- .is_some_and(|entry| entry.null_count > 0 && entry.key_min.is_empty());
- if null_only {
- let run = if nulls_first { 0 } else { partition_count - 1 };
- remapped[run].push(file);
- continue;
- }
+ // The NULL run belongs wholly to the partition the `nulls_first`
+ // end names. `key_min`/`key_max` are value extremes, so no
+ // overlap check on them can find that partition — a file whose
+ // values and NULLs belong to different consumers has to be
+ // delivered to both, and is below. Every producer routes its own
+ // NULLs to its own last slot, so without this only the one file
+ // whose values happen to land on the run's partition delivers
+ // them and the rest of the run is silently dropped.
+ let null_part_idx = entry
+ .is_some_and(|entry| entry.null_count > 0)
+ .then(|| if nulls_first { 0 } else { partition_count - 1 });
let range = entry.and_then(|entry| {
let lo = entry.key_min.first()?;
let hi = entry.key_max.first()?;
Some((lo, hi))
});
let Some((min_proto, max_proto)) = range else {
- // No routing info. Safe to skip only if the file has zero rows
- if file.partition_stats.num_rows != Some(0) {
- return internal_err!(
- "range-repartition remap: file has num_rows={:?} but no usable key range (task_id={task_id}, sub_part_id={sub_part_id})",
- file.partition_stats.num_rows
- );
+ // No value range: an all-NULL file, or an empty one. Anything
+ // else has rows this remap cannot address.
+ match null_part_idx {
+ Some(part_idx) => remapped[part_idx].push(file),
+ None if file.partition_stats.num_rows != Some(0) => {
+ return internal_err!(
+ "range-repartition remap: file has num_rows={:?} but no usable key range (task_id={task_id}, sub_part_id={sub_part_id})",
+ file.partition_stats.num_rows
+ );
+ }
+ None => {}
}
continue;
};
@@ -1003,6 +1009,13 @@ pub fn cut_partitions(
let reach_hi = widen_above(&sketch_max, halo_lo)?;
let b_lo = global_cuts.partition_point(|cut| cut <= &reach_lo);
let b_hi = global_cuts.partition_point(|cut| cut <= &reach_hi);
+ // Delivered for its NULL run only when the value range didn't
+ // already reach that partition: two copies in one consumer would
+ // be read twice and counted twice.
+ if let Some(part) = null_part_idx.filter(|part| !(b_lo..=b_hi).contains(part))
+ {
+ remapped[part].push(file.clone());
+ }
for bucket in &mut remapped[b_lo..=b_hi] {
bucket.push(file.clone());
}
@@ -1971,6 +1984,123 @@ mod overlap_remap_tests {
assert_eq!(remapped[1][0].file_id, Some(300));
}
+ /// [`sketch_report`] with a NULL count per sub-part, taken as
+ /// `(values, nulls)` pairs so the two can't fall out of step.
+ fn sketch_report_with_nulls(
+ producer_task_id: usize,
+ per_sub_part: Vec<(Vec<f64>, u64)>,
+ ) -> TaskRuntimeStats {
+ let null_counts: Vec<u64> = per_sub_part.iter().map(|(_, n)| *n).collect();
+ let mut stats = sketch_report(
+ producer_task_id,
+ per_sub_part.into_iter().map(|(values, _)| values).collect(),
+ );
+ for (entry, nulls) in stats.report.partitions.iter_mut().zip(null_counts) {
+ entry.null_count = nulls;
+ entry.row_count += nulls;
+ }
+ stats
+ }
+
+ /// A file holding both values and NULLs belongs to two consumers: the one
+ /// its value range overlaps, and the one holding the NULL run. Reported
+ /// extremes are value extremes, so an overlap check alone never finds the
+ /// second — and every producer writes its own NULLs to its own last slot,
+ /// so dropping it loses all but one producer's share of the run.
+ #[test]
+ fn overlap_remap_mixed_file_reaches_both_its_range_and_the_null_run() {
+ // Values [1, 2] fall left of the cut; the run sorts last, so it
+ // belongs to partition 1.
+ let reports = vec![sketch_report_with_nulls(100, vec![(vec![1.0, 2.0], 2)])];
+ let cuts = cuts_f64([10.0]);
+ let original_partitions = vec![vec![location(0, 100)]];
+
+ let remapped = cut_partitions(
+ original_partitions,
+ &reports,
+ &cuts,
+ &ZERO_HALO,
+ &ZERO_HALO,
+ false,
+ )
+ .unwrap();
+ assert_eq!(remapped[0].len(), 1, "values [1, 2] land left of the cut");
+ assert_eq!(remapped[1].len(), 1, "the NULL run sorts last");
+ }
+
+ /// Same, mirrored: `nulls_first` puts the run in partition 0 while the
+ /// values belong to the last partition.
+ #[test]
+ fn overlap_remap_mixed_file_reaches_the_null_run_at_the_front() {
+ let reports = vec![sketch_report_with_nulls(100, vec![(vec![20.0, 25.0], 2)])];
+ let cuts = cuts_f64([10.0]);
+ let original_partitions = vec![vec![location(0, 100)]];
+
+ let remapped = cut_partitions(
+ original_partitions,
+ &reports,
+ &cuts,
+ &ZERO_HALO,
+ &ZERO_HALO,
+ true,
+ )
+ .unwrap();
+ assert_eq!(remapped[0].len(), 1, "the NULL run sorts first");
+ assert_eq!(
+ remapped[1].len(),
+ 1,
+ "values [20, 25] land right of the cut"
+ );
+ }
+
+ /// When the value range already reaches the partition holding the run,
+ /// the file must be delivered once and not twice — a second copy in one
+ /// consumer is read twice and so counted twice.
+ #[test]
+ fn overlap_remap_mixed_file_in_the_run_partition_is_delivered_once() {
+ let reports = vec![sketch_report_with_nulls(100, vec![(vec![20.0, 25.0], 2)])];
+ let cuts = cuts_f64([10.0]);
+ let original_partitions = vec![vec![location(0, 100)]];
+
+ let remapped = cut_partitions(
+ original_partitions,
+ &reports,
+ &cuts,
+ &ZERO_HALO,
+ &ZERO_HALO,
+ false,
+ )
+ .unwrap();
+ assert!(remapped[0].is_empty(), "no values left of the cut");
+ assert_eq!(
+ remapped[1].len(),
+ 1,
+ "values and run share partition 1, so one delivery"
+ );
+ }
+
+ /// An all-NULL file has a row count but no value range, so the run is the
+ /// only thing placing it.
+ #[test]
+ fn overlap_remap_all_null_file_goes_only_to_the_null_run() {
+ let reports = vec![sketch_report_with_nulls(100, vec![(vec![], 3)])];
+ let cuts = cuts_f64([10.0]);
+ let mut only_nulls = location(0, 100);
+ only_nulls.partition_stats = PartitionStats::new(Some(3), None, None);
+
+ let remapped = cut_partitions(
+ vec![vec![only_nulls]],
+ &reports,
+ &cuts,
+ &ZERO_HALO,
+ &ZERO_HALO,
+ false,
+ )
+ .unwrap();
+ assert!(remapped[0].is_empty(), "no values to place");
+ assert_eq!(remapped[1].len(), 1, "the run sorts last");
+ }
+
/// A producer file without `file_id` means the writer that produced it
/// wasn't the passthrough writer — remap can't identify the file, so
/// error rather than silently misroute.
diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs
index df3cdcab9..99f9a7c5f 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;
@@ -498,6 +498,37 @@ 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.
+ ///
+ /// Errors rather than dropping a capture it cannot place. 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<Vec<(usize, ObservedWindowState)>> {
+ 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
@@ -954,6 +985,23 @@ 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.
+fn collect_window_state_operators<'a>(
+ plan: &'a Arc<dyn ExecutionPlan>,
+ out: &mut Vec<&'a PartitionedBoundedWindowAggExec>,
+) {
+ if let Some(op) = plan.downcast_ref::<PartitionedBoundedWindowAggExec>() {
+ 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..469de39cf
--- /dev/null
+++ b/ballista/core/src/execution_plans/window_state.rs
@@ -0,0 +1,554 @@
+// 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.
+
+//! Module of helpers for 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`.
+
+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.
+/// Preserves all the info from the callback for storage.
+#[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<dyn WindowExpr>`; 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<ScalarValue>,
+}
+
+/// 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<dyn WindowExpr>` 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<Arc<dyn WindowExpr>>,
+ observed: Mutex<Vec<ObservedWindowState>>,
+}
+
+impl WindowStateCollector {
+ /// Build a collector resolving callbacks against `window_expr`.
+ pub fn new(window_expr: Vec<Arc<dyn WindowExpr>>) -> 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<ObservedWindowState> {
+ 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<dyn WindowExpr>,
+ partition_key: &PartitionKey,
+ state: Vec<ScalarValue>,
+ ) -> 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<WindowStateReport> {
+ 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<ScalarValue>, Vec<ScalarValue>)> {
+ 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<Vec<datafusion_proto_common::ScalarValue>> {
+ scalars
+ .iter()
+ .map(datafusion_proto_common::ScalarValue::try_from)
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .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<Vec<ScalarValue>> {
+ proto
+ .iter()
+ .map(ScalarValue::try_from)
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .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<dyn WindowExpr>],
+ partition_count: usize,
+) -> Result<Vec<FinalizedPartitionState>> {
+ // (global partition, window expression) -> that group's finalized state.
+ let mut states: HashMap<(usize, usize), Vec<ScalarValue>> = 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<FinalizedPartitionState> = 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::<PlainAggregateWindowExpr>()
+ 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::<Result<Vec<_>>>()?;
+ 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::array::{ArrayRef, Int64Array};
+ use datafusion::arrow::compute::SortOptions;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::functions_aggregate::approx_distinct::approx_distinct_udaf;
+ use datafusion::functions_aggregate::sum::sum_udaf;
+ use datafusion::logical_expr::{
+ WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
+ };
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+ use datafusion::physical_expr::expressions::{Column, col};
+ use datafusion::physical_plan::PhysicalExpr;
+ use datafusion::physical_plan::windows::create_window_expr;
+ use prost::Message;
+
+ /// Encoded size of one `approx_distinct` state observed over `rows`
+ /// distinct values, as it would ride `SuccessfulTask`.
+ fn approx_distinct_wire_size(rows: i64) -> Result<usize> {
+ let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
+ let column: Arc<dyn PhysicalExpr> = Arc::new(Column::new("v", 0));
+ let aggregate = AggregateExprBuilder::new(approx_distinct_udaf(), vec![column])
+ .schema(Arc::clone(&schema))
+ .alias("wire_size")
+ .build()?;
+ let mut accumulator = aggregate.create_accumulator()?;
+ let values: ArrayRef = Arc::new(Int64Array::from_iter_values(0..rows));
+ accumulator.update_batch(&[values])?;
+ let report = window_state_to_proto(
+ 0,
+ &ObservedWindowState {
+ partition_idx: 0,
+ window_expr_index: 0,
+ partition_key: vec![],
+ state: accumulator.state()?,
+ },
+ )?;
+ Ok(report.encoded_len())
+ }
+
+ /// A sketch state's wire size is fixed by the sketch, not by the rows it
+ /// saw: `approx_distinct` keeps a dense HLL register array. That is what
+ /// makes a task's payload predictable — it is (partitions in the task's
+ /// slice) × (aggregate window expressions) × this, with no term for input
+ /// size.
+ ///
+ /// The ceiling is deliberately loose. It is not a budget, only a tripwire
+ /// for a representation change that grows the payload by an order of
+ /// magnitude.
+ #[test]
+ fn sketch_state_wire_size_does_not_grow_with_rows() -> Result<()> {
+ let one_row = approx_distinct_wire_size(1)?;
+ let many_rows = approx_distinct_wire_size(100_000)?;
+ assert_eq!(
+ one_row, many_rows,
+ "a dense HLL encodes to the same size at any cardinality"
+ );
+ assert!(
+ many_rows < 32 * 1024,
+ "one sketch state encodes to {many_rows} bytes, and a task \
+ carries one per partition in its slice per aggregate window \
+ expression"
+ );
+ Ok(())
+ }
+
+ /// `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<dyn WindowExpr> {
+ 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<Option<f64>> {
+ 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 acc22bd55..dbff74c76 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.
@@ -259,6 +261,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<WindowApplyNode>,
+ /// 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<FinalizedPartitionStateNode>,
+}
+/// 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<window_apply_node::Apply>,
+}
+/// 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<AggregateStateSlotNode>,
+}
+#[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<AggregateStateNode>,
+}
+#[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.
@@ -1040,6 +1119,47 @@ pub struct SuccessfulTask {
/// to combine reports across tasks/executors.
#[prost(message, repeated, tag = "3")]
pub runtime_stats: ::prost::alloc::vec::Vec<RuntimeStatsReport>,
+ /// 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.
+ #[prost(message, repeated, tag = "4")]
+ pub window_state: ::prost::alloc::vec::Vec<WindowStateReport>,
+}
+/// 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.
+ #[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)]
@@ -1640,6 +1760,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<Self> {
+ 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 db28fb077..3017404d6 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -56,11 +56,12 @@ use std::{convert::TryInto, io::Cursor};
use crate::execution_plans::sort_shuffle::SortShuffleConfig;
use crate::execution_plans::{
- BufferExec, BufferMode, ChaosExec, CoalescePlan, InputOrder,
+ BufferExec, BufferMode, ChaosExec, CoalescePlan, FinalizedPartitionState, InputOrder,
OrderedRangeRepartitionExec, PartitionGroup, PartitionedBoundedWindowAggExec,
- PerPartitionFilterExec, RangeFilterExec, RangeShuffleReaderExec, RuntimeStatsExec,
- ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec,
- UnorderedRangeRepartitionExec, UnresolvedShuffleExec,
+ PerPartitionFilterExec, PrefixMergeExec, RangeFilterExec, RangeShuffleReaderExec,
+ RuntimeStatsExec, ScalarOp, ShuffleReaderExec, ShuffleWriterExec,
+ SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec,
+ WindowApply,
};
use crate::serde::protobuf::{
ballista_logical_plan_node::LogicalPlanType,
@@ -68,6 +69,7 @@ use crate::serde::protobuf::{
};
use crate::serde::scheduler::PartitionLocation;
use datafusion::arrow::compute::SortOptions;
+use datafusion::execution::FunctionRegistry;
pub use generated::ballista as protobuf;
/// Generated protobuf code from Ballista protocol definitions.
@@ -367,6 +369,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<Vec<datafusion_proto_common::ScalarValue>, 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<Vec<datafusion::scalar::ScalarValue>, 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<Vec<protobuf::FinalizedPartitionStateNode>, 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<Vec<FinalizedPartitionState>, 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<Vec<protobuf::WindowApplyNode>, 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<Vec<WindowApply>, 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,
@@ -708,6 +895,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!(
@@ -1105,6 +1306,32 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
))
})?;
Ok(())
+ } else if let Some(exec) = node.downcast_ref::<PrefixMergeExec>() {
+ // 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::<RangeFilterExec>() {
let raw_bounds = exec.raw_bounds().ok_or_else(|| {
DataFusionError::Internal(
@@ -1441,6 +1668,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<ScalarValue>` 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<u8> = 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::<PrefixMergeExec>()
+ .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..914dc85d8 100644
--- a/ballista/executor/src/execution_engine.rs
+++ b/ballista/executor/src/execution_engine.rs
@@ -100,6 +100,25 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display {
) -> Vec<ballista_core::serde::protobuf::RuntimeStatsReport> {
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.
+ fn collect_window_state_reports(
+ &self,
+ ) -> Result<Vec<ballista_core::serde::protobuf::WindowStateReport>> {
+ Ok(Vec::new())
+ }
}
/// Default execution engine using DataFusion's ShuffleWriterExec.
@@ -340,6 +359,27 @@ impl QueryStageExecutor for DefaultQueryStageExec {
}
}
}
+
+ fn collect_window_state_reports(
+ &self,
+ ) -> Result<Vec<ballista_core::serde::protobuf::WindowStateReport>> {
+ // 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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
.collect::<Result<Vec<_>, 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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
TaskCompletionExtras {
operator_metrics,
runtime_stats,
+ window_state,
},
));
diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs
index 0d56417e8..6c009ce2c 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -460,6 +460,18 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
.collect::<Result<Vec<_>, 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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
TaskCompletionExtras {
operator_metrics,
runtime_stats,
+ window_state,
},
);
diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs
index cfc8ad808..9499005a0 100644
--- a/ballista/executor/src/lib.rs
+++ b/ballista/executor/src/lib.rs
@@ -64,7 +64,7 @@ use log::info;
use crate::shutdown::Shutdown;
use ballista_core::serde::protobuf::{
FailedTask, OperatorMetricsSet, RuntimeStatsReport, ShuffleWritePartition,
- SuccessfulTask, TaskStatus, task_status,
+ SuccessfulTask, TaskStatus, WindowStateReport, task_status,
};
use ballista_core::serde::scheduler::TaskKey;
use ballista_core::utils::GrpcServerConfig;
@@ -111,6 +111,10 @@ pub struct TaskCompletionExtras {
pub operator_metrics: Option<Vec<OperatorMetricsSet>>,
/// Runtime-stats reports harvested from `RuntimeStatsExec` taps in the plan.
pub runtime_stats: Vec<RuntimeStatsReport>,
+ /// 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<WindowStateReport>,
}
/// 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 2e8684dbf..e74d0939d 100644
--- a/ballista/scheduler/src/scheduler_server/mod.rs
+++ b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -688,6 +688,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 05ff59274..c01e84211 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -870,6 +870,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 89a654043..eae9380b9 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};
@@ -359,6 +361,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::<PrefixMergeExec>() 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,
@@ -378,6 +439,7 @@ impl AdaptiveExecutionGraph {
debug!("ignoring completion for superseded stage {stage_id}");
return Ok(HashSet::new());
}
+ 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
@@ -913,6 +975,7 @@ impl ExecutionGraph for AdaptiveExecutionGraph {
let ballista_core::serde::protobuf::SuccessfulTask {
partitions,
runtime_stats,
+ window_state,
..
} = successful_task;
debug!(
@@ -924,6 +987,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(
@@ -1524,3 +1589,45 @@ fn downstream_halos(
))
})
}
+
+/// 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<dyn ExecutionPlan>) -> Option<usize> {
+ let mut node = Arc::clone(start);
+ loop {
+ if let Some(exchange) = node.downcast_ref::<ExchangeExec>() {
+ 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.
+fn descend_to_window_expr(
+ start: &Arc<dyn ExecutionPlan>,
+) -> Option<Vec<Arc<dyn WindowExpr>>> {
+ let mut found = None;
+ start
+ .apply(|node| {
+ if let Some(op) = node.downcast_ref::<PartitionedBoundedWindowAggExec>() {
+ 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 431e6ae3c..b21eed099 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
@@ -22,6 +22,7 @@ pub mod distributed_exchange;
pub mod join_selection;
pub mod normalize_interleave;
pub mod parallel_window;
+pub mod prefix_window;
pub mod propagate_empty;
pub use coalesce_partitions::*;
@@ -30,4 +31,5 @@ pub use distributed_exchange::*;
pub use join_selection::*;
pub use normalize_interleave::*;
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..23abfbf69
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/prefix_window.rs
@@ -0,0 +1,547 @@
+// 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`].
+//!
+//! # The shape it gates on
+//!
+//! No PARTITION BY, a single ascending ORDER BY column, and an ever-expanding
+//! frame. The ORDER BY column's type is gated by asking
+//! [`SortKeyCodec::try_new`] rather than by a type list here, because that
+//! codec is what ORRE's routing sketch encodes keys with — widening the codec
+//! widens this rule.
+//!
+//! # Actual rule input
+//!
+//! Captured by logging `optimize`'s argument on h2o q7. 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]
+//! ```
+//!
+//! # 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. 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, to exchange prefixes with scheduler
+//! 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]
+//! <source>
+//! ```
+//!
+//! ## 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.
+
+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::{
+ InputOrder, OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec,
+ PrefixMergeExec, RangeFilterExec, RuntimeStatsExec, WindowApply,
+};
+use ballista_core::sort_key::SortKeyCodec;
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::physical_expr::LexOrdering;
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::window::PlainAggregateWindowExpr;
+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.
+#[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<AtomicUsize>,
+}
+
+impl PrefixWindowRule {
+ pub(crate) fn new(plan_id_generator: Arc<AtomicUsize>) -> Self {
+ Self { plan_id_generator }
+ }
+}
+
+impl PhysicalOptimizerRule for PrefixWindowRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ let bc = config
+ .extensions
+ .get::<BallistaConfig>()
+ .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<dyn ExecutionPlan>]) -> bool {
+ for node in nodes {
+ if node.is::<OrderedRangeRepartitionExec>()
+ || node.is::<RangeFilterExec>()
+ || node.is::<PrefixMergeExec>()
+ {
+ return true;
+ }
+ if subtree_contains_our_rewrite(node.children().as_slice()) {
+ return true;
+ }
+ }
+ 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 → <source>` 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<dyn ExecutionPlan>,
+ output_partitions: usize,
+ plan_id_generator: &Arc<AtomicUsize>,
+) -> datafusion::common::Result<Option<Arc<dyn ExecutionPlan>>> {
+ let Some(window) = node.downcast_ref::<BoundedWindowAggExec>() 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::<Column>() 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<dyn ExecutionPlan> = (*immediate).clone();
+ while base_source.is::<SortPreservingMergeExec>() || base_source.is::<SortExec>() {
+ 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 via a `SortKeySketch`, so the gate
+ // is whatever that sketch's codec can encode — asked here rather than
+ // spelled as a type list, so widening the codec widens this rule with no
+ // edit. `SortKeyCodec` gives NULLs a position per `nulls_first`, so a
+ // nullable key needs no separate gate.
+ let routing_type = order.expr.data_type(&source_schema)?;
+ if SortKeyCodec::try_new(&routing_type, order.options).is_none() {
+ return Ok(None);
+ }
+
+ let sort_expr = order.clone();
+
+ // 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(sort_lex, rse1).with_preserve_partitioning(true));
+ let orre: Arc<dyn ExecutionPlan> = Arc::new(OrderedRangeRepartitionExec::try_new(
+ sorted_over_rse1,
+ vec![sort_expr.clone()],
+ output_partitions,
+ )?);
+ let rse2: Arc<dyn ExecutionPlan> = Arc::new(RuntimeStatsExec::try_new(
+ orre,
+ Some(vec![sort_expr.clone()]),
+ )?);
+ let trim: Arc<dyn ExecutionPlan> = Arc::new(RangeFilterExec::try_new_pending(
+ rse2,
+ sort_expr.expr.clone(),
+ ScalarValue::new_zero(&routing_type)?,
+ ScalarValue::new_zero(&routing_type)?,
+ Some(InputOrder::Ordered(sort_expr.options)),
+ )?);
+ // 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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.
+ //
+ // For now 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::<PlainAggregateWindowExpr>() 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<dyn ExecutionPlan> =
+ 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::{DataType, 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<Arc<dyn ExecutionPlan>> {
+ 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),
+ // Not in h2o's schema. The sketch codec places a NULL run per
+ // `nulls_first`, so nullability is not part of the type gate —
+ // `rewrites_nullable_routing` is what holds that open.
+ Field::new("nullable_key", DataType::Int64, true),
+ ]));
+ 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<dyn ExecutionPlan>,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ 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.
+ ///
+ /// 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, ordering by `Int64 id3`. The rewrite must not be
+ /// specific to the `Float64` key `rewrites_unbounded_preceding_shape`
+ /// uses: the gate is what the routing sketch's codec encodes, and that
+ /// covers every fixed-width type.
+ #[tokio::test]
+ async fn rewrites_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"),
+ "an Int64 routing key is one the sketch codec encodes:\n{rendered}"
+ );
+ Ok(())
+ }
+
+ /// A nullable routing key. NULLs do not break the prefix: they sort as one
+ /// run to the end `nulls_first` names, so the run lands wholly inside one
+ /// partition and "every partition before k" stays well defined.
+ #[tokio::test]
+ async fn rewrites_nullable_routing() -> datafusion::common::Result<()> {
+ let plan = plan(
+ "SELECT sum(v2) OVER (ORDER BY nullable_key \
+ 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"),
+ "a nullable routing key is routable — the codec gives the NULL \
+ run a position:\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 0f0e41cd4..2a01df0f9 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -23,7 +23,7 @@ use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule;
use crate::state::aqe::optimizer_rule::{
CoalescePartitionsRule, DelayJoinSelectionRule, DemoteUnsafeBroadcastJoinRule,
DistributedExchangeRule, NormalizeInterleaveRule, ParallelWindowRule,
- PropagateEmptyExecRule, SelectJoinRule,
+ PrefixWindowRule, PropagateEmptyExecRule, SelectJoinRule,
};
use crate::state::distributed_explain::handle_explain_plan;
use crate::state::execution_stage::StageOutput;
@@ -573,6 +573,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 220a2b38a..1ee1e191a 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<TaskRuntimeStats>,
+ /// 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<TaskWindowState>,
}
/// 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,30 @@ 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<WindowStateReport>,
+ ) {
+ for report in &reports {
+ 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 +1077,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 +1118,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 +1226,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 +1447,7 @@ mod tests {
executor_id: "executor-1".to_string(),
partitions: vec![],
runtime_stats: vec![],
+ window_state: vec![],
})),
metrics: vec![],
}
@@ -1783,6 +1826,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.
@@ -1805,6 +1883,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 879d62e35..eeac4f039 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,42 +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::<RangeFilterExec>() {
+ // 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::<datafusion::common::Result<_>>()?;
- return Ok(Arc::new(RangeFilterExec::try_new_resolved(
- new_child,
- rf.filter_expr().clone(),
- rf.halo_lo().clone(),
- rf.halo_hi().clone(),
- rf.input_order(),
- 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 25604673b..57eb51918 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![],
})),
});
}
@@ -1269,6 +1270,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![],
})),
}
}
diff --git a/chaos-testing/src/cluster.rs b/chaos-testing/src/cluster.rs
index d1972cd57..3a19efad5 100644
--- a/chaos-testing/src/cluster.rs
+++ b/chaos-testing/src/cluster.rs
@@ -24,13 +24,24 @@ use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, OwnedMutexGuard};
-/// Reserve a free TCP port by binding to :0 and immediately releasing it.
+/// Reserve `N` free TCP ports by binding each to :0, releasing them only once
+/// every port has been read.
///
-/// Inherently racy, but adequate here: the child binds within milliseconds and
-/// the tests are the only thing running.
-fn free_port() -> u16 {
- let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
- listener.local_addr().expect("local addr").port()
+/// Holding all `N` listeners is what makes the ports distinct from each other.
+/// Releasing each before taking the next lets the kernel hand the same number
+/// back, and a caller that asked for two ports then gives its child one port
+/// for two servers: the second bind fails `AddrInUse` and the child exits
+/// during startup.
+///
+/// Still racy against the rest of the machine, which is adequate here: the
+/// child binds within milliseconds and the tests are the only thing running.
+fn free_ports<const N: usize>() -> [u16; N] {
+ let listeners: [TcpListener; N] = std::array::from_fn(|_| {
+ TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port")
+ });
+ listeners
+ .each_ref()
+ .map(|listener| listener.local_addr().expect("local addr").port())
}
/// Open a child process log file for append.
@@ -163,7 +174,7 @@ impl TestClusterBuilder {
let temp = tempfile::tempdir().map_err(|e| e.to_string())?;
let log_dir = temp.path().join("logs");
std::fs::create_dir_all(&log_dir).map_err(|e| e.to_string())?;
- let scheduler_port = free_port();
+ let [scheduler_port] = free_ports();
let scheduler_log = log_dir.join("scheduler.log");
let scheduler_stdout = open_log(&scheduler_log).map_err(|e| {
@@ -304,8 +315,9 @@ impl TestCluster {
}
pub(crate) fn spawn_executor(&mut self, index: usize) -> Result<(), String> {
- let port = free_port();
- let grpc_port = free_port();
+ // One call, so the two are guaranteed distinct: the executor binds
+ // its flight server to `port` and its gRPC server to `grpc_port`.
+ let [port, grpc_port] = free_ports();
let work_dir = self.temp.path().join(format!("executor-{index}"));
std::fs::create_dir_all(&work_dir).map_err(|e| e.to_string())?;
@@ -789,6 +801,24 @@ impl Drop for TestCluster {
mod tests {
use super::*;
+ /// Ports handed out together must differ. An executor gets two of them, one
+ /// for its flight server and one for its gRPC server, and a repeat makes
+ /// the second bind fail `AddrInUse` and kills the executor during startup,
+ /// which surfaces as the whole cluster timing out on registration.
+ ///
+ /// Guards the requirement rather than reproducing a failure: holding every
+ /// listener makes this hold by construction, since the OS will not assign
+ /// one port to two live sockets. It is here to fail if anyone returns to
+ /// reserving the ports one at a time, which is what allows a repeat.
+ #[test]
+ fn ports_reserved_together_are_distinct() {
+ let ports = free_ports::<16>();
+ let mut seen = std::collections::HashSet::new();
+ for port in ports {
+ assert!(seen.insert(port), "port {port} handed out twice: {ports:?}");
+ }
+ }
+
#[tokio::test]
async fn cluster_starts_with_the_requested_executors_registered() {
let cluster = TestCluster::builder()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org