avantgardnerio commented on code in PR #2211: URL: https://github.com/apache/datafusion-ballista/pull/2211#discussion_r3786288542
########## ballista/core/src/execution_plans/prefix_merge.rs: ########## @@ -0,0 +1,1481 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cross-partition state merge for windowed aggregates in an AQE range-shuffle +//! pipeline. +//! +//! Range-shuffle produces `N` ordered disjoint partitions of a stream sorted +//! by the window's ORDER BY key. Each executor task then runs a windowed +//! aggregate over its slice, producing correct per-row values *within* the +//! slice but not across slices — task `k`'s running SUM starts at zero, not +//! at the sum of everything in partitions `[0..k)`. +//! +//! # Division of labor +//! +//! The prefix-merge splits cleanly across the scheduler/executor boundary: +//! +//! - **Scheduler (global step):** collects each upstream task's finalized +//! [`Accumulator::state`] via task-status transport, then computes the +//! prefix-merge — for each partition `k`, combining the individual states +//! of partitions `[0..k)` into a single already-merged state per window +//! expression. This is the step that requires global visibility across +//! tasks; only the scheduler has it. +//! - **Executor (local step, this operator):** receives the *already-merged* +//! state for its partition in its constructor and folds it row-wise into +//! the window-aggregate columns. Aggregate-agnostic by construction — the +//! fold is the same [`Accumulator::merge_batch`]-shaped composition the +//! group-by `Partial → Final` protocol uses, so SUM adds, MIN/MAX take the +//! extreme, sketches (KLL, TDigest, HLL) merge as sketches, without this +//! operator knowing which aggregate is which. +//! +//! # Apply descriptors +//! +//! `try_new` takes a `Vec<`[`WindowApply`]`>` — one entry per output column +//! that needs cross-partition correction, telling the operator *how* to +//! rewrite that column. Two shapes: +//! +//! - [`WindowApply::Scalar`] — fast path. Combines each row's existing value +//! with a scheduler-provided scalar via [`ScalarOp`] (`Add`/`Min`/`Max` for +//! SUM/COUNT/MIN/MAX and `row_number`; `Overwrite` for `first_value` / +//! `last_value`). No `Accumulator` constructed. +//! - [`WindowApply::Aggregate`] — fallback. Constructs a fresh `Accumulator` +//! seeded from the pre-merged state, feeds `args` per row, overwrites the +//! column with `evaluate()`. Fits AVG (without decomposition), sketch-backed +//! windows (APPROX_DISTINCT, APPROX_QUANTILE), and statistical aggregates. +//! +//! Non-corrected window functions don't appear in `applies` at all. `lead` / +//! `lag` / `nth_value` are solved by halo rows in the shuffle layer. +//! `rank` / `dense_rank` / `percent_rank` / `cume_dist` / `ntile` need a +//! separate segment-tree-plus-broadcast design and are out of scope here. +//! +//! # Prefix-state input +//! +//! [`FinalizedPartitionState`] — one entry per input partition, holding +//! the *pre-merged* [`Accumulator::state`] for each aggregate window +//! expression (indexed by position in `BoundedWindowAggExec::window_expr()`). +//! Only consumed by [`WindowApply::Aggregate`] entries; +//! [`WindowApply::Scalar`] carries its own offsets inline. No PARTITION BY +//! dimension is exposed: DataFusion publishes state per closed group, but the +//! rewrite that plants this operator only fires on windows without a +//! PARTITION BY, and the scheduler rejects any report carrying a key rather +//! than flattening two groups together. A window that does have a PARTITION +//! BY needs no help from this operator, since `BoundedWindowAggExec` asks for +//! `KeyPartitioned` input and each partition's window is already independent. +//! The scheduler bakes the state when it constructs the downstream stage +//! after the upstream stage's tasks complete. +//! +//! **Status.** Both apply paths are implemented. +//! +//! - [`WindowApply::Aggregate`] builds a fresh `Accumulator` per partition, +//! seeds it via `merge_batch` from the offset state, and replays each row +//! through `update_batch` + `evaluate` to overwrite the output column. +//! - [`WindowApply::Scalar`] applies the [`ScalarOp`] batch-at-a-time via +//! arrow kernels — `numeric::add` for `Add`, `cmp::lt_eq`/`gt_eq` + `zip` +//! for `Min`/`Max`, and a constant-fill for `Overwrite`. +//! +//! Inherits the DF-side getter's at-most-one-PARTITION-BY-group scoping +//! (matches the AQE synthetic-PARTITION-BY pattern) — multi-key queries +//! are handled by DataFusion's normal key-partitioned distribution and +//! don't route through this operator. +//! +//! # Relation to DataFusion +//! +//! The upstream tasks' finalized state — which the scheduler prefix-merges +//! before handing the result to this operator — comes from the accumulators +//! inside `BoundedWindowAggExec`, captured by a `WindowStateObserver` and +//! shipped to the scheduler on task completion. See +//! [`window_state`](super::window_state) for that half. +//! +//! An empty `applies` list makes this operator a passthrough. A non-empty one +//! always rewrites its output columns, including when a state slot is `None` — +//! that seeds the accumulator with nothing rather than skipping the column, so +//! the result is the partition-local aggregate, not the input value. +//! +//! [`Accumulator::state`]: datafusion::logical_expr::Accumulator::state +//! [`Accumulator::merge_batch`]: datafusion::logical_expr::Accumulator::merge_batch + +use std::fmt::{self, Debug, Formatter}; +use std::sync::Arc; + +use log::debug; +use parking_lot::Mutex; + +use datafusion::arrow::array::{ArrayRef, RecordBatch}; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Result, ScalarValue, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{Accumulator, AggregateUDF}; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::{ + ColumnarValue, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, StatisticsArgs, + apply_expression_roots, statistics::ChildStats, +}; +use futures::{Stream, StreamExt, ready}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use crate::execution_plans::plan_algebra::{ + PartitionSliceable, slice_by_global_partition, +}; + +/// A single already-prefix-merged window-aggregate state. Indexed by window +/// expression (same order the upstream `BoundedWindowAggExec` reports in +/// `window_expr()`); `None` at a slot indicates a non-aggregate window +/// function (`row_number`, `rank`, `lead`/`lag`, ...) that contributes no +/// state. The inner `Vec<ScalarValue>` is whatever `Accumulator::state()` +/// returned for that aggregate (1 for SUM/COUNT/MIN/MAX, 2 for AVG's +/// `(sum, count)`, N for sketch-backed / higher-moment aggregates). +/// +/// The scheduler produces one of these per input partition, having already +/// combined the individual states from every prior partition into one merged +/// value per window expression — see the `prefix_merge` module docs for the +/// division of labor. This operator applies it; it does not compute it. +/// +/// A newtype rather than an alias for `Vec<Option<Vec<ScalarValue>>>`. That +/// type appears in this operator's public signatures, where it is neither +/// readable nor searchable, and it spells "no state for this window +/// expression" two ways — a missing index and a `None` — which every caller +/// then has to handle. [`Self::slot`] collapses both into one answer. Any +/// later change to the representation also stays internal rather than +/// breaking whoever wrote the concrete type. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct FinalizedPartitionState { + /// Indexed by position in the upstream operator's `window_expr()` list. + per_window_expr: Vec<Option<Vec<ScalarValue>>>, +} Review 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
