From 5cdf7d5f44cc3d6f5902f9fb2a2b4513a240d869 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:43:28 +0300 Subject: [PATCH 01/42] copy paste aggregates as is to aggregates_blocked --- .../aggregate_hash_table/common.rs | 728 ++ .../aggregate_hash_table/common_ordered.rs | 468 + .../aggregate_hash_table/final_table.rs | 85 + .../aggregate_hash_table/mod.rs | 160 + .../ordered_final_table.rs | 92 + .../ordered_partial_table.rs | 114 + .../ordered_single_table.rs | 93 + .../partial_reduce_table.rs | 79 + .../aggregate_hash_table/partial_table.rs | 239 + .../aggregate_hash_table/single_table.rs | 84 + .../aggregates_blocked/aggregate_stream.rs | 478 + .../group_values/metrics.rs | 579 ++ .../aggregates_blocked/group_values/mod.rs | 217 + .../group_values/multi_group_by/boolean.rs | 495 + .../group_values/multi_group_by/bytes.rs | 702 ++ .../group_values/multi_group_by/bytes_view.rs | 1023 +++ .../group_values/multi_group_by/dictionary.rs | 861 ++ .../multi_group_by/fixed_size_binary.rs | 516 ++ .../group_values/multi_group_by/mod.rs | 2613 ++++++ .../group_values/multi_group_by/primitive.rs | 654 ++ .../group_values/multi_group_by/row_backed.rs | 1133 +++ .../group_values/null_builder.rs | 67 + .../aggregates_blocked/group_values/row.rs | 414 + .../group_values/single_group_by/boolean.rs | 153 + .../group_values/single_group_by/bytes.rs | 128 + .../single_group_by/bytes_view.rs | 130 + .../group_values/single_group_by/mod.rs | 23 + .../group_values/single_group_by/primitive.rs | 296 + .../aggregates_blocked/grouped_hash_stream.rs | 1695 ++++ .../aggregates_blocked/grouped_topk_stream.rs | 403 + .../src/aggregates_blocked/hash_stream.rs | 1838 ++++ .../src/aggregates_blocked/mod.rs | 8176 +++++++++++++++++ .../src/aggregates_blocked/order/full.rs | 156 + .../src/aggregates_blocked/order/mod.rs | 219 + .../src/aggregates_blocked/order/partial.rs | 358 + .../ordered_final_stream.rs | 903 ++ .../ordered_partial_stream.rs | 352 + .../ordered_single_stream.rs | 887 ++ .../partial_reduce_stream.rs | 385 + .../src/aggregates_blocked/single_stream.rs | 834 ++ .../src/aggregates_blocked/skip_partial.rs | 305 + .../src/aggregates_blocked/topk/hash_table.rs | 697 ++ .../src/aggregates_blocked/topk/heap.rs | 779 ++ .../src/aggregates_blocked/topk/mod.rs | 22 + .../aggregates_blocked/topk/priority_map.rs | 805 ++ datafusion/physical-plan/src/lib.rs | 1 + 46 files changed, 31439 insertions(+) create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common_ordered.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/final_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_final_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_partial_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_single_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_reduce_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/single_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/aggregate_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/metrics.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/boolean.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/bytes.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/bytes_view.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/dictionary.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/fixed_size_binary.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/primitive.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/multi_group_by/row_backed.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/null_builder.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/row.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/single_group_by/boolean.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/single_group_by/bytes.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/single_group_by/bytes_view.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/single_group_by/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/group_values/single_group_by/primitive.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/grouped_hash_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/grouped_topk_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/hash_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/order/full.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/order/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/order/partial.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/ordered_final_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/ordered_partial_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/ordered_single_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/partial_reduce_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/single_stream.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/skip_partial.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/topk/hash_table.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/topk/heap.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/topk/mod.rs create mode 100644 datafusion/physical-plan/src/aggregates_blocked/topk/priority_map.rs diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common.rs new file mode 100644 index 0000000000000..6aa246beedb9d --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common.rs @@ -0,0 +1,728 @@ +// 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. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; + +use crate::PhysicalExpr; +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + GroupByMetrics, GroupValues, new_group_values, +}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, +}; + +use super::AggregateTableMetrics; + +/// Marker for raw rows -> partial state aggregation. +pub(in crate::aggregates) struct PartialMarker; +/// Marker for raw rows -> final value aggregation. +pub(in crate::aggregates) struct SingleMarker; +/// Marker for partial state -> partial state aggregation. +pub(in crate::aggregates) struct PartialReduceMarker; +/// Marker for raw rows -> partial state conversion without aggregation. +pub(in crate::aggregates) struct PartialSkipMarker; +/// Marker for partial state -> final value aggregation. +pub(in crate::aggregates) struct FinalMarker; + +/// Grouped hash table shared by the partial and final paths. +/// +/// While building, it consumes input batches and updates group / accumulator +/// state. While outputting, it incrementally drains that state into output +/// batches. +/// +/// # Logical and Physical Model +/// +/// Logically, this is a hash table that maps { group keys -> accumulator states } +/// For example, `AVG(v) GROUP BY k` stores one entry per `k`, where each +/// entry owns the `sum(v)` and `count(v)` state needed to compute the final +/// average. +/// +/// Physically, the group keys and accumulators are backed by [`GroupValues`] and +/// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay +/// vectorized. +/// +/// # Marker Type +/// `AggrMode` selects the aggregate semantics. +/// +/// e.g. `AggregateHashTable::::new(...)` creates an aggregate hash table +/// for the partial hash aggregate stage, the input schema is raw rows and output +/// schema is intermediate states. +/// +/// It is a zero-sized compile-time marker, so each stage keeps its update logic +/// in a separate impl block, to make the behavior difference explicit. +pub(in crate::aggregates) struct AggregateHashTable { + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Per-aggregate timing metrics for evaluating aggregate arguments. + pub(super) aggregate_argument_metrics: AggregateArgumentMetrics, + + /// Per-aggregate timing metrics for accumulator operations. + pub(super) aggregate_accumulator_metrics: Arc, + + /// Raw input schema, used to evaluate expressions and synthesize empty + /// grouping-set rows. + pub(super) input_schema: SchemaRef, + + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Lifecycle-specific state: building stage / outputting stage. + pub(super) state: AggregateHashTableState, + + pub(super) _mode: PhantomData, +} + +/// Methods shared by all aggregate hash table modes. +impl AggregateHashTable { + pub(super) fn new_with_filters( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + filters: Vec>>, + ) -> Result { + if batch_size == 0 { + return internal_err!("AggregateHashTable requires config batch_size >= 1"); + } + + let input_schema = agg.input().schema(); + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + &agg.mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators: Vec<_> = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(HashAggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::>()?; + + let group_schema = agg.group_by.group_schema(&input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + + let metrics = AggregateTableMetrics::new(agg, partition); + + Ok(Self { + group_by_metrics: metrics.group_by, + aggregate_argument_metrics: metrics.aggregate_arguments, + aggregate_accumulator_metrics: metrics.accumulator, + input_schema, + output_schema, + state_schema, + batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// See comments in [`EvaluatedAggregateBatch`] + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result { + let state = self.state.building(); + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + // outer vec: one per each grouping set + // inner vec: all group by exprs for the current grouping set + let grouping_set_args = evaluate_group_by(&state.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + // The evaluated args for each accumulator + let accumulator_args = state + .accumulators + .iter() + .enumerate() + .map(|(idx, acc)| { + self.aggregate_argument_metrics + .time(idx, || acc.evaluate_acc_args(batch)) + }) + .collect::>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Aggregates one input batch after selecting the mode-specific accumulator + /// operation. + /// + /// Each aggregation mode chooses a different `aggregate_fn` according to its + /// semantics. For example, partial aggregation takes raw inputs, and update them + /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used. + pub(super) fn aggregate_batch_inner( + &mut self, + batch: &RecordBatch, + aggregate_fn: AggregateBatchFn, + accumulator_phase: AccumulatorPhase, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let state = self.state.building_mut(); + + let _timer = self.group_by_metrics.aggregation_time.timer(); + for group_values in &evaluated_batch.grouping_set_args { + state + .group_values + .intern(group_values, &mut state.batch_group_indices)?; + let group_indices = &state.batch_group_indices; + let total_num_groups = state.group_values.len(); + + for (idx, (acc, values)) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + .enumerate() + { + accumulator_metrics.time(idx, accumulator_phase, || { + aggregate_fn(acc, values, group_indices, total_num_groups) + })?; + } + } + + Ok(()) + } + + /// Materializes the full output once, then returns it downstream incrementally + /// by slicing it into `batch_size` chunks. + /// + /// Each aggregation mode chooses a different `materialize_accumulator_fn` + /// according to its semantics. For example, partial aggregation emits + /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`]. + /// + /// This is a temporary solution until blocked state management is implemented: + /// Issue: + pub(super) fn next_output_batch_inner( + &mut self, + materialize_accumulator_fn: MaterializeAccumulatorFn, + accumulator_phase: AccumulatorPhase, + ) -> Result> { + let output_schema = Arc::clone(&self.output_schema); + let batch_size = self.batch_size; + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + + let mut output = + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let emit_to = EmitTo::All; + let timer = self.group_by_metrics.emitting_time.timer(); + let mut columns = state.group_values.emit(emit_to)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + columns.extend(accumulator_metrics.time( + idx, + accumulator_phase, + || materialize_accumulator_fn(acc, emit_to), + )?); + } + drop(timer); + + let batch = RecordBatch::try_new(output_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + MaterializedAggregateOutput::new(batch) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( + "next_output_batch must be called in the outputting state" + ); + } + }; + + let batch = output.next_batch(batch_size); + if output.is_exhausted() { + self.state = AggregateHashTableState::Done; + } else { + self.state = AggregateHashTableState::OutputtingMaterialized(output); + } + Ok(batch) + } + + pub(in crate::aggregates) fn memory_size(&self) -> usize { + match &self.state { + AggregateHashTableState::Building(state) + | AggregateHashTableState::Outputting(state) => { + let acc = state + .accumulators + .iter() + .map(|acc| acc.accumulator.size()) + .sum::(); + + acc + state.group_values.size() + + state.batch_group_indices.allocated_size() + } + AggregateHashTableState::OutputtingMaterialized(output) => { + output.memory_size() + } + AggregateHashTableState::Done => 0, + } + } + + /// Returns the number of distinct groups accumulated so far. + pub(in crate::aggregates) fn building_group_count(&self) -> usize { + self.state.building().group_values.len() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue accumulating raw input. + /// + /// Unlike normal single aggregation output, this materializes intermediate + /// states rather than final values. The states can therefore be merged after + /// spilling without finalizing the same group more than once. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result> { + let state_schema = Arc::clone(&self.state_schema); + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let state = self.state.building_mut(); + if state.group_values.is_empty() { + return Ok(None); + } + + let mut output = state.group_values.emit(EmitTo::All)?; + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + output.extend(accumulator_metrics.time( + idx, + AccumulatorPhase::State, + || acc.state(EmitTo::All), + )?); + } + + let batch = RecordBatch::try_new(state_schema, output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is sorted for spilling. + state.group_values.clear_shrink(0); + state.batch_group_indices.clear(); + state.batch_group_indices.shrink_to_fit(); + + Ok(Some(batch)) + } + + pub(in crate::aggregates) fn is_building(&self) -> bool { + matches!(self.state, AggregateHashTableState::Building(_)) + } + + pub(in crate::aggregates) fn is_done(&self) -> bool { + matches!(self.state, AggregateHashTableState::Done) + } + + pub(super) fn start_outputting(&mut self) { + let AggregateHashTableState::Building(mut state) = + std::mem::replace(&mut self.state, AggregateHashTableState::Done) + else { + unreachable!("hash aggregate table is not building") + }; + + state.batch_group_indices = Vec::new(); + self.state = AggregateHashTableState::Outputting(state); + } +} + +/// State and argument information for a single Aggregate +/// +/// For example, for `SELECT COUNT(x), SUM(y WHERE z > 10) ...` there would be two +/// `HashAggregateAccumulator`, one each for `COUNT(x)` and `SUM(y WHERE z > 10)` +pub(super) struct HashAggregateAccumulator { + /// Aggregate expression used to create a fresh accumulator for related + /// hash tables, such as the partial-skip table. + aggregate_expr: Arc, + + /// Arguments to pass to this accumulator. + /// + /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one. + arguments: Vec>, + + /// Optional `FILTER` expression for this accumulator. + /// + /// Example: `SUM(x) FILTER (WHERE x > 10)` stores the `x > 10` predicate. + filter: Option>, + + /// Accumulator state for all groups for one aggregate expression. + accumulator: Box, +} + +pub(super) type AggregateAccumulator = HashAggregateAccumulator; + +/// Function used by [`AggregateHashTable::aggregate_batch_inner`] to update one +/// accumulator with one evaluated input batch. +/// +/// Arguments: +/// * accumulator to update. +/// * accumulator's evaluated arguments and optional filter. +/// * one group index per input row, mapping each row to its interned group. +/// * total number of groups currently interned in that buffer, including newly +/// interned groups. +pub(super) type AggregateBatchFn = fn( + &mut AggregateAccumulator, + &EvaluatedAccumulatorArgs, + &[usize], + usize, +) -> Result<()>; + +/// Function used by [`AggregateHashTable::next_output_batch_inner`] to +/// materialize one accumulator's output columns. +/// +/// Arguments: +/// * accumulator to materialize. +/// * group range to emit from the accumulator. +pub(super) type MaterializeAccumulatorFn = + fn(&mut AggregateAccumulator, EmitTo) -> Result>; + +/// Evaluated aggregate arguments and filter for one input batch. +/// +/// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` +/// and `x > 0`. +/// +/// These arrays can be passed directly to [`GroupsAccumulator`]. +pub(super) struct EvaluatedAccumulatorArgs { + /// Evaluated argument arrays. Some aggregate functions take multiple arguments. + pub(super) arguments: Vec, + /// Evaluated filter array, `Some` if the aggregate has a `FILTER` expression. + pub(super) filter: Option, +} + +/// Evaluated all group by keys and accumulator args. +/// +/// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function evaluates +/// `k+1`, `v*v` +pub(super) struct EvaluatedAggregateBatch { + /// One entry per grouping set; each entry contains all evaluated group key + /// arrays for the current input batch. + pub(super) grouping_set_args: Vec>, + + /// Evaluated arguments and filters, one entry per aggregate expression. + pub(super) accumulator_args: Vec, +} + +/// Buffer for the aggregate hash table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits final results during the +/// outputting stage. +/// +/// [`GroupValues`] stores the physical group-key layout, while +/// [`GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct AggregateHashTableBuffer { + /// GROUP BY expressions evaluated for each input batch. + pub(super) group_by: Arc, + + /// Interned group keys. Accumulator state is stored separately by group index. + pub(super) group_values: Box, + + /// Group index for each row in the current input batch. + /// + /// Each value indexes into `group_values`, and the same index is used by every + /// accumulator to update that group's aggregate state. + pub(super) batch_group_indices: Vec, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec, +} + +pub(super) enum AggregateHashTableState { + /// Accumulating input rows into group keys and aggregate state. + Building(AggregateHashTableBuffer), + /// Emitting results directly from group keys and aggregate state. + Outputting(AggregateHashTableBuffer), + /// Materialize all the output results, and then incrementally output in the `OutputtingMaterialized` state. + /// + /// Note this is a temporary solution until the `GroupValues` issue is solved: + /// Issue: + OutputtingMaterialized(MaterializedAggregateOutput), + Done, +} + +/// Fully evaluated aggregate output and the next row offset to emit. +/// +/// Final aggregate evaluation consumes accumulator state, and partial terminal +/// output should not repeatedly renumber group values with `EmitTo::First`. +/// Materialize once and then slice to honor `batch_size` across output polls. +pub(super) struct MaterializedAggregateOutput { + batch: RecordBatch, + offset: usize, +} + +impl MaterializedAggregateOutput { + pub(super) fn new(batch: RecordBatch) -> Self { + Self { batch, offset: 0 } + } + + pub(super) fn next_batch(&mut self, batch_size: usize) -> Option { + debug_assert!(batch_size > 0); + if self.is_exhausted() { + return None; + } + + let length = batch_size.min(self.batch.num_rows() - self.offset); + let batch = self.batch.slice(self.offset, length); + self.offset += length; + Some(batch) + } + + pub(super) fn is_exhausted(&self) -> bool { + self.offset >= self.batch.num_rows() + } + + pub(super) fn memory_size(&self) -> usize { + self.batch.get_array_memory_size() + } +} + +impl HashAggregateAccumulator { + pub(super) fn new( + aggregate_expr: Arc, + arguments: Vec>, + filter: Option>, + accumulator: Box, + ) -> Self { + Self { + aggregate_expr, + arguments, + filter, + accumulator, + } + } + + /// Construct a new accumulator with the same definition, but with empty internal + /// state buffers (empty [`GroupsAccumulator`]). + pub(super) fn empty_like(&self) -> Result { + let accumulator = create_group_accumulator(&self.aggregate_expr)?; + Ok(Self::new( + Arc::clone(&self.aggregate_expr), + self.arguments.clone(), + self.filter.clone(), + accumulator, + )) + } + + /// Evaluate aggregate arguments and filter for one input batch. + /// + /// For example, `AVG(2 / x) FILTER (WHERE x > 0)` evaluates `x > 0` + /// first, then evaluates `2 / x` only for selected rows. + /// Filtered rows will be evaluated to `NULL`, and won't trigger errors + /// such as divide by zero. + /// + /// These arrays can be passed directly to [`GroupsAccumulator`] next. + pub(super) fn evaluate_acc_args( + &self, + batch: &RecordBatch, + ) -> Result { + let filter = self + .filter + .as_ref() + .map(|filter| { + filter + .evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .transpose()?; + let selection = filter.as_ref().map(|filter| filter.as_boolean()); + let arguments = self + .arguments + .iter() + .map(|expr| { + selection + .map_or_else( + || expr.evaluate(batch), + |selection| expr.evaluate_selection(batch, selection), + ) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .collect::>()?; + + Ok(EvaluatedAccumulatorArgs { arguments, filter }) + } + + pub(super) fn size(&self) -> usize { + self.accumulator.size() + } + + pub(super) fn update_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + let filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator.update_batch( + &values.arguments, + group_indices, + filter, + total_num_groups, + ) + } + + pub(super) fn merge_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + debug_assert!(values.filter.is_none()); + self.accumulator + .merge_batch(&values.arguments, group_indices, total_num_groups) + } + + /// Evaluating final aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `evaluate(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn evaluate(&mut self, emit_to: EmitTo) -> Result { + self.accumulator.evaluate(emit_to) + } + + pub(super) fn evaluate_to_columns( + &mut self, + emit_to: EmitTo, + ) -> Result> { + Ok(vec![self.evaluate(emit_to)?]) + } + + /// Evaluating partial aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `state(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn state(&mut self, emit_to: EmitTo) -> Result> { + self.accumulator.state(emit_to) + } + + pub(super) fn convert_to_state( + &mut self, + values: &EvaluatedAccumulatorArgs, + ) -> Result> { + let opt_filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator + .convert_to_state(&values.arguments, opt_filter) + } + + pub(super) fn null_arguments( + &self, + input_schema: &SchemaRef, + ) -> Result> { + self.arguments + .iter() + .map(|expr| { + let data_type = expr.data_type(input_schema)?; + Ok(new_null_array(&data_type, 1)) + }) + .collect() + } +} + +impl AggregateHashTableState { + pub(super) fn building(&self) -> &AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } + + pub(super) fn building_mut(&mut self) -> &mut AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Array, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + + use super::*; + + #[test] + fn materialized_aggregate_output_slices_batches_until_exhausted() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + )?; + let mut output = MaterializedAggregateOutput::new(batch); + + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]); + assert!(output.next_batch(2).is_none()); + assert!(output.is_exhausted()); + + Ok(()) + } + + fn int32_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|idx| array.value(idx)).collect() + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common_ordered.rs new file mode 100644 index 0000000000000..4ced967a0977b --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/common_ordered.rs @@ -0,0 +1,468 @@ +// 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. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::assert_or_internal_err; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + GroupByMetrics, GroupValues, new_group_values, +}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::AggregateTableMetrics; +use super::common::{ + AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, + MaterializeAccumulatorFn, +}; + +#[derive(Clone)] +pub(in crate::aggregates) struct OrderedAggregateTableMetrics { + pub(super) group_by: GroupByMetrics, + pub(super) aggregate_arguments: AggregateArgumentMetrics, + pub(super) accumulator: Arc, +} + +impl OrderedAggregateTableMetrics { + pub(in crate::aggregates) fn new(agg: &AggregateExec, partition: usize) -> Self { + let metrics = AggregateTableMetrics::new(agg, partition); + Self { + group_by: metrics.group_by, + aggregate_arguments: metrics.aggregate_arguments, + accumulator: metrics.accumulator, + } + } + + pub(in crate::aggregates) fn from_hash_table( + table: &AggregateHashTable, + ) -> Self { + Self { + group_by: table.group_by_metrics.clone(), + aggregate_arguments: table.aggregate_argument_metrics.clone(), + accumulator: Arc::clone(&table.aggregate_accumulator_metrics), + } + } +} + +/// Aggregate table shared by the ordered single, partial and final paths. +/// +/// # Ordering optimization +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Single, partial and final variant difference +/// +/// The partial and final aggregate tables implement the two stages of grouped +/// aggregation, while the single aggregate table implements both stages in one +/// table. See +/// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream) +/// for the high-level plan shape. +/// +/// Example: `AVG(v) FILTER (WHERE v>0) GROUP BY k` +/// +/// Partial table ([`AggregateMode::Partial`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, sum(v), count(v)` +/// +/// Final table ([`AggregateMode::Final`], no filters): +/// - Input rows: `k, sum(v), count(v)` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// +/// Single table ([`AggregateMode::Single`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl`; single/partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to pass through or spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Per-aggregate timing metrics for evaluating aggregate arguments. + pub(super) aggregate_argument_metrics: AggregateArgumentMetrics, + + /// Per-aggregate timing metrics for accumulator operations. + pub(super) aggregate_accumulator_metrics: Arc, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec, +} + +/// Methods shared by all aggregate modes +impl OrderedAggregateTable { + #[expect( + clippy::too_many_arguments, + reason = "keeps ordered single, partial and final table construction explicit" + )] + pub(super) fn new_for_mode( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + aggregate_mode: &AggregateMode, + filters: Vec>>, + metrics: OrderedAggregateTableMetrics, + ) -> Result { + assert_or_internal_err!( + batch_size > 0, + "OrderedAggregateTable requires config batch_size >= 1" + ); + + let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_schema = agg.group_by.group_schema(input_schema)?; + let group_values = new_group_values(group_schema, &group_ordering)?; + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + aggregate_mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(AggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::>()?; + + Ok(Self { + output_schema, + state_schema, + batch_size, + group_by_metrics: metrics.group_by, + aggregate_argument_metrics: metrics.aggregate_arguments, + aggregate_accumulator_metrics: metrics.accumulator, + buffer: OrderedAggregateTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_ordering, + group_values, + group_indices: vec![], + accumulators, + }, + _mode: PhantomData, + }) + } + + /// Evaluates all group by keys and accumulator args. + /// + /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function + /// evaluates `k+1`, `v*v`. + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result { + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let accumulator_args = self + .buffer + .accumulators + .iter() + .enumerate() + .map(|(idx, acc)| { + self.aggregate_argument_metrics + .time(idx, || acc.evaluate_acc_args(batch)) + }) + .collect::>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Called after the input stream is exhausted and the last batch has been + /// aggregated. + /// + /// Updates the internal `GroupOrdering` so it can continue emitting until + /// the buffer is empty. + pub(in crate::aggregates) fn input_done(&mut self) { + self.buffer.group_ordering.input_done(); + } + + /// Returns the ordering state used to decide how memory pressure is handled. + pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering { + &self.buffer.group_ordering + } + + /// Number of groups currently buffered. + pub(in crate::aggregates) fn num_groups(&self) -> usize { + self.buffer.group_values.len() + } + + /// Check if there is zero groups accumulated so far. + pub(in crate::aggregates) fn is_empty(&self) -> bool { + self.num_groups() == 0 + } + + /// All internal buffer's memory size. + pub(in crate::aggregates) fn memory_size(&self) -> usize { + self.buffer + .accumulators + .iter() + .map(|acc| acc.size()) + .sum::() + + self.buffer.group_values.size() + + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() + } + + pub(in crate::aggregates) fn metrics(&self) -> OrderedAggregateTableMetrics { + OrderedAggregateTableMetrics { + group_by: self.group_by_metrics.clone(), + aggregate_arguments: self.aggregate_argument_metrics.clone(), + accumulator: Arc::clone(&self.aggregate_accumulator_metrics), + } + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue with a new ordered input segment. + /// + /// Unlike normal ordered emission, this operation is allowed to take the + /// active (incomplete) groups. Partial aggregation can pass those states to + /// its final stage, while single and final aggregation sort and spill them + /// before replay. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let mut output = self.buffer.group_values.emit(EmitTo::All)?; + for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { + output.extend(accumulator_metrics.time( + idx, + AccumulatorPhase::State, + || acc.state(EmitTo::All), + )?); + } + + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is passed downstream or sorted for spilling. + self.buffer.group_values.clear_shrink(0); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); + self.buffer.group_ordering.reset(); + + Ok(Some(batch)) + } + + /// Returns the [`EmitTo`], clamped to the specified batch size + /// + /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number + /// of groups to emit from `GroupValues` / accumulators, and + /// `should_remove_groups` indicates whether `GroupOrdering` must also shift + /// its tracked indexes. + pub(super) fn clamp_emit_to( + &self, + group_count: usize, + emit_to: EmitTo, + ) -> (EmitTo, bool) { + match emit_to { + EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), + EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), + EmitTo::All => (EmitTo::First(self.batch_size), false), + } + } + + /// Aggregates one evaluated input batch after selecting the mode-specific + /// accumulator operation. + /// + /// Each aggregation mode chooses a different `aggregate_fn` according to its + /// semantics. For example, partial aggregation takes raw inputs and updates + /// stored partial states, so it uses + /// [`datafusion_expr::GroupsAccumulator::update_batch`]. + pub(super) fn aggregate_evaluated_batch( + &mut self, + evaluated_batch: &EvaluatedAggregateBatch, + aggregate_fn: AggregateBatchFn, + accumulator_phase: AccumulatorPhase, + ) -> Result<()> { + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (idx, (acc, values)) in self + .buffer + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + .enumerate() + { + accumulator_metrics.time(idx, accumulator_phase, || { + aggregate_fn( + acc, + values, + &self.buffer.group_indices, + total_num_groups, + ) + })?; + } + drop(timer); + } + + Ok(()) + } + + /// Emits groups allowed by `GroupOrdering`, leaving only the current + /// unfinished ordered-key range buffered. + /// + /// Each aggregation mode chooses a different `materialize_accumulator_fn` + /// according to its semantics. For example, partial aggregation emits + /// partial states to feed the final stage, so it uses + /// [`datafusion_expr::GroupsAccumulator::state`]. + pub(super) fn next_output_batch_inner( + &mut self, + materialize_accumulator_fn: MaterializeAccumulatorFn, + accumulator_phase: AccumulatorPhase, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let Some(emit_to) = self.buffer.group_ordering.emit_to() else { + return Ok(None); + }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); + + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + if should_remove_groups { + match emit_to { + EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + // `EmitTo::All` is only used after `input_done`, when all + // buffered groups are known complete and the ordering state is + // no longer needed. + EmitTo::All => {} + } + } + + for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() { + output.extend(accumulator_metrics.time(idx, accumulator_phase, || { + materialize_accumulator_fn(acc, emit_to) + })?); + } + drop(timer); + + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(Some(batch)) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/final_table.rs new file mode 100644 index 0000000000000..307215cfd2797 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/final_table.rs @@ -0,0 +1,85 @@ +// 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. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; +use crate::aggregates::group_values::AccumulatorPhase; + +use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + output_schema, + Arc::clone(&agg.input().schema()), + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::evaluate_to_columns, + AccumulatorPhase::Evaluate, + ) + } + + /// Final aggregation consumes partial aggregate states and merges them into + /// the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner( + batch, + HashAggregateAccumulator::merge_batch, + AccumulatorPhase::Merge, + ) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/mod.rs new file mode 100644 index 0000000000000..902a859bac96d --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/mod.rs @@ -0,0 +1,160 @@ +// 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. + +mod common; +mod common_ordered; +mod final_table; +mod ordered_final_table; +mod ordered_partial_table; +mod ordered_single_table; +mod partial_reduce_table; +mod partial_table; +mod single_table; + +use std::sync::Arc; + +use crate::aggregates::group_values::{ + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + GroupByMetrics, +}; +use crate::aggregates::{AggregateExec, AggregateMode, aggregate_metric_label}; + +pub(super) fn accumulator_phases(mode: &AggregateMode) -> &'static [AccumulatorPhase] { + match mode { + AggregateMode::Partial => &[ + AccumulatorPhase::Update, + AccumulatorPhase::State, + AccumulatorPhase::ConvertToState, + ], + AggregateMode::PartialReduce => { + &[AccumulatorPhase::Merge, AccumulatorPhase::State] + } + // Final and single aggregation emit intermediate states when spilling, + // then replay them through a final aggregate table. + AggregateMode::Final | AggregateMode::FinalPartitioned => &[ + AccumulatorPhase::Merge, + AccumulatorPhase::State, + AccumulatorPhase::Evaluate, + ], + AggregateMode::Single | AggregateMode::SinglePartitioned => &[ + AccumulatorPhase::Update, + AccumulatorPhase::State, + AccumulatorPhase::Merge, + AccumulatorPhase::Evaluate, + ], + } +} + +pub(super) struct AggregateTableMetrics { + pub(super) group_by: GroupByMetrics, + pub(super) aggregate_arguments: AggregateArgumentMetrics, + pub(super) accumulator: Arc, +} + +impl AggregateTableMetrics { + pub(super) fn new(agg: &AggregateExec, partition: usize) -> Self { + let aggregate_labels = agg + .aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)) + .collect::>(); + + Self { + group_by: GroupByMetrics::new(&agg.metrics, partition), + aggregate_arguments: AggregateArgumentMetrics::new( + &agg.metrics, + partition, + aggregate_labels.clone(), + ), + accumulator: Arc::new(AggregateAccumulatorMetrics::new( + &agg.metrics, + partition, + aggregate_labels, + accumulator_phases(&agg.mode), + )), + } + } +} + +pub(super) use common::{ + AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker, + PartialSkipMarker, SingleMarker, +}; +pub(super) use common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; + +#[cfg(test)] +mod tests { + use super::accumulator_phases; + use crate::aggregates::AggregateMode; + use crate::aggregates::group_values::AccumulatorPhase; + + #[test] + fn accumulator_phases_match_aggregate_mode() { + let cases: [(AggregateMode, &[AccumulatorPhase]); 6] = [ + ( + AggregateMode::Partial, + &[ + AccumulatorPhase::Update, + AccumulatorPhase::State, + AccumulatorPhase::ConvertToState, + ], + ), + ( + AggregateMode::PartialReduce, + &[AccumulatorPhase::Merge, AccumulatorPhase::State], + ), + ( + AggregateMode::Final, + &[ + AccumulatorPhase::Merge, + AccumulatorPhase::State, + AccumulatorPhase::Evaluate, + ], + ), + ( + AggregateMode::FinalPartitioned, + &[ + AccumulatorPhase::Merge, + AccumulatorPhase::State, + AccumulatorPhase::Evaluate, + ], + ), + ( + AggregateMode::Single, + &[ + AccumulatorPhase::Update, + AccumulatorPhase::State, + AccumulatorPhase::Merge, + AccumulatorPhase::Evaluate, + ], + ), + ( + AggregateMode::SinglePartitioned, + &[ + AccumulatorPhase::Update, + AccumulatorPhase::State, + AccumulatorPhase::Merge, + AccumulatorPhase::Evaluate, + ], + ), + ]; + + for (mode, expected) in cases { + assert!(accumulator_phases(&mode) == expected, "{mode:?}"); + } + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_final_table.rs new file mode 100644 index 0000000000000..d0d0c99bb5bd8 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_final_table.rs @@ -0,0 +1,92 @@ +// 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. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase}; + +use super::common::HashAggregateAccumulator; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + metrics: OrderedAggregateTableMetrics, + ) -> Result { + Self::new_for_mode( + agg, + input_schema, + output_schema, + Arc::clone(input_schema), + batch_size, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + metrics, + ) + } + + /// Merges one partial-state input batch and updates ordering information for + /// any newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + // `PhysicalGroupBy::as_final()` removes grouping sets while planning + // final aggregation, so final ordered aggregation sees one grouping. + debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::merge_batch, + AccumulatorPhase::Merge, + ) + } + + /// See comments in `ordered_partial_stream::next_output_batch` + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::evaluate_to_columns, + AccumulatorPhase::Evaluate, + ) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_partial_table.rs new file mode 100644 index 0000000000000..6ed93e59f3296 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_partial_table.rs @@ -0,0 +1,114 @@ +// 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. + +//! Aggregate table for partial aggregation when input is ordered by group keys. +//! +//! See the [`super::common_ordered`] comments for the high-level ideas. +//! +//! This operator handles input that is ordered by group keys: +//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` +//! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` +//! +//! When a group key combination is exhausted, this table eagerly flushes the +//! completed groups to improve memory efficiency. +//! +//! The implementation is separated from other aggregate tables because this +//! execution path is likely to be optimized further in the future. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::{ + AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, + group_values::AccumulatorPhase, +}; + +use super::common::HashAggregateAccumulator; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + let input_schema = agg.input().schema(); + let state_schema = Arc::clone(&output_schema); + let metrics = OrderedAggregateTableMetrics::new(agg, partition); + Self::new_for_mode( + agg, + &input_schema, + output_schema, + state_schema, + batch_size, + &agg.input_order_mode, + &AggregateMode::Partial, + agg.filter_expr.iter().cloned().collect(), + metrics, + ) + } + + /// Aggregates one raw input batch and updates ordering information for any + /// newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::update_batch, + AccumulatorPhase::Update, + ) + } + + /// Emits the next batch of partial state rows for groups proven complete by + /// the input ordering. + /// + /// For example, when the query is `GROUP BY a` and the input is ordered by + /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` + /// are complete and safe to emit. + /// + /// Key steps: + /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. + /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and + /// all `GroupsAccumulator`s. + /// + /// This may output small batches. Avoiding tiny batches is left to future + /// ordered-aggregation optimizations. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::state, + AccumulatorPhase::State, + ) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_single_table.rs new file mode 100644 index 0000000000000..ce1ce647b46fe --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/ordered_single_table.rs @@ -0,0 +1,93 @@ +// 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. + +//! Aggregate table for single aggregation when raw input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::aggregate_hash_table::SingleMarker; +use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase}; + +use super::common::HashAggregateAccumulator; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; + +/// Implementation specific to single aggregation, where the table stores final +/// aggregate values and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, avg(x)` +/// - Input rows: `k, x` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + )); + + let input_schema = agg.input().schema(); + let metrics = OrderedAggregateTableMetrics::new(agg, partition); + Self::new_for_mode( + agg, + &input_schema, + output_schema, + state_schema, + batch_size, + &agg.input_order_mode, + &agg.mode, + agg.filter_expr.iter().cloned().collect(), + metrics, + ) + } + + /// Aggregates one raw input batch and updates ordering information for any + /// newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + self.aggregate_evaluated_batch( + &evaluated_batch, + HashAggregateAccumulator::update_batch, + AccumulatorPhase::Update, + ) + } + + /// Emits the next batch of final aggregate values for groups proven complete + /// by the input ordering. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::evaluate_to_columns, + AccumulatorPhase::Evaluate, + ) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_reduce_table.rs new file mode 100644 index 0000000000000..2892d059332cf --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_reduce_table.rs @@ -0,0 +1,79 @@ +// 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. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; +use crate::aggregates::group_values::AccumulatorPhase; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker}; + +/// Methods specific to the aggregate hash table used in the partial-reduce stage. +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::state, + AccumulatorPhase::State, + ) + } + + /// Partial-reduce aggregation consumes partial aggregate states and merges + /// them into the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner( + batch, + HashAggregateAccumulator::merge_batch, + AccumulatorPhase::Merge, + ) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_table.rs new file mode 100644 index 0000000000000..54997f0537b87 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/partial_table.rs @@ -0,0 +1,239 @@ +// 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. + +use std::collections::HashMap; +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; + +use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; + +use super::common::{ + AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, + EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, +}; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::state, + AccumulatorPhase::State, + ) + } + + /// In skip-partial-aggregation optimization, when a decision has been made to skip + /// partial stage, build a typed hash table only for aggregation state conversion + /// row-by-row. + pub(in crate::aggregates) fn partial_skip_table( + &self, + ) -> Result> { + let state = self.state.building(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::>>()?; + + Ok(AggregateHashTable { + group_by_metrics: self.group_by_metrics.clone(), + aggregate_argument_metrics: self.aggregate_argument_metrics.clone(), + aggregate_accumulator_metrics: Arc::clone( + &self.aggregate_accumulator_metrics, + ), + input_schema: Arc::clone(&self.input_schema), + output_schema: Arc::clone(&self.output_schema), + state_schema: Arc::clone(&self.state_schema), + batch_size: self.batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&state.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// Partial aggregation consumes raw input rows and updates the table's + /// partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner( + batch, + HashAggregateAccumulator::update_batch, + AccumulatorPhase::Update, + ) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.init_empty_grouping_sets()?; + self.start_outputting(); + Ok(()) + } + + /// Creates the required empty grouping-set rows when the input is empty. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. + fn init_empty_grouping_sets(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + accumulator_metrics.time(idx, AccumulatorPhase::Update, || { + acc.update_batch(&values, &[0], total_groups) + })?; + } + } + + Ok(()) + } +} + +impl AggregateHashTable { + pub(in crate::aggregates) fn convert_batch_to_state( + &mut self, + batch: &RecordBatch, + ) -> Result { + let evaluated_batch = self.evaluate_batch(batch)?; + + assert_eq_or_internal_err!( + evaluated_batch.grouping_set_args.len(), + 1, + "group_values expected to have single element" + ); + let mut output = evaluated_batch + .grouping_set_args + .into_iter() + .next() + .unwrap_or_default(); + + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let state = self.state.building_mut(); + for (idx, (acc, values)) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + .enumerate() + { + output.extend(accumulator_metrics.time( + idx, + AccumulatorPhase::ConvertToState, + || acc.convert_to_state(values), + )?); + } + + Ok(RecordBatch::try_new( + Arc::clone(&self.output_schema), + output, + )?) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/single_table.rs new file mode 100644 index 0000000000000..1ff05fc79d224 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_hash_table/single_table.rs @@ -0,0 +1,84 @@ +// 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. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; +use crate::aggregates::group_values::AccumulatorPhase; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, SingleMarker}; + +/// Implementation specific to single aggregation, where the table stores final +/// aggregate values and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, avg(x)` +/// - Input rows: `k, x` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + output_schema, + state_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner( + HashAggregateAccumulator::evaluate_to_columns, + AccumulatorPhase::Evaluate, + ) + } + + /// Single aggregation consumes raw input rows and updates the table's + /// final-value accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner( + batch, + HashAggregateAccumulator::update_batch, + AccumulatorPhase::Update, + ) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates_blocked/aggregate_stream.rs new file mode 100644 index 0000000000000..ac7727b459300 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/aggregate_stream.rs @@ -0,0 +1,478 @@ +// 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. + +//! Aggregate without grouping columns + +use crate::aggregates::{ + AccumulatorItem, AggrDynFilter, AggregateInputMode, AggregateMode, + DynamicFilterAggregateType, aggregate_expressions, create_accumulators, + finalize_aggregation, +}; +use crate::metrics::{BaselineMetrics, RecordOutput}; +use crate::stream::EmptyRecordBatchStream; +use crate::{RecordBatchStream, SendableRecordBatchStream}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, lit}; +use futures::stream::BoxStream; +use std::borrow::Cow; +use std::cmp::Ordering; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::AggregateExec; +use crate::filter::batch_filter; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::stream::{Stream, StreamExt}; + +/// stream struct for aggregation without grouping columns +pub(crate) struct AggregateStream { + stream: BoxStream<'static, Result>, + schema: SchemaRef, +} + +/// Actual implementation of [`AggregateStream`]. +/// +/// This is wrapped into yet another struct because we need to interact with the async memory management subsystem +/// during poll. To have as little code "weirdness" as possible, we chose to just use [`BoxStream`] together with +/// [`futures::stream::unfold`]. +/// +/// The latter requires a state object, which is [`AggregateStreamInner`]. +struct AggregateStreamInner { + // ==== Properties ==== + schema: SchemaRef, + mode: AggregateMode, + input: SendableRecordBatchStream, + aggregate_expressions: Vec>>, + filter_expressions: Arc<[Option>]>, + + // ==== Runtime States/Buffers ==== + accumulators: Vec, + // None if the dynamic filter is not applicable. See details in `AggrDynFilter`. + agg_dyn_filter_state: Option>, + finished: bool, + + // ==== Execution Resources ==== + baseline_metrics: BaselineMetrics, + reservation: MemoryReservation, +} + +impl AggregateStreamInner { + // TODO: check if we get Null handling correct + /// # Examples + /// - Example 1 + /// Accumulators: min(c1) + /// Current Bounds: min(c1)=10 + /// --> dynamic filter PhysicalExpr: c1 < 10 + /// + /// - Example 2 + /// Accumulators: min(c1), max(c1), min(c2) + /// Current Bounds: min(c1)=10, max(c1)=100, min(c2)=20 + /// --> dynamic filter PhysicalExpr: (c1 < 10) OR (c1>100) OR (c2 < 20) + /// + /// # Errors + /// Returns internal errors if the dynamic filter is not enabled, or other + /// invariant check fails. + fn build_dynamic_filter_from_accumulator_bounds( + &self, + ) -> Result> { + let Some(filter_state) = self.agg_dyn_filter_state.as_ref() else { + return internal_err!( + "`build_dynamic_filter_from_accumulator_bounds()` is only called when dynamic filter is enabled" + ); + }; + + let mut predicates: Vec> = + Vec::with_capacity(filter_state.supported_accumulators_info.len()); + + for acc_info in &filter_state.supported_accumulators_info { + // Skip if we don't yet have a meaningful bound + let bound = { + let guard = acc_info.shared_bound.lock(); + if (*guard).is_null() { + continue; + } + guard.clone() + }; + + let agg_exprs = self + .aggregate_expressions + .get(acc_info.aggr_index) + .ok_or_else(|| { + internal_datafusion_err!( + "Invalid aggregate expression index {} for dynamic filter", + acc_info.aggr_index + ) + })?; + // Only aggregates with a single argument are supported. + let column_expr = agg_exprs.first().ok_or_else(|| { + internal_datafusion_err!( + "Aggregate expression at index {} expected a single argument", + acc_info.aggr_index + ) + })?; + + let literal = lit(bound); + let predicate: Arc = match acc_info.aggr_type { + DynamicFilterAggregateType::Min => Arc::new(BinaryExpr::new( + Arc::clone(column_expr), + Operator::Lt, + literal, + )), + DynamicFilterAggregateType::Max => Arc::new(BinaryExpr::new( + Arc::clone(column_expr), + Operator::Gt, + literal, + )), + }; + predicates.push(predicate); + } + + let combined = predicates.into_iter().reduce(|acc, pred| { + Arc::new(BinaryExpr::new(acc, Operator::Or, pred)) as Arc + }); + + Ok(combined.unwrap_or_else(|| lit(true))) + } + + // If the dynamic filter is enabled, update it using the current accumulator's + // values + fn maybe_update_dyn_filter(&mut self) -> Result<()> { + // Step 1: Update each partition's current bound + let Some(filter_state) = self.agg_dyn_filter_state.as_ref() else { + return Ok(()); + }; + + let mut bounds_changed = false; + + for acc_info in &filter_state.supported_accumulators_info { + let acc = + self.accumulators + .get_mut(acc_info.aggr_index) + .ok_or_else(|| { + internal_datafusion_err!( + "Invalid accumulator index {} for dynamic filter", + acc_info.aggr_index + ) + })?; + // First get current partition's bound, then update the shared bound among + // all partitions. + let current_bound = acc.evaluate()?; + { + let mut bound = acc_info.shared_bound.lock(); + let new_bound = match acc_info.aggr_type { + DynamicFilterAggregateType::Max => { + scalar_max(&bound, ¤t_bound)? + } + DynamicFilterAggregateType::Min => { + scalar_min(&bound, ¤t_bound)? + } + }; + if new_bound != *bound { + *bound = new_bound; + bounds_changed = true; + } + } + } + + // Step 2: Sync the dynamic filter physical expression with reader, + // but only if any bound actually changed. + if bounds_changed { + let predicate = self.build_dynamic_filter_from_accumulator_bounds()?; + filter_state.filter.update(predicate)?; + } + + Ok(()) + } +} + +/// Returns the element-wise minimum of two `ScalarValue`s. +/// +/// # Null semantics +/// - `min(NULL, NULL) = NULL` +/// - `min(NULL, x) = x` +/// - `min(x, NULL) = x` +/// +/// # Errors +/// Returns internal error if v1 and v2 has incompatible types. +fn scalar_min(v1: &ScalarValue, v2: &ScalarValue) -> Result { + if let Some(result) = scalar_cmp_null_short_circuit(v1, v2) { + return Ok(result); + } + + match v1.partial_cmp(v2) { + Some(Ordering::Less | Ordering::Equal) => Ok(v1.clone()), + Some(Ordering::Greater) => Ok(v2.clone()), + None => datafusion_common::internal_err!( + "cannot compare values of different or incompatible types: {v1:?} vs {v2:?}" + ), + } +} + +/// Returns the element-wise maximum of two `ScalarValue`s. +/// +/// # Null semantics +/// - `max(NULL, NULL) = NULL` +/// - `max(NULL, x) = x` +/// - `max(x, NULL) = x` +/// +/// # Errors +/// Returns internal error if v1 and v2 has incompatible types. +fn scalar_max(v1: &ScalarValue, v2: &ScalarValue) -> Result { + if let Some(result) = scalar_cmp_null_short_circuit(v1, v2) { + return Ok(result); + } + + match v1.partial_cmp(v2) { + Some(Ordering::Greater | Ordering::Equal) => Ok(v1.clone()), + Some(Ordering::Less) => Ok(v2.clone()), + None => datafusion_common::internal_err!( + "cannot compare values of different or incompatible types: {v1:?} vs {v2:?}" + ), + } +} + +fn scalar_cmp_null_short_circuit( + v1: &ScalarValue, + v2: &ScalarValue, +) -> Option { + match (v1, v2) { + (ScalarValue::Null, ScalarValue::Null) => Some(ScalarValue::Null), + (ScalarValue::Null, other) | (other, ScalarValue::Null) => Some(other.clone()), + _ => None, + } +} + +/// Prepend the grouping ID column to the output columns if present. +/// +/// For GROUPING SETS with no GROUP BY expressions, the schema includes a `__grouping_id` +/// column that must be present in the output. This function inserts it at the beginning +/// of the columns array to maintain schema alignment. +fn prepend_grouping_id_column( + mut columns: Vec>, + grouping_id: Option<&ScalarValue>, +) -> Result>> { + if let Some(id) = grouping_id { + let num_rows = columns.first().map(|array| array.len()).unwrap_or(1); + let grouping_ids = id.to_array_of_size(num_rows)?; + columns.insert(0, grouping_ids); + } + Ok(columns) +} + +impl AggregateStream { + /// Create a new AggregateStream + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + let agg_schema = Arc::clone(&agg.schema); + let agg_filter_expr = Arc::clone(&agg.filter_expr); + + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let input = agg.input.execute(partition, Arc::clone(context))?; + + let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; + let filter_expressions = match agg.mode.input_mode() { + AggregateInputMode::Raw => agg_filter_expr, + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), + }; + let accumulators = create_accumulators(&agg.aggr_expr)?; + + let reservation = MemoryConsumer::new(format!("AggregateStream[{partition}]")) + .register(context.memory_pool()); + + // Enable dynamic filter if: + // 1. AggregateExec did the check and ensure it supports the dynamic filter + // (its dynamic_filter field will be Some(..)) + // 2. Aggregate dynamic filter is enabled from the config + let mut maybe_dynamic_filter = match agg.dynamic_filter.as_ref() { + Some(filter) => Some(Arc::clone(filter)), + _ => None, + }; + + if !context + .session_config() + .options() + .optimizer + .enable_aggregate_dynamic_filter_pushdown + { + maybe_dynamic_filter = None; + } + + let inner = AggregateStreamInner { + schema: Arc::clone(&agg.schema), + mode: agg.mode, + input, + baseline_metrics, + aggregate_expressions, + filter_expressions, + accumulators, + reservation, + finished: false, + agg_dyn_filter_state: maybe_dynamic_filter, + }; + + let stream = futures::stream::unfold(inner, |mut this| async move { + if this.finished { + return None; + } + + loop { + let result = match this.input.next().await { + Some(Ok(batch)) => { + let result = { + let elapsed_compute = this.baseline_metrics.elapsed_compute(); + let _timer = elapsed_compute.timer(); // Stops on drop + aggregate_batch( + &this.mode, + &batch, + &mut this.accumulators, + &this.aggregate_expressions, + &this.filter_expressions, + ) + }; + + let result = result.and_then(|allocated| { + this.maybe_update_dyn_filter()?; + Ok(allocated) + }); + + // allocate memory + // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with + // overshooting a bit. Also this means we either store the whole record batch or not. + match result + .and_then(|allocated| this.reservation.try_grow(allocated)) + { + Ok(_) => continue, + Err(e) => Err(e), + } + } + Some(Err(e)) => Err(e), + None => { + this.finished = true; + // Release the input pipeline's resources before finalization. + let input_schema = this.input.schema(); + this.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + let timer = this.baseline_metrics.elapsed_compute().timer(); + let result = + finalize_aggregation(&mut this.accumulators, &this.mode) + .and_then(|columns| { + prepend_grouping_id_column(columns, None) + }) + .and_then(|columns| { + RecordBatch::try_new( + Arc::clone(&this.schema), + columns, + ) + .map_err(Into::into) + }) + .record_output(&this.baseline_metrics); + + timer.done(); + + result + } + }; + + this.finished = true; + return Some((result, this)); + } + }); + + // seems like some consumers call this stream even after it returned `None`, so let's fuse the stream. + let stream = stream.fuse(); + let stream = Box::pin(stream); + + Ok(Self { + schema: agg_schema, + stream, + }) + } +} + +impl Stream for AggregateStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + let this = &mut *self; + this.stream.poll_next_unpin(cx) + } +} + +impl RecordBatchStream for AggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Perform group-by aggregation for the given [`RecordBatch`]. +/// +/// If successful, this returns the additional number of bytes that were allocated during this process. +/// +/// TODO: Make this a member function +fn aggregate_batch( + mode: &AggregateMode, + batch: &RecordBatch, + accumulators: &mut [AccumulatorItem], + expressions: &[Vec>], + filters: &[Option>], +) -> Result { + let mut allocated = 0usize; + + // 1.1 iterate accumulators and respective expressions together + // 1.2 filter the batch if necessary + // 1.3 evaluate expressions + // 1.4 update / merge accumulators with the expressions' values + + // 1.1 + accumulators + .iter_mut() + .zip(expressions) + .zip(filters) + .try_for_each(|((accum, expr), filter)| { + // 1.2 + let batch = match filter { + Some(filter) => Cow::Owned(batch_filter(batch, filter)?), + None => Cow::Borrowed(batch), + }; + + // 1.3 + let values = evaluate_expressions_to_arrays(expr, batch.as_ref())?; + + // 1.4 + let size_pre = accum.size(); + let res = match mode.input_mode() { + AggregateInputMode::Raw => accum.update_batch(&values), + AggregateInputMode::Partial => accum.merge_batch(&values), + }; + let size_post = accum.size(); + allocated += size_post.saturating_sub(size_pre); + res + })?; + + Ok(allocated) +} diff --git a/datafusion/physical-plan/src/aggregates_blocked/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates_blocked/group_values/metrics.rs new file mode 100644 index 0000000000000..21e19027cebcf --- /dev/null +++ b/datafusion/physical-plan/src/aggregates_blocked/group_values/metrics.rs @@ -0,0 +1,579 @@ +// 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. + +//! Metrics for the various group-by implementations. + +use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; + +#[derive(Clone)] +pub(crate) struct AggregateArgumentMetrics { + argument_times: Vec