feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates - #2211
feat: parallel prefix scan for UNBOUNDED PRECEDING window aggregates#2211avantgardnerio wants to merge 5 commits into
Conversation
|
CI is red on two jobs here, but only one of them is yours.
The
Design feedback coming in a separate comment. |
|
Design feedback, separate from the CI note above. The overall shape reads well to me, and the module header is genuinely good documentation. Splitting the global prefix merge onto the scheduler and leaving a row wise apply on the executor is the right decomposition, and the APPROX_DISTINCT case makes a convincing argument for why the The per row accumulator replay looks like a performance trap. Serde is deferred, and it's the hard part. #2255 landed its proto message and codec arm in the same PR, and I'd like this one to end up there too, since The related question is the transport the design picks. The description says upstream state reaches the scheduler over task status. That's a hot and frequent message, and HLL, KLL or TDigest state per task per window expression is not small. Since you describe the division of labor as fixed at design time, I'd rather pressure test that choice now than after the scheduler side is built on top of it. Partition index coupling has no guard. Both Smaller things:
|
906eacf to
5a5f695
Compare
…e 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) <noreply@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) <noreply@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) <noreply@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) <noreply@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) <noreply@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) <noreply@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) <noreply@anthropic.com>
8696f35 to
f85b988
Compare
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) <noreply@anthropic.com>
avantgardnerio
changed the title
| for i in 0..num_rows { | ||
| let row_args: Vec<ArrayRef> = | ||
| arg_arrays.iter().map(|a| a.slice(i, 1)).collect(); | ||
| self.accumulator.update_batch(&row_args)?; | ||
| new_values.push(self.accumulator.evaluate()?); | ||
| } |
There was a problem hiding this comment.
it re derives work the upstream BWAG already did
Yes, and that is exactly what the second pass is doing: the upstream window already advanced an accumulator through every row, and this walks it again from the carry-in.
The reason it is not avoidable today is where the state is exposed. apache/datafusion#24035 publishes accumulator state at PARTITION BY group close, one state per group, so there is no per row state to carry into a batch at a time correction.
there may be a middle path where the upstream emits partial state columns
I think that is worth pursuing, but it needs more than the API. DataFusion's dense HLL is 16 KiB per sketch (approx_distinct.rs#L257), so materializing per row state for a sketch is 16 KiB times row count as a column. A workable version probably needs a cheaper per row representation than raw accumulator state, which makes it a real piece of design rather than a switch to flip.
Worth being precise about the current state: this PR does not select the scalar path yet. WindowApply::Scalar exists on the operator and applies batch at a time via arrow kernels, but the rule currently builds an Aggregate apply for every window expression, including SUM. Selecting Scalar where it applies will land on this PR before merge, and for SUM, COUNT, MIN and MAX it removes this cost rather than halving it.
Could you run it over a realistic partition
Yes. A three row test does not surface any of this. The trade we are making is 2x local work divided across cores (theoretically 16x on 32, to be measured shortly). Benchmarks will land in this thread before anything else is built on the shape.
| // 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. |
There was a problem hiding this comment.
HLL, KLL or TDigest state per task per window expression is not small
Agreed, and it is worth sizing. DataFusion's dense HLL is 16 KiB per sketch (approx_distinct.rs#L257). A task carries one state per (partition in its slice, aggregate window expression), so max_partitions_per_task multiplies it: a 32 partition slice with a single sketch aggregate puts 512 KiB on that task's completion message.
Two things make me think the current shape is defensible as a starting point rather than a commitment. RuntimeStatsExec already ships quantile sketches over this exact path, so this adds a second payload of a class task status already carries rather than introducing one. And the escape hatch is cheap: the state can be written as a sidecar beside the shuffle files, the way sort shuffle already writes <data>.arrow.index next to its data, with only a reference on the message. That keeps completion fixed size regardless of aggregate, and it is a change to this one field rather than to the design around it.
What I have not done is measure it. The e2e only exercises a Float64 SUM, where the payload is a handful of bytes. If you would rather see a real number for approx_distinct before this merges, I can add a test that reports the encoded size and post it here.
| 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.routing_expr().clone(), | ||
| rf.halo_lo().clone(), | ||
| rf.halo_hi().clone(), | ||
| sliced_bounds, | ||
| )?)); | ||
| return sliceable.slice_to_partitions(new_child, partitions); | ||
| } |
There was a problem hiding this comment.
any rule that repartitions the input to the same partition count would silently attach each partition's offsets to the wrong rows
This was real and it fired on the first end to end run, though not by the route you predicted. Task specialization already restricts a stage plan to one task's partition slice, so PrefixMergeExec was rebuilt over a narrower input while still holding state keyed by global partition index. It failed loudly only because try_new validates the length: per_partition_state.len() 4 does not match input partition count 1. Without that check it is exactly the silent misattachment you describe, since a task's execute(k) numbers its own slice from zero.
Fixed by slicing state and Scalar.offset parallel to the input restriction, following the precedent RangeFilterExec already set for its bounds. It is behind a PartitionSliceable trait so each operator implements the slicing next to the fields being sliced, rather than the task builder knowing their internals: add a partition indexed field and the code that has to slice it is in the same file.
The e2e runs with max_partitions_per_task = 2, so tasks carry a genuine multi partition slice rather than the degenerate one partition case.
| 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)?, | ||
| )?)) |
There was a problem hiding this comment.
Arc<AggregateUDF>,Vec<Arc<dyn PhysicalExpr>>, andVec<ScalarValue>sketch state all have to cross the wire
Done on this PR. The UDAF crosses by name and is resolved from the executor's FunctionRegistry on decode, args as PhysicalExprNode, and state as ScalarValue so sketch state works unchanged. A round trip test covers both WindowApply shapes, the registry lookup, and the distinction between a slot with no state and a slot with empty state.
Encoding refuses while the prefix state is unresolved, matching RangeFilterExec's refusal on unresolved bounds. An executor has no route to that state, so a plan reaching the wire without it could only emit partition-local aggregates.
…ype 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) <noreply@anthropic.com>
fe70013 to
7ae3ca8
Compare
| /// `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, |
There was a problem hiding this comment.
Which frame is that aimed at?
None, you are right. Over an ever-expanding frame last_value is the current row's own value and needs no correction, so overwriting every row with a single scalar would be wrong. The doc now claims first_value only and says why last_value is excluded, so the next reader does not have to re-derive it.
| /// `#[non_exhaustive]`: the two shapes here cover what the prefix rewrite | ||
| /// plants today, and a third (a segment-tree broadcast for the ranking | ||
| /// family) is anticipated. | ||
| #[derive(Debug, Clone)] | ||
| #[non_exhaustive] |
There was a problem hiding this comment.
Marking both
#[non_exhaustive]now costs nothing and saves a breaking change on the first new variant
Done, on both. Construction is unaffected so the rewrite rule still builds an Aggregate apply directly; it is exhaustive matching that now needs a wildcard, which is the part that would have broken.
| fn rebuild_batch( | ||
| schema: SchemaRef, | ||
| columns: Vec<ArrayRef>, | ||
| apply_index: usize, | ||
| output_column: usize, | ||
| ) -> Result<RecordBatch> { | ||
| let replaced = columns[output_column].data_type().clone(); |
There was a problem hiding this comment.
RecordBatch::try_newfails with an opaque arrow error rather than something naming the offendingapplies[i]
Fixed. Both apply paths now rebuild through this helper, which reports applies[3] produced Int64 for column 1, which the schema declares as Float64. AggregateApply already carried apply_index for exactly this and was not using it at the failure site. There is a test on the message, since it is otherwise an untested branch.
avantgardnerio
changed the title
|
Heads up that this is functional now rather than a scaffold. It runs end to end and Leaving it in draft: benchmarks for the per-row replay are not in yet, and the scalar path selection lands before merge. Both are tracked in the threads. |
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) <noreply@anthropic.com>
| struct ApplyStream { | ||
| input: SendableRecordBatchStream, | ||
| appliers: Vec<PreparedApply>, | ||
| schema: SchemaRef, | ||
| baseline: BaselineMetrics, | ||
| path_metrics: PathMetrics, | ||
| } |
There was a problem hiding this comment.
PrefixMergeExecdoesn't implementmetrics()andApplyStreamhas noBaselineMetrics
Added. BaselineMetrics for elapsed_compute and output_rows, per partition since execute(partition) builds a stream each, plus a rows_corrected counter.
The two apply paths are timed separately rather than together. WindowApply::Scalar is an arrow kernel over a whole batch and WindowApply::Aggregate seeds an accumulator and replays every row, so a single number would cover two quite different costs for a query that mixes them. It also makes the replay you flagged something the operator reports in production rather than something only a benchmark can see.
I will post these numbers along with the benchmarking results.
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) <noreply@anthropic.com>
| pub struct FinalizedPartitionState { | ||
| /// Indexed by position in the upstream operator's `window_expr()` list. | ||
| per_window_expr: Vec<Option<Vec<ScalarValue>>>, | ||
| } |
There was a problem hiding this comment.
A newtype now would keep that swap internal
Done, though not for the reason you gave, and the difference seems worth recording.
apache/datafusion#24007 did not land. What shipped in DataFusion 55 is apache/datafusion#24035, a WindowStateObserver callback rather than a getter, and it publishes no type equivalent to this one. So there is no upstream type to swap to and the break you were guarding against cannot happen.
The newtype earns its place anyway. Vec<Option<Vec<ScalarValue>>> appears in this operator's public signatures, where it is neither readable nor searchable, and it spells "no state for this window expression" two ways: a missing index, and a None at a present index. Every caller had to handle both. slot(window_expr_index) now collapses them into one answer, and a later change to the representation stays internal.
While looking at this I removed two comments claiming the DataFusion side guarantees at most one PARTITION BY group per partition. It does not, and never did — the callback is keyed by group, so that invariant is ours. The scheduler now enforces it by rejecting any report that carries a key rather than flattening two groups together. Worth noting a window that does have a PARTITION BY needs nothing from this operator: BoundedWindowAggExec asks for KeyPartitioned input, so each partition's window is already independent and there is no serial bottleneck to remove.
2 participants
Summary
BoundedWindowAggExecdeclaresSinglePartitionwhen there is no PARTITION BY, because a window frame spans rows across the whole input. For an ever-expanding frame that means the entire window computation runs on one core, however wide the cluster.Measured on h2o
window.sqlQ7 at 1e7, 8 partitions, 2 executors:elapsed_computeSortPreservingMergeExeccollapses 8 to 1, then the windowThis PR removes the collapse. Each task computes a partition-local running aggregate over a range-disjoint slice, and a downstream operator corrects it with the merged accumulator state of every prior partition. A halo cannot cover this case, since an
UNBOUNDED PRECEDINGframe reaches back to the first row of the dataset.Shape
Boundary 2 carries no repartition. It exists so the scheduler has a point where every upstream task has published its state.
State flow
DataFusion 55 exposes accumulator state through a callback (apache/datafusion#24035). A
WindowStateCollectorcatches it,ShuffleWriterExectranslates each capture's task-local partition index to a stage-global one, and it ridesSuccessfulTaskto the scheduler. The scheduler prefix-merges the reports into one carry-in per partition and binds them to the downstreamPrefixMergeExec.Merging goes through the aggregate's own
merge_batch, so non-decomposable aggregates work: twoapprox_distinctHLL sketches combine correctly where two distinct counts could not.Correctness
ballista/client/tests/prefix_window.rsruns a distributed running sum against an independently computed oracle through the real scheduler, shuffle and executor, withmax_partitions_per_task = 2so tasks carry genuine multi-partition slices.Safety
Gated behind
ballista.planner.parallel_window.enabled, off by default. The rewrite fires only on: no PARTITION BY, a single ascending ORDER BY column,Float64routing, and an ever-expanding frame. Everything else falls through untouched, including the existing halo rewrite, whose sliding frames DataFusion refuses to publish state for.Known limits
Float64routing only.OrderedRangeRepartitionExecroutes on a T-Digest, which is Float64-only until the KLL migration. h2o Q7 orders byid3(Int64) and so cannot take this path yet — the blocker is the sketch, not this rule.WindowApply::Scalaris implemented but not yet selected; every apply currently routes through the accumulator path.Still draft
Benchmarks are not in yet, and the scalar-path selection lands before merge. Review threads on the earlier feedback are open inline.