avantgardnerio commented on code in PR #2211: URL: https://github.com/apache/datafusion-ballista/pull/2211#discussion_r3785952837
########## ballista/core/src/execution_plans/prefix_merge.rs: ########## @@ -0,0 +1,1259 @@ +// 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. The DF-side getter +//! already commits to at-most-one PARTITION BY group per DF partition (see +//! `apache/datafusion#24007`), so no per-key dimension is exposed here. 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::{ + 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. +/// +/// The type mirrors `datafusion::physical_plan::windows::FinalizedPartitionState` +/// from [apache/datafusion#24007]; the alias here is a local stand-in so +/// this crate compiles against stable DataFusion 54 until that PR lands. +/// +/// [apache/datafusion#24007]: https://github.com/apache/datafusion/pull/24007 +pub type FinalizedPartitionState = Vec<Option<Vec<ScalarValue>>>; + +/// How to combine each row's existing value in an output column with a +/// scheduler-provided scalar offset. The result overwrites the column. +/// +/// [`Overwrite`] ignores the row's existing value and just writes the offset; +/// it's the shape needed for `first_value` / `last_value`, where the scheduler +/// picks the correct global value once and every row gets a copy. +/// +/// [`Overwrite`]: ScalarOp::Overwrite +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +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` / + /// `last_value`, where the scheduler picks a single global scalar and + /// every row gets the same corrected value. + Overwrite, +} + +/// How to correct one window-function output column at row-apply time. Each +/// entry describes exactly one column PrefixMergeExec should rewrite. +/// +/// Two shapes are covered by construction; anything else is out of scope +/// (`lead`/`lag`/`nth_value` are solved by halo rows in the shuffle layer, and +/// ranking-family functions like `rank`/`percent_rank`/`ntile` want a separate +/// segment-tree-plus-broadcast infrastructure). +#[derive(Debug, Clone)] +pub enum WindowApply { + /// Fast path: monoidal op between each row's existing value and a + /// scheduler-provided scalar. No `Accumulator` is constructed. + /// + /// Fits every aggregate whose per-row output is itself a valid partial + /// state that composes via a single scalar op — SUM, COUNT, MIN, MAX — + /// plus ranking functions like `row_number` (offset = prior row count) + /// and value-selection functions like `first_value` / `last_value` + /// (offset = scheduler-picked global value; op = [`ScalarOp::Overwrite`]). + Scalar { + /// Combining op between `row_value` and `offset`. + op: ScalarOp, + /// One scalar per input partition. `offset[k]` combines with every + /// row passing through partition `k`. Length must match the input's + /// output partition count. + offset: Vec<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. + /// + /// [`Accumulator::merge_batch`]: datafusion::logical_expr::Accumulator::merge_batch + /// [`Accumulator::evaluate`]: datafusion::logical_expr::Accumulator::evaluate + Aggregate { + /// UDF used to construct a fresh `Accumulator` per partition. + udf: Arc<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, + }, +} + +impl WindowApply { + fn output_column(&self) -> usize { + match self { + WindowApply::Scalar { output_column, .. } + | WindowApply::Aggregate { output_column, .. } => *output_column, + } + } +} + +/// Apply pre-merged window-aggregate state (computed by the scheduler) to +/// each row of the current partition's output. See the `prefix_merge` module +/// docs for the division of labor between scheduler and executor and the AQE +/// pipeline this fits into. +/// +/// [`WindowApply::Aggregate`] entries are applied per row via a seeded +/// `Accumulator`; [`WindowApply::Scalar`] entries are applied per batch via +/// arrow kernels. +pub struct PrefixMergeExec { + input: Arc<dyn ExecutionPlan>, + /// One entry per window-function output column that needs cross-partition + /// correction. Non-corrected columns (e.g. `lead`/`lag` handled by halos, + /// or ranking functions left to segment-tree infrastructure) don't appear + /// here. + applies: Vec<WindowApply>, + /// `per_partition_state[k]` is the *already-merged* state summarising + /// every input partition in `[0..k)`. Only consumed by + /// [`WindowApply::Aggregate`] entries — [`WindowApply::Scalar`] carries + /// its own offsets. Length equals + /// `input.output_partitioning().partition_count()`. + /// + /// Late-bound: `None` until [`PrefixMergeExec::resolve_state`]. The rule + /// plants this operator at plan time, but the state only exists once + /// every upstream task has closed its window and published accumulator + /// state. `execute` refuses while unresolved rather than treating an + /// absent carry-in as zero, which would silently emit partition-local + /// aggregates. + per_partition_state: Arc<Mutex<Option<Vec<FinalizedPartitionState>>>>, + properties: Arc<PlanProperties>, +} + +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)) + } + + /// Errors on any of: + /// - `per_partition_state.len()` != input's partition count, when given. + /// - Any [`WindowApply::Scalar`]'s `offset.len()` != input's partition + /// count. + /// - Any entry's `output_column` outside the input schema's field range. + fn try_new_inner( + input: Arc<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, + }) + } + + /// Bind the prefix state. Called by the scheduler once the upstream + /// stage's tasks have all reported their finalized accumulator state and + /// it has been prefix-merged into one carry-in per partition. + /// + /// Idempotent overwrite, matching `RangeFilterExec::resolve_bounds`: AQE + /// re-plans, and a later pass may resolve the same operator again with + /// the same values. + pub fn resolve_state(&self, state: Vec<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. The scheduler descends from here to find the + /// state-sync boundary and the window operator whose expressions the + /// reported state is indexed against. + pub fn input(&self) -> &Arc<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. [`WindowApply::Aggregate`] carries `window_expr_index`, + /// which indexes window expressions rather than partitions, so it rides + /// over unchanged. + fn slice_to_partitions( + &self, + child: Arc<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, + }), + 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)] + } + + /// 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(); + + // The DF-side getter (BoundedWindowAggExec::finalized_partition_state) + // already commits to at-most-one PARTITION BY group per DF partition, + // so this operator receives the state indexed only by window + // expression — no PARTITION BY 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: Option<&Vec<ScalarValue>> = key_state + .get(*window_expr_index) + .and_then(|slot| slot.as_ref()); + 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), + }; + 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; + Ok(RecordBatch::try_new(batch.schema(), columns)?) + } +} + +/// 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>, + 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 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()?); + } Review 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`](https://github.com/apache/datafusion/blob/55.0.0-rc2/datafusion/functions-aggregate/src/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. -- 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]
