diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 1e23ec894c8e2..fea133a84a201 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1133,6 +1133,24 @@ config_namespace! { /// aggregation ratio check and trying to switch to skipping aggregation mode pub skip_partial_aggregation_probe_rows_threshold: usize, default = 100_000 + /// (experimental) Number of groups above which a hash aggregation + /// stops growing a single hash table, so that its tables stay small + /// enough to be cache friendly. A partial aggregation then emits the + /// state of its table and starts over, as long as the emitted groups + /// do not come back. A final aggregation splits the groups seen so + /// far and all further input into hash buckets, which are aggregated + /// one after another and can be spilled and released independently; + /// it does so at a quarter of this number when its input holds about + /// one row per group. Aggregations of millions of groups per partition + /// run faster and with less memory. Only aggregations whose group keys + /// and aggregate state are all fixed-width are affected: a string or + /// other variable-length value costs about as much to move into its + /// bucket as the smaller table saves. Moving rows has a cost of its + /// own, so aggregations that end at a few times this number of groups + /// can run a few percent slower, and input that repeats its groups can + /// use more memory. Set to 0 to disable. + pub hash_aggregate_bucket_threshold: usize, default = 0 + /// Should DataFusion use row number estimates at the input to decide /// whether increasing parallelism is beneficial or not. By default, /// only exact row numbers (not estimates) are used for this decision. diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index b25caf815eed1..d18a8fc08f023 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -40,8 +40,8 @@ use crate::aggregates::group_values::{ }; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - group_id_array, max_duplicate_ordinal, + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, group_id_array, max_duplicate_ordinal, }; use super::AggregateTableMetrics; @@ -136,6 +136,13 @@ pub(in crate::aggregates) struct AggregateHashTable { /// Lifecycle-specific state: building stage / outputting stage. pub(super) state: AggregateHashTableState, + /// Set by tables that are used for one set of groups after another (see + /// `AggregateHashTable::::restart`): when the output is + /// materialized, the emptied buffer is kept in `recycled_buffer` with + /// its allocations instead of being dropped. + pub(super) recycle_buffer: bool, + pub(super) recycled_buffer: Option, + pub(super) _mode: PhantomData, } @@ -148,16 +155,44 @@ impl AggregateHashTable { state_schema: SchemaRef, batch_size: usize, filters: Vec>>, + ) -> Result { + Self::new_for_input( + agg, + agg.input().schema(), + &agg.mode, + partition, + output_schema, + state_schema, + batch_size, + filters, + ) + } + + /// Like [`Self::new_with_filters`], for a table whose input is not the + /// input of `agg`: `input_schema` is the schema of the batches it + /// aggregates and `mode` decides how they are read, as raw rows or as + /// partial state. `agg.group_by` must refer to `input_schema`. + #[expect(clippy::too_many_arguments)] + pub(super) fn new_for_input( + agg: &AggregateExec, + input_schema: SchemaRef, + mode: &AggregateMode, + 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 metrics = AggregateTableMetrics::new(agg, partition); + // `mode` is a parameter rather than `agg.mode`: a bucket's table reads + // the partial state of an aggregation that is not its own. let aggregate_arguments = aggregate_expressions( agg.aggr_expr(), - &agg.mode, + mode, agg.group_by().num_group_exprs(), )?; let accumulators: Vec<_> = agg @@ -197,6 +232,8 @@ impl AggregateHashTable { batch_group_indices: Default::default(), accumulators, }), + recycle_buffer: false, + recycled_buffer: None, _mode: PhantomData, }) } @@ -319,6 +356,13 @@ impl AggregateHashTable { let batch = RecordBatch::try_new(output_schema, columns)?; debug_assert!(batch.num_rows() > 0); + if self.recycle_buffer { + // Keep the hash table's capacity for as many groups + // as it just held: the next set is likely as large. + state.group_values.clear_shrink(batch.num_rows()); + state.batch_group_indices.clear(); + self.recycled_buffer = Some(state); + } MaterializedAggregateOutput::new(batch) } AggregateHashTableState::OutputtingMaterialized(output) => output, @@ -340,23 +384,25 @@ impl AggregateHashTable { } 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::(); + let buffer_size = |state: &AggregateHashTableBuffer| { + 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() + acc + state.group_values.size() + state.batch_group_indices.allocated_size() + }; + let recycled = self.recycled_buffer.as_ref().map_or(0, buffer_size); + recycled + + match &self.state { + AggregateHashTableState::Building(state) + | AggregateHashTableState::Outputting(state) => buffer_size(state), + AggregateHashTableState::OutputtingMaterialized(output) => { + output.memory_size() + } + AggregateHashTableState::Done => 0, } - AggregateHashTableState::Done => 0, - } } /// Returns the number of distinct groups accumulated so far. @@ -372,6 +418,21 @@ impl AggregateHashTable { /// spilling without finalizing the same group more than once. pub(in crate::aggregates) fn take_state_batch( &mut self, + ) -> Result> { + self.take_state_batch_inner(false) + } + + /// Like [`Self::take_state_batch`], but keeps the table's capacity for as + /// many groups as it held, for a table that is filled again right away. + pub(in crate::aggregates) fn take_state_batch_keep_capacity( + &mut self, + ) -> Result> { + self.take_state_batch_inner(true) + } + + fn take_state_batch_inner( + &mut self, + keep_capacity: bool, ) -> Result> { let state_schema = Arc::clone(&self.state_schema); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); @@ -399,7 +460,9 @@ impl AggregateHashTable { // `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 + .group_values + .clear_shrink(if keep_capacity { batch.num_rows() } else { 0 }); state.batch_group_indices.clear(); state.batch_group_indices.shrink_to_fit(); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 4efd2a1c3ad75..60a4581a689f7 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -21,10 +21,12 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; -use crate::aggregates::AggregateExec; use crate::aggregates::group_values::AccumulatorPhase; +use crate::aggregates::{AggregateExec, AggregateMode}; -use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator}; +use super::common::{ + AggregateHashTable, AggregateHashTableState, FinalMarker, HashAggregateAccumulator, +}; /// Implementation specific to final aggregation, where the table stores partial /// aggregate states and the input rows are also partial states. @@ -50,6 +52,30 @@ impl AggregateHashTable { ) } + /// A table that merges partial state rows of `state_schema` for `agg`, + /// which does not have to be a final aggregation itself: a single stage + /// aggregation passes a copy of itself whose `group_by` refers to the + /// state columns (see `PhysicalGroupBy::as_final`), as it does to replay + /// its spills. + pub(in crate::aggregates) fn new_over_state( + agg: &AggregateExec, + state_schema: &SchemaRef, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_for_input( + agg, + Arc::clone(state_schema), + &AggregateMode::Final, + partition, + output_schema, + Arc::clone(state_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`. @@ -78,6 +104,29 @@ impl AggregateHashTable { ) } + /// Makes the table reusable through [`Self::restart`]. + pub(in crate::aggregates) fn with_restart(mut self) -> Self { + self.recycle_buffer = true; + self + } + + /// After all output has been taken, goes back to aggregating a new set of + /// groups with the allocations of the previous one: a fresh table grows + /// from a few entries, rehashing every group it holds each time it + /// doubles. Returns false if the table cannot be reused. + pub(in crate::aggregates) fn restart(&mut self) -> bool { + match (&self.state, self.recycled_buffer.take()) { + (AggregateHashTableState::Done, Some(buffer)) => { + self.state = AggregateHashTableState::Building(buffer); + true + } + (_, buffer) => { + self.recycled_buffer = buffer; + false + } + } + } + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { self.start_outputting(); Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 397b766f41697..13ef352ff869d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -102,6 +102,8 @@ impl AggregateHashTable { batch_group_indices: Default::default(), accumulators, }), + recycle_buffer: false, + recycled_buffer: None, _mode: PhantomData, }) } diff --git a/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs b/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs new file mode 100644 index 0000000000000..ff3e650eb4b04 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs @@ -0,0 +1,407 @@ +// 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. + +//! Aggregating partial state that was split into [`FinalBuckets`]. +//! +//! See [`BucketedAggregation`]. + +use std::sync::Arc; + +use arrow::compute::BatchCoalescer; +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{TryEmitter, async_try_stream}; +use futures::StreamExt; +use futures::future::BoxFuture; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, FinalMarker}; +use super::final_buckets::{FinalBuckets, MAX_BUCKET_LEVELS}; +use crate::SendableRecordBatchStream; +use crate::metrics; +use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::RecordBatchStreamAdapter; + +/// What a hash aggregation stream needs to split the groups of a hash table +/// that has reached `hash_aggregate_bucket_threshold` into [`FinalBuckets`], +/// and to aggregate those buckets one after another. +/// +/// The buckets hold partial state rows, whatever the input of the stream is: +/// a final aggregation routes its input as is, a single stage aggregation +/// routes the state of the table that aggregates its raw input. +pub(super) struct BucketedAggregation { + /// Number of groups in one hash table that triggers bucketing + threshold: usize, + /// Aggregate configuration used to construct the table of each bucket: + /// its `group_by` refers to the columns of `state_schema`. + agg: AggregateExec, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Schema of the partial state rows held by the buckets. + state_schema: SchemaRef, + /// Schema of the aggregation's output. + output_schema: SchemaRef, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_manager: Option, + /// Number of times a hash table was split into buckets + bucket_splits: metrics::Count, + /// Number of times the rows of a bucket were replaced by their aggregated state + bucket_compactions: metrics::Count, +} + +/// What [`BucketedAggregation::output_stream`] works with while it runs. +struct BucketOutput { + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + /// Combines the output of small buckets into batches of the target size. + coalescer: BatchCoalescer, +} + +impl BucketedAggregation { + /// `agg` must aggregate rows of `state_schema` as partial state: see + /// [`AggregateHashTable::::new_over_state`]. + pub(super) fn new( + threshold: usize, + agg: AggregateExec, + partition: usize, + batch_size: usize, + state_schema: SchemaRef, + output_schema: SchemaRef, + spill_manager: Option, + ) -> Self { + let bucket_splits = + MetricBuilder::new(&agg.metrics).counter("bucket_splits", partition); + let bucket_compactions = + MetricBuilder::new(&agg.metrics).counter("bucket_compactions", partition); + Self { + threshold, + agg, + partition, + batch_size, + state_schema, + output_schema, + spill_manager, + bucket_splits, + bucket_compactions, + } + } + + /// True if a schema of partial state rows can be bucketed. + /// + /// Bucketing moves every row once more before it is aggregated, which only + /// pays when moving a row is cheap compared to the cache miss it saves. + /// That holds for fixed-width columns. Variable-length keys or state + /// (strings, lists, ...) cost as much to move as the smaller tables save, + /// so such aggregations keep their single table. + pub(super) fn supports_state(state_schema: &SchemaRef) -> bool { + state_schema.fields().iter().all(|field| { + let data_type = field.data_type(); + data_type.is_primitive() || *data_type == DataType::Boolean + }) + } + + pub(super) fn threshold(&self) -> usize { + self.threshold + } + + fn new_table(&self) -> Result> { + AggregateHashTable::::new_over_state( + &self.agg, + &self.state_schema, + self.partition, + Arc::clone(&self.output_schema), + self.batch_size, + ) + } + + /// Starts the buckets of `level` with `state`, the state of the table + /// that is split. `kept` is the share of its input rows which that table + /// kept as groups. + pub(super) fn split( + &self, + level: u32, + state: Option, + kept: f64, + ) -> Result { + let mut buckets = FinalBuckets::new( + &self.state_schema, + self.agg.group_by().num_group_exprs(), + self.batch_size, + level, + self.spill_manager.clone(), + ); + self.bucket_splits.add(1); + buckets.expect_kept(kept); + if let Some(state) = state { + buckets.route(&state)?; + } + Ok(buckets) + } + + /// Replaces the rows of every bucket that is due for it by their aggregated + /// state. See the compaction section of [`FinalBuckets`]. + pub(super) fn compact( + &self, + buckets: &mut FinalBuckets, + table: &mut Option>, + ) -> Result<()> { + while let Some(index) = buckets.bucket_to_compact() { + let table = match table { + Some(table) => table, + None => table.insert(self.new_table()?), + }; + let mut input_rows = 0; + for batch in buckets.take_bucket(index)? { + input_rows += batch.num_rows(); + table.aggregate_batch(&batch)?; + } + buckets.put_compacted( + index, + table.take_state_batch_keep_capacity()?, + input_rows, + ); + self.bucket_compactions.add(1); + } + Ok(()) + } + + /// Reserves `other_bytes` plus the memory of `buckets`, spilling buckets + /// for as long as the reservation does not fit. + pub(super) fn reserve( + &self, + reservation: &MemoryReservation, + other_bytes: usize, + buckets: &mut FinalBuckets, + ) -> Result<()> { + loop { + let size = other_bytes.saturating_add(buckets.memory_size()); + match reservation.try_resize(size) { + Ok(()) => return Ok(()), + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + if buckets.spill_largest()? { + continue; + } + // Every bucket is on disk. What is left is the fixed cost + // of routing a batch, which no spill can release, so go on + // like the sort based spill path does after it has spilled + // its table: with what the pool still grants, which covers + // at least the memory held besides these buckets. + if buckets.is_fully_spilled() + && reservation.try_resize(other_bytes).is_ok() + { + return Ok(()); + } + return Err(e.context("Hash aggregate has no more buckets to spill")); + } + Err(e) => return Err(e), + } + } + } + + /// Aggregates the buckets one after another, each with a table of its + /// own size, and emits the groups of a bucket before reading the next one. + /// + /// A bucket whose table reaches the bucketing threshold again, or does + /// not fit in memory, is split into buckets of the next level. + pub(super) fn output_stream( + self: Arc, + buckets: FinalBuckets, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + ) -> SendableRecordBatchStream { + let schema = Arc::clone(&self.output_schema); + let batch_size = self.batch_size; + let stream = async_try_stream(move |mut emitter| async move { + let mut output = BucketOutput { + reservation, + baseline_metrics, + coalescer: BatchCoalescer::new( + Arc::clone(&self.output_schema), + batch_size, + ) + .with_biggest_coalesce_batch_size(Some(batch_size / 2)), + }; + self.produce_output(buckets, &mut output, &mut emitter) + .await?; + output.reservation.try_resize(0)?; + + output.coalescer.finish_buffered_batch()?; + while let Some(batch) = output.coalescer.next_completed_batch() { + emitter + .emit(batch.record_output(&output.baseline_metrics)) + .await; + } + Ok(()) + }); + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) + } + + fn produce_output<'a>( + &'a self, + buckets: FinalBuckets, + output: &'a mut BucketOutput, + emitter: &'a mut TryEmitter, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let elapsed_compute = output.baseline_metrics.elapsed_compute().clone(); + let next_level = buckets.level() + 1; + + let sources = buckets.into_sources()?; + // Memory of the buckets that wait for their turn + let mut waiting_bytes: usize = + sources.iter().map(|source| source.memory_size()).sum(); + + // One table aggregates all the buckets, one after another + let mut reusable_table = None; + for source in sources { + let mut source_bytes = source.memory_size(); + waiting_bytes -= source_bytes; + let mut input = source.into_stream(&self.state_schema); + + let mut timer = elapsed_compute.timer(); + let mut hash_table = match reusable_table.take() { + Some(hash_table) => hash_table, + None => self.new_table()?.with_restart(), + }; + let mut table_rows = 0usize; + let mut sub_buckets: Option = None; + let mut compaction_table = None; + + while let Some(batch) = input.next().await.transpose()? { + // Batches of an in-memory bucket are released as they are read + source_bytes = + source_bytes.saturating_sub(batch.get_array_memory_size()); + let held_bytes = waiting_bytes + source_bytes; + + if let Some(sub_buckets) = sub_buckets.as_mut() { + sub_buckets.route(&batch)?; + self.compact(sub_buckets, &mut compaction_table)?; + self.reserve(&output.reservation, held_bytes, sub_buckets)?; + continue; + } + + hash_table.aggregate_batch(&batch)?; + table_rows += batch.num_rows(); + + let can_split = next_level < MAX_BUCKET_LEVELS; + let split = match output + .reservation + .try_resize(held_bytes + hash_table.memory_size()) + { + Ok(()) => { + can_split + && hash_table.building_group_count() >= self.threshold + } + Err(DataFusionError::ResourcesExhausted(_)) if can_split => true, + Err(e) => return Err(e), + }; + if split { + let kept = hash_table.building_group_count() as f64 + / table_rows.max(1) as f64; + let mut new_buckets = + self.split(next_level, hash_table.take_state_batch()?, kept)?; + self.reserve( + &output.reservation, + held_bytes + hash_table.memory_size(), + &mut new_buckets, + )?; + sub_buckets = Some(new_buckets); + } + } + drop(input); + + if let Some(sub_buckets) = sub_buckets { + // The table handed its groups over and is empty again + reusable_table = Some(hash_table); + timer.done(); + self.produce_output(sub_buckets, output, emitter).await?; + continue; + } + + hash_table.start_output()?; + while let Some(batch) = hash_table.next_output_batch()? { + output + .reservation + .try_resize(waiting_bytes + hash_table.memory_size())?; + output.coalescer.push_batch(batch)?; + while let Some(batch) = output.coalescer.next_completed_batch() { + timer.done(); + emitter + .emit(batch.record_output(&output.baseline_metrics)) + .await; + timer = elapsed_compute.timer(); + } + } + if hash_table.restart() { + reusable_table = Some(hash_table); + } + timer.done(); + } + + Ok(()) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::datatypes::{Field, Fields, Schema}; + + fn schema(types: Vec) -> SchemaRef { + let fields: Fields = types + .into_iter() + .enumerate() + .map(|(i, data_type)| Field::new(format!("c{i}"), data_type, true)) + .collect(); + Arc::new(Schema::new(fields)) + } + + #[test] + fn only_fixed_width_state_is_bucketed() { + assert!(BucketedAggregation::supports_state(&schema(vec![ + DataType::Int64, + DataType::Boolean, + DataType::Decimal128(38, 10), + DataType::Float64, + DataType::Date32, + ]))); + + // A variable-length group key or state column turns bucketing off + for variable_length in [ + DataType::Utf8, + DataType::Utf8View, + DataType::LargeBinary, + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))), + ] { + assert!(!BucketedAggregation::supports_state(&schema(vec![ + DataType::Int64, + variable_length.clone(), + ]))); + assert!(!BucketedAggregation::supports_state(&schema(vec![ + variable_length, + DataType::Int64, + ]))); + } + } +} diff --git a/datafusion/physical-plan/src/aggregates/final_buckets.rs b/datafusion/physical-plan/src/aggregates/final_buckets.rs new file mode 100644 index 0000000000000..63c9601ed414b --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/final_buckets.rs @@ -0,0 +1,530 @@ +// 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. + +//! Hash buckets of partial aggregate state for the final hash aggregation. +//! +//! See [`FinalBuckets`]. + +use std::collections::VecDeque; +use std::sync::Arc; + +use arrow::array::{ArrayRef, PrimitiveArray}; +use arrow::compute::{BatchCoalescer, take_arrays}; +use arrow::datatypes::{SchemaRef, UInt32Type}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use datafusion_common::Result; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_execution::async_try_stream; +use futures::StreamExt; + +use crate::SendableRecordBatchStream; +use crate::spill::in_progress_spill_file::InProgressSpillFile; +use crate::spill::spill_manager::SpillManager; +use crate::stream::RecordBatchStreamAdapter; + +/// Number of bits of the routing hash consumed by one bucketing level. +const BUCKET_BITS: u32 = 6; + +/// Number of buckets rows are split into at each level. +const NUM_BUCKETS: usize = 1 << BUCKET_BITS; + +/// A bucket that is still too large is split again with the next hash bits, +/// up to this many levels. +pub(super) const MAX_BUCKET_LEVELS: u32 = 4; + +/// Seed of the routing hash. It differs from both the `RepartitionExec` seed, +/// whose hash is the same for every row of one final partition modulo the +/// partition count, and the aggregation seed, so that the rows of one bucket +/// still spread evenly over the bucket's own hash table. +const BUCKET_HASH_SEED: RandomState = RandomState::with_seed(5364907223173859721); + +/// A bucket is compacted for the first time once it holds this many batches +/// worth of rows. With 64 buckets per partition this floor is what the +/// buckets hold when every bucket has few groups, so it is kept small. +const MIN_COMPACTION_BATCHES: usize = 1; + +/// A compaction that kept more than this share of its rows found (almost) no +/// repeated groups. +const POOR_COMPACTION: f64 = 0.8; + +/// After a compaction that paid off, the bucket is compacted again once it +/// has received at most this many times its own rows. Measured on unreduced +/// input with 5 to 8 rows per group: 2 holds peak memory at about 1.35x of a +/// single table, 16 lets it reach 2.3x. +const MAX_COMPACTION_FACTOR: f64 = 2.0; + +/// After a poor compaction the bucket waits for this many times its own rows. +/// Aggregating rows that are already unique again and again is pure rework +/// (1.1x -> 1.5x run time on a wide string key with a factor of 2). But a +/// small sample of a bucket with many groups looks unique even when every +/// group repeats, so the wait must stay short enough to find that out before +/// the whole input is buffered (16 let peak memory reach 1.8x of a single +/// table on input with 5 rows per group). +const POOR_COMPACTION_FACTOR: f64 = 4.0; + +/// One hash bucket: its completed batches, held in memory or appended to a +/// spill file, plus the rows not yet forming a full batch. +struct Bucket { + coalescer: BatchCoalescer, + batches: VecDeque, + /// Memory held by `batches` + batches_bytes: usize, + /// Set once the bucket has been spilled: completed batches go here + spill_file: Option, + /// Rows routed to the bucket that are still held in memory + rows: usize, + /// The bucket is due for compaction once it holds this many rows + compact_at: usize, +} + +/// Rows of partial aggregate state, split by a hash of their group keys. +/// +/// Once the table of a final hash aggregation has outgrown the CPU caches, +/// every probe and every accumulator update of a further row is a cache +/// miss. The stream then stops growing that table: it moves the table's +/// state in here, routes the rest of its input here without aggregating it, +/// and finally aggregates the buckets one after another, each with a table +/// that is `NUM_BUCKETS` times smaller. Buckets are independent, so they can +/// be spilled, read back and released one at a time. +/// +/// # Compaction +/// +/// The input may hold the same group many times (for example when the partial +/// aggregation was skipped), so buffering it as is could take far more memory +/// than the single table it replaces. A bucket that has grown enough is +/// therefore *compacted*: the stream aggregates its rows with a small table +/// and puts the resulting state, one row per group, back in their place. The +/// next compaction waits for as many new rows as made the last one pay off, +/// so the buffered rows stay within a small multiple of the groups they hold, +/// and waits much longer after a compaction that found nothing to merge, so +/// input that does not reduce is not aggregated over and over. +pub(super) struct FinalBuckets { + /// Bucketing level, selects which bits of the routing hash are used + level: u32, + /// Number of leading columns of a batch that are the group keys + num_group_columns: usize, + buckets: Vec, + /// `None` if spilling is not supported by the configured `DiskManager` + spill_manager: Option, + batch_size: usize, + min_compaction_rows: usize, + /// Reused buffers + hashes: Vec, + bucket_sizes: Vec, + reordered_indices: Vec, +} + +impl FinalBuckets { + pub(super) fn new( + schema: &SchemaRef, + num_group_columns: usize, + batch_size: usize, + level: u32, + spill_manager: Option, + ) -> Self { + debug_assert!(level < MAX_BUCKET_LEVELS); + let min_compaction_rows = MIN_COMPACTION_BATCHES * batch_size; + let buckets = (0..NUM_BUCKETS) + .map(|_| Bucket { + coalescer: BatchCoalescer::new(Arc::clone(schema), batch_size), + batches: VecDeque::new(), + batches_bytes: 0, + spill_file: None, + rows: 0, + compact_at: min_compaction_rows, + }) + .collect(); + Self { + level, + num_group_columns, + buckets, + spill_manager, + batch_size, + min_compaction_rows, + hashes: vec![], + bucket_sizes: vec![0; NUM_BUCKETS], + reordered_indices: vec![], + } + } + + pub(super) fn level(&self) -> u32 { + self.level + } + + /// Tells the buckets what share of its input rows the table they are + /// split from kept as groups. If that table found (almost) no repeated + /// groups, the first compaction of a bucket would most likely find none + /// either, so it waits as long as after a poor compaction. + pub(super) fn expect_kept(&mut self, kept: f64) { + if kept > POOR_COMPACTION { + let compact_at = + (self.min_compaction_rows as f64 * POOR_COMPACTION_FACTOR) as usize; + for bucket in &mut self.buckets { + bucket.compact_at = compact_at; + } + } + } + + /// Appends every row of `batch` to the bucket of its group key. + pub(super) fn route(&mut self, batch: &RecordBatch) -> Result<()> { + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + let group_columns: &[ArrayRef] = &batch.columns()[..self.num_group_columns]; + self.hashes.clear(); + self.hashes.resize(num_rows, 0); + create_hashes(group_columns, &BUCKET_HASH_SEED, &mut self.hashes)?; + + // Counting sort of the row indices by bucket + let shift = self.level * BUCKET_BITS; + let bucket_of = |hash: u64| ((hash >> shift) as usize) & (NUM_BUCKETS - 1); + self.bucket_sizes.fill(0); + for &hash in &self.hashes { + self.bucket_sizes[bucket_of(hash)] += 1; + } + let mut starts = [0u32; NUM_BUCKETS]; + let mut next = 0; + for (start, size) in starts.iter_mut().zip(&self.bucket_sizes) { + *start = next; + next += size; + } + self.reordered_indices.clear(); + self.reordered_indices.resize(num_rows, 0); + let mut cursors = starts; + for (row, &hash) in self.hashes.iter().enumerate() { + let cursor = &mut cursors[bucket_of(hash)]; + self.reordered_indices[*cursor as usize] = row as u32; + *cursor += 1; + } + + // One gather for the whole batch, then a slice per bucket. The + // coalescer copies the slice, so a bucket owns its batches. + let indices: PrimitiveArray = + std::mem::take(&mut self.reordered_indices).into(); + let columns = take_arrays(batch.columns(), &indices, None)?; + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + let reordered = + RecordBatch::try_new_with_options(batch.schema(), columns, &options)?; + + for (bucket, (&start, &size)) in self + .buckets + .iter_mut() + .zip(starts.iter().zip(&self.bucket_sizes)) + { + if size == 0 { + continue; + } + let slice = reordered.slice(start as usize, size as usize); + match &mut bucket.spill_file { + // A spilled bucket holds no memory: its rows go straight to + // the file and are combined into batches when read back. + Some(spill_file) => { + spill_file.append_batch(&slice)?; + } + None => { + bucket.coalescer.push_batch(slice)?; + bucket.rows += size as usize; + bucket.collect_completed()?; + } + } + } + Ok(()) + } + + /// Returns a bucket that is due for compaction, if any. + pub(super) fn bucket_to_compact(&self) -> Option { + self.buckets.iter().position(|bucket| { + bucket.spill_file.is_none() && bucket.rows >= bucket.compact_at + }) + } + + /// Removes and returns the rows of an in-memory bucket, to be followed by + /// [`Self::put_compacted`]. + pub(super) fn take_bucket(&mut self, index: usize) -> Result> { + let bucket = &mut self.buckets[index]; + debug_assert!(bucket.spill_file.is_none()); + bucket.coalescer.finish_buffered_batch()?; + bucket.collect_completed()?; + bucket.batches_bytes = 0; + Ok(bucket.batches.drain(..).collect()) + } + + /// Puts back the aggregated state of the `input_rows` rows that + /// [`Self::take_bucket`] returned. + pub(super) fn put_compacted( + &mut self, + index: usize, + state: Option, + input_rows: usize, + ) { + let bucket = &mut self.buckets[index]; + let state_rows = state.as_ref().map_or(0, |state| state.num_rows()); + if let Some(state) = state { + bucket.batches_bytes += state.get_array_memory_size(); + bucket.batches.push_back(state); + } + bucket.rows = state_rows; + + // Compacting again costs about `state_rows + new_rows` and removes + // about `(1 - kept) * new_rows` rows, where `kept` is the share of rows + // that survived this time. Wait for enough new rows to pay for it. + let kept = state_rows as f64 / input_rows.max(1) as f64; + let factor = if kept > POOR_COMPACTION { + POOR_COMPACTION_FACTOR + } else { + (kept / (1.0 - kept)).clamp(1.0, MAX_COMPACTION_FACTOR) + }; + let new_rows = + ((state_rows as f64 * factor) as usize).max(self.min_compaction_rows); + bucket.compact_at = state_rows + new_rows; + } + + /// Memory held by the buckets + pub(super) fn memory_size(&self) -> usize { + self.buckets + .iter() + .map(|bucket| bucket.batches_bytes + bucket.coalescer.size()) + .sum::() + + self.hashes.capacity() * size_of::() + } + + /// Spills the in-memory batches of the largest bucket, which from then on + /// appends its batches to the spill file. Returns false if nothing can be + /// spilled. + pub(super) fn spill_largest(&mut self) -> Result { + let Some(spill_manager) = &self.spill_manager else { + return Ok(false); + }; + let Some(bucket) = self + .buckets + .iter_mut() + .filter(|bucket| bucket.spill_file.is_none() && bucket.rows > 0) + .max_by_key(|bucket| bucket.batches_bytes + bucket.coalescer.size()) + else { + return Ok(false); + }; + bucket.spill_file = + Some(spill_manager.create_in_progress_file("FinalHashAggregateBucket")?); + // Rows that do not form a full batch yet are spilled as well, and the + // coalescer is replaced to give up the buffers it allocated. + bucket.coalescer.finish_buffered_batch()?; + bucket.collect_completed()?; + bucket.coalescer = BatchCoalescer::new(bucket.coalescer.schema(), 1); + Ok(true) + } + + /// True if no bucket holds rows in memory. + pub(super) fn is_fully_spilled(&self) -> bool { + self.buckets.iter().all(|bucket| bucket.rows == 0) + } + + /// Finishes the buckets and returns the non-empty ones. + pub(super) fn into_sources(self) -> Result> { + let mut sources = Vec::with_capacity(NUM_BUCKETS); + for mut bucket in self.buckets { + bucket.coalescer.finish_buffered_batch()?; + bucket.collect_completed()?; + match bucket.spill_file { + Some(mut spill_file) => { + let spill_manager = self + .spill_manager + .as_ref() + .expect("a spilled bucket has a spill manager"); + if let Some(file) = spill_file.finish()? { + sources.push(BucketSource::Spilled { + stream: spill_manager.read_spill_as_stream(file, None)?, + batch_size: self.batch_size, + }); + } + } + None if bucket.batches.is_empty() => {} + None => sources.push(BucketSource::Memory { + batches: bucket.batches, + bytes: bucket.batches_bytes, + }), + } + } + Ok(sources) + } +} + +impl Bucket { + /// Moves the batches the coalescer has completed to where the bucket + /// lives: the spill file if it has been spilled, memory otherwise. + fn collect_completed(&mut self) -> Result<()> { + while let Some(batch) = self.coalescer.next_completed_batch() { + self.batches_bytes += batch.get_array_memory_size(); + self.batches.push_back(batch); + } + if let Some(spill_file) = &mut self.spill_file { + for batch in self.batches.drain(..) { + spill_file.append_batch(&batch)?; + } + self.batches_bytes = 0; + self.rows = 0; + } + Ok(()) + } +} + +/// The rows of one finished bucket +pub(super) enum BucketSource { + Memory { + batches: VecDeque, + bytes: usize, + }, + /// Rows were appended to the file in pieces much smaller than a batch + Spilled { + stream: SendableRecordBatchStream, + batch_size: usize, + }, +} + +impl BucketSource { + /// Memory held by the bucket's rows + pub(super) fn memory_size(&self) -> usize { + match self { + Self::Memory { bytes, .. } => *bytes, + Self::Spilled { .. } => 0, + } + } + + pub(super) fn into_stream(self, schema: &SchemaRef) -> SendableRecordBatchStream { + match self { + Self::Memory { batches, .. } => Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(schema), + futures::stream::iter(batches.into_iter().map(Ok)), + )), + Self::Spilled { stream, batch_size } => { + coalesce_stream(stream, Arc::clone(schema), batch_size) + } + } + } +} + +/// Combines the batches of `input` into batches of `batch_size` rows. +fn coalesce_stream( + mut input: SendableRecordBatchStream, + schema: SchemaRef, + batch_size: usize, +) -> SendableRecordBatchStream { + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + let stream = async_try_stream(|mut emitter| async move { + while let Some(batch) = input.next().await.transpose()? { + coalescer.push_batch(batch)?; + while let Some(batch) = coalescer.next_completed_batch() { + emitter.emit(batch).await; + } + } + coalescer.finish_buffered_batch()?; + while let Some(batch) = coalescer.next_completed_batch() { + emitter.emit(batch).await; + } + Ok(()) + }); + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{AsArray, Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; + use std::collections::HashMap; + + fn test_batch(schema: &SchemaRef, keys: std::ops::Range) -> RecordBatch { + let names: StringArray = keys + .clone() + .map(|k| Some(format!("key-{}", k % 97))) + .collect(); + let values: Int64Array = keys.clone().map(|k| k * 10).collect::>().into(); + let keys: Int64Array = keys.collect::>().into(); + RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(keys), Arc::new(names), Arc::new(values)], + ) + .unwrap() + } + + #[tokio::test] + async fn routes_every_row_once_and_keeps_keys_together() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("v", DataType::Int64, false), + ])); + let mut buckets = FinalBuckets::new(&schema, 2, 1024, 0, None); + // the same keys arrive in two different batches + buckets.route(&test_batch(&schema, 0..20_000))?; + buckets.route(&test_batch(&schema, 10_000..30_000))?; + assert!(buckets.memory_size() > 0); + assert!( + !buckets.spill_largest()?, + "no spill manager, nothing spilled" + ); + + let sources = buckets.into_sources()?; + assert!( + sources.len() > NUM_BUCKETS / 2, + "rows spread over the buckets" + ); + + let mut bucket_of_key: HashMap = HashMap::new(); + let mut rows = 0; + for (bucket, source) in sources.into_iter().enumerate() { + let mut stream = source.into_stream(&schema); + while let Some(batch) = stream.next().await.transpose()? { + rows += batch.num_rows(); + let keys = batch.column(0).as_primitive::(); + let values = batch.column(2).as_primitive::(); + for (key, value) in keys.values().iter().zip(values.values()) { + assert_eq!(*value, key * 10, "columns stay aligned"); + let seen = *bucket_of_key.entry(*key).or_insert(bucket); + assert_eq!(seen, bucket, "key {key} in two buckets"); + } + } + } + assert_eq!(rows, 40_000); + assert_eq!(bucket_of_key.len(), 30_000); + Ok(()) + } + + #[test] + fn levels_use_different_hash_bits() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("v", DataType::Int64, false), + ])); + // Rows of one level-0 bucket must spread again at level 1 + let mut level0 = FinalBuckets::new(&schema, 1, 1024, 0, None); + level0.route(&test_batch(&schema, 0..50_000))?; + let source = level0.into_sources()?.swap_remove(0); + let BucketSource::Memory { batches, .. } = source else { + unreachable!() + }; + let mut level1 = FinalBuckets::new(&schema, 1, 1024, 1, None); + for batch in &batches { + level1.route(batch)?; + } + assert!(level1.into_sources()?.len() > NUM_BUCKETS / 2); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2f1bb9b638a4c..9635b636aff92 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -20,11 +20,13 @@ //! See comments in [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] //! for details. +use std::collections::{HashMap, VecDeque}; use std::mem::size_of; use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::hash_utils::{RandomState, create_hashes}; use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, }; @@ -37,6 +39,8 @@ use super::aggregate_hash_table::{ AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, PartialSkipMarker, }; +use super::bucketed_aggregation::BucketedAggregation; +use super::final_buckets::FinalBuckets; use super::skip_partial::SkipAggregationProbe; use super::spill::AggregateSpill; use crate::metrics::{ @@ -159,6 +163,9 @@ pub(crate) struct PartialHashAggregateStream { /// Number of times accumulated states were emitted due to memory pressure. early_emit_count: metrics::Count, + /// `None` unless `hash_aggregate_bucket_threshold` is set and applies. + table_flush: Option, + /// Tracks whether partial aggregation should switch to direct state conversion. skip_aggregation_probe: Option, @@ -199,6 +206,8 @@ pub(crate) struct FinalHashAggregateStream { hash_table: Option>, /// `None` if spilling is not supported by the configured `DiskManager`. spill_context: Option>, + /// `None` unless `hash_aggregate_bucket_threshold` is set and applies. + bucketing: Option>, } #[derive(PartialEq)] @@ -207,9 +216,143 @@ enum HandleInputResult { ReachedLimit, #[expect(clippy::upper_case_acronyms)] OOM, + /// The table reached the bucket threshold, see [`PartialTableFlush`] + TableFull, SwitchToSkipAggregation, } +/// A final table whose input is [`BARELY_REDUCING`] moves into buckets once +/// it holds this share of the bucket threshold. +/// +/// Moving a table into buckets aggregates its groups a second time. An input +/// of about one row per group never earns that back unless the table is +/// small compared to the input, and such an input gains nothing from a larger +/// first table either. +const EARLY_BUCKETS_DIVISOR: usize = 4; + +/// Share of the rows seen so far that started a new group above which the +/// input is taken to hold about one row per group, the same ratio as the +/// default of `skip_partial_aggregation_probe_ratio_threshold`. +const BARELY_REDUCING: f64 = 0.8; + +/// Whether a final table of `groups` groups built from `rows` rows moves into +/// buckets. +/// +/// An input that repeats its groups keeps its table up to the full threshold: +/// the table reduces that input, which buckets only do by compacting. +fn starts_buckets(groups: usize, rows: usize, threshold: usize) -> bool { + groups >= threshold + || (groups >= (threshold / EARLY_BUCKETS_DIVISOR).max(1) + && groups as f64 >= BARELY_REDUCING * rows as f64) +} + +/// Number of flushed groups whose hashes are kept to detect recurring groups. +const FLUSH_SAMPLE_SIZE: usize = 1024; + +/// Number of earlier flushes whose samples are kept, so that groups which +/// only come back after many flushes (keys that cycle with a long period) are +/// noticed as well. Costs at most `64 * 1024` remembered hashes. +const FLUSH_SAMPLES_KEPT: usize = 64; + +/// Flushing stops once more than this share of the sampled groups of an +/// earlier flush shows up again in a later one. +const MAX_RECURRING_GROUPS: f64 = 0.2; + +/// Seed for the hashes of sampled groups, only compared with each other. +const FLUSH_SAMPLE_SEED: RandomState = RandomState::with_seed(8122871950429871369); + +/// Keeps the table of a partial hash aggregation small: once it holds +/// `hash_aggregate_bucket_threshold` groups, its state is emitted downstream +/// and the table starts over, the same way it does under memory pressure. +/// +/// A table that has outgrown the CPU caches pays a cache miss for every probe +/// and every accumulator update, so capping it keeps the partial aggregation +/// fast; the final aggregation merges whatever is emitted more than once. +/// +/// That only pays while flushing is free, that is while a flushed group does +/// not come back: a group that returns is emitted again, and the reduction +/// the partial stage exists for is lost. Each flush therefore remembers a +/// sample of its groups and the following flushes count how many of them they +/// hold again. Sorted, clustered or mostly unique keys never recur and keep being +/// flushed; keys that recur turn flushing off for the rest of the stream, +/// which then grows one table as before. +struct PartialTableFlush { + threshold: usize, + /// Number of leading columns of the state batch that are the group keys + num_group_columns: usize, + /// Set once flushed groups were seen to recur + disabled: bool, + /// Hashes of a sample of the groups emitted by the last flushes, with the + /// number of the flush that emitted them + sampled_groups: HashMap, + /// `(flush number, sample size)` of the flushes in `sampled_groups` + sampled_flushes: VecDeque<(usize, usize)>, + /// Number of flushes so far + num_flushes: usize, + /// Groups emitted so far, which the skip aggregation probe must still count + flushed_groups: usize, + flush_count: metrics::Count, + hashes: Vec, +} + +impl PartialTableFlush { + fn should_flush(&self, num_groups: usize) -> bool { + !self.disabled && num_groups >= self.threshold + } + + /// Records the flush of `state`, the emitted groups and their states. + fn record_flush(&mut self, state: &RecordBatch) -> Result<()> { + let num_groups = state.num_rows(); + self.flush_count.add(1); + self.flushed_groups += num_groups; + + self.hashes.clear(); + self.hashes.resize(num_groups, 0); + create_hashes( + &state.columns()[..self.num_group_columns], + &FLUSH_SAMPLE_SEED, + &mut self.hashes, + )?; + + // How many sampled groups of each earlier flush are in this one? + let oldest = self.sampled_flushes.front().map_or(0, |(flush, _)| *flush); + let mut recurring = vec![0usize; self.sampled_flushes.len()]; + for hash in &self.hashes { + if let Some(flush) = self.sampled_groups.get(hash) { + recurring[flush - oldest] += 1; + } + } + let groups_recur = recurring.iter().zip(&self.sampled_flushes).any( + |(recurring, (_, sample_size))| { + *recurring as f64 > MAX_RECURRING_GROUPS * *sample_size as f64 + }, + ); + if groups_recur { + self.disabled = true; + self.sampled_groups = HashMap::new(); + self.sampled_flushes = VecDeque::new(); + self.hashes = vec![]; + return Ok(()); + } + + if self.sampled_flushes.len() == FLUSH_SAMPLES_KEPT + && let Some((evicted, _)) = self.sampled_flushes.pop_front() + { + self.sampled_groups.retain(|_, flush| *flush != evicted); + } + let step = num_groups.div_ceil(FLUSH_SAMPLE_SIZE).max(1); + let mut sample_size = 0; + for hash in self.hashes.iter().step_by(step) { + self.sampled_groups.insert(*hash, self.num_flushes); + sample_size += 1; + } + self.sampled_flushes + .push_back((self.num_flushes, sample_size)); + self.num_flushes += 1; + Ok(()) + } +} + impl PartialHashAggregateStream { pub fn new( agg: &AggregateExec, @@ -232,6 +375,32 @@ impl PartialHashAggregateStream { let early_emit_count = MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition); + let group_values_soft_limit = agg.limit_options().map(|config| config.limit()); + let bucket_threshold = context + .session_config() + .options() + .execution + .hash_aggregate_bucket_threshold; + let num_group_columns = agg.group_by().num_group_exprs(); + // Same conditions as for bucketing in the final aggregation, which + // receives what is flushed here: see `FinalHashAggregateStream::new`. + let fixed_width_state = BucketedAggregation::supports_state(&schema); + let table_flush = (bucket_threshold > 0 + && group_values_soft_limit.is_none() + && fixed_width_state) + .then(|| PartialTableFlush { + threshold: bucket_threshold, + num_group_columns, + disabled: false, + sampled_groups: HashMap::new(), + sampled_flushes: VecDeque::new(), + num_flushes: 0, + flushed_groups: 0, + flush_count: MetricBuilder::new(&agg.metrics) + .counter("table_flush_count", partition), + hashes: vec![], + }); + let hash_table = AggregateHashTable::::new( agg, partition, @@ -279,8 +448,9 @@ impl PartialHashAggregateStream { reservation, reduction_factor, early_emit_count, + table_flush, skip_aggregation_probe, - group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + group_values_soft_limit, hash_table: Some(hash_table), }) } @@ -315,14 +485,19 @@ impl PartialHashAggregateStream { | HandleInputResult::SwitchToSkipAggregation => { break; } - HandleInputResult::OOM => { + HandleInputResult::OOM | HandleInputResult::TableFull => { let materialized_group_states = hash_table.take_state_batch()?.ok_or_else(|| { internal_datafusion_err!( "Partial hash aggregate ran out of memory with no aggregated groups" ) })?; - self.early_emit_count.add(1); + match (&last_state, self.table_flush.as_mut()) { + (HandleInputResult::TableFull, Some(table_flush)) => { + table_flush.record_flush(&materialized_group_states)? + } + _ => self.early_emit_count.add(1), + } timer.done(); self.emit_on_memory_pressure( materialized_group_states, @@ -410,7 +585,14 @@ impl PartialHashAggregateStream { // ---------------------------------------------- // Step 3: Skip partial aggregation optimization // ---------------------------------------------- - self.update_skip_aggregation_probe(input_rows, hash_table.building_group_count()); + let flushed_groups = self + .table_flush + .as_ref() + .map_or(0, |table_flush| table_flush.flushed_groups); + self.update_skip_aggregation_probe( + input_rows, + flushed_groups + hash_table.building_group_count(), + ); // True branch: a decision has been made to skip partial aggregation. if self.should_skip_aggregation() { @@ -422,10 +604,24 @@ impl PartialHashAggregateStream { // ------------------------------------------------- let resize_result = self.reservation.try_resize(hash_table.memory_size()); match resize_result { - Ok(()) => Ok(HandleInputResult::ProcessNext), - Err(DataFusionError::ResourcesExhausted(_)) => Ok(HandleInputResult::OOM), - Err(e) => Err(e), + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + return Ok(HandleInputResult::OOM); + } + Err(e) => return Err(e), + } + + // ----------------------------------------------------------- + // Step 5: Keep the table small while that is free (see + // `PartialTableFlush`) + // ----------------------------------------------------------- + let table_full = self.table_flush.as_ref().is_some_and(|table_flush| { + table_flush.should_flush(hash_table.building_group_count()) + }); + if table_full { + return Ok(HandleInputResult::TableFull); } + Ok(HandleInputResult::ProcessNext) } /// emit a materialized partial-state on memory pressure @@ -623,14 +819,40 @@ impl FinalHashAggregateStream { // Reserve memory for the initial hash table. that we hold for the lifetime of the stream. reservation.try_grow(hash_table.memory_size())?; + let group_values_soft_limit = agg.limit_options().map(|config| config.limit()); + + let bucket_threshold = context + .session_config() + .options() + .execution + .hash_aggregate_bucket_threshold; + // A soft limit stops reading input early, which bucketing cannot do. + let bucketing = (bucket_threshold > 0 + && group_values_soft_limit.is_none() + && BucketedAggregation::supports_state(&input_schema)) + .then(|| { + Arc::new(BucketedAggregation::new( + bucket_threshold, + agg.clone(), + partition, + batch_size, + Arc::clone(&input_schema), + Arc::clone(&schema), + spill_context + .as_ref() + .map(|context| context.spill_manager().clone()), + )) + }); + Ok(Self { schema, input, baseline_metrics, reservation, - group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + group_values_soft_limit, hash_table: Some(hash_table), spill_context, + bucketing, }) } @@ -652,10 +874,29 @@ impl FinalHashAggregateStream { let mut spill_context = self.spill_context.take(); - self.consume_input(&mut hash_table, &mut spill_context) + let buckets = self + .consume_input(&mut hash_table, &mut spill_context) .await?; self.close_input(); + if let (Some(buckets), Some(bucketing)) = (buckets, self.bucketing.clone()) { + // The table handed its groups over to the buckets, whose + // memory this stream's reservation already covers. + drop(hash_table); + let empty = self.reservation.new_empty(); + let bucket_reservation = std::mem::replace(&mut self.reservation, empty); + let mut emitter = emitter; + let mut output = bucketing.output_stream( + buckets, + bucket_reservation, + self.baseline_metrics.clone(), + ); + while let Some(batch) = output.next().await.transpose()? { + emitter.emit(batch).await; + } + return Ok(()); + } + match spill_context.filter(|s| s.has_spills()) { // - If spilled before, perform merging spill runs Some(spill_context) => { @@ -705,16 +946,35 @@ impl FinalHashAggregateStream { /// /// Spilling: The table cannot reserve enough memory. /// Move all current states into one fully group-key-sorted spill run. + /// + /// Bucketing: The table has reached `hash_aggregate_bucket_threshold` groups. + /// Move all current states and the rest of the input into hash + /// buckets, which are returned for [`BucketedAggregation::output_stream`]. async fn consume_input( &mut self, hash_table: &mut AggregateHashTable, spill_context: &mut Option>, - ) -> Result<()> { + ) -> Result> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let mut buckets: Option = None; + let mut compaction_table = None; + // Rows aggregated by `hash_table` + let mut table_rows = 0usize; while let Some(batch) = self.input.next().await.transpose()? { let _timer = elapsed_compute.timer(); + + if let (Some(buckets), Some(bucketing)) = + (buckets.as_mut(), self.bucketing.as_ref()) + { + buckets.route(&batch)?; + bucketing.compact(buckets, &mut compaction_table)?; + bucketing.reserve(&self.reservation, 0, buckets)?; + continue; + } + hash_table.aggregate_batch(&batch)?; + table_rows += batch.num_rows(); // Soft group limits are usually small and rarely coincide with // spilling. Once spilling has occurred, skip this optimization to @@ -726,6 +986,29 @@ impl FinalHashAggregateStream { break; } + // Once sorted runs exist the output comes from merging them, so + // bucketing only starts from a table that has never spilled. + if let Some(bucketing) = &self.bucketing + && !spilled + && starts_buckets( + hash_table.building_group_count(), + table_rows, + bucketing.threshold(), + ) + { + let kept = + hash_table.building_group_count() as f64 / table_rows.max(1) as f64; + let mut new_buckets = + bucketing.split(0, hash_table.take_state_batch()?, kept)?; + bucketing.reserve( + &self.reservation, + hash_table.memory_size(), + &mut new_buckets, + )?; + buckets = Some(new_buckets); + continue; + } + // Check memory reservation, and potentially spill. let resize_result = self.reservation @@ -778,7 +1061,7 @@ impl FinalHashAggregateStream { } } - Ok(()) + Ok(buckets) } /// Produce output from spills @@ -1192,6 +1475,351 @@ mod tests { Ok(()) } + /// Runs the final hash aggregation of `SELECT group_col, COUNT(value_col) + /// .. GROUP BY group_col` over partial state that holds every group + /// `num_partitions` times with a count of 1, as if that many partial + /// aggregations had fed it. The input is not charged to the memory pool, + /// so a memory limit only constrains the final aggregation. + /// + /// Returns the `(group, count)` rows sorted by group, and the final + /// aggregation's `bucket_splits` and `spill_count` metrics. + async fn run_final_hash_aggregate( + num_groups: usize, + num_partitions: usize, + bucket_threshold: usize, + memory_limit: Option, + ) -> Result<(Vec<(i32, i64)>, usize, usize)> { + use datafusion_common::ScalarValue; + + let batch_size = 1024; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + let mut runtime = RuntimeEnvBuilder::default(); + if let Some(memory_limit) = memory_limit { + runtime = runtime.with_memory_limit(memory_limit, 1.0); + } + let task_ctx = TaskContext::default().with_runtime(runtime.build_arc()?); + let session_config = task_ctx + .session_config() + .clone() + .set( + "datafusion.execution.batch_size", + &ScalarValue::UInt64(Some(batch_size as u64)), + ) + .set( + "datafusion.execution.hash_aggregate_bucket_threshold", + &ScalarValue::UInt64(Some(bucket_threshold as u64)), + ); + let task_ctx = Arc::new(task_ctx.with_session_config(session_config)); + + let group_by = PhysicalGroupBy::new_single(vec![( + col("group_col", &schema)?, + "group_col".to_string(), + )]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + // The partial aggregation is only built for its output schema + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let state_schema = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + empty, + Arc::clone(&schema), + )? + .schema(); + let mut state_batches = vec![]; + for _ in 0..num_partitions { + for start in (0..num_groups).step_by(batch_size) { + let end = (start + batch_size).min(num_groups); + let groups: Vec = (start as i32..end as i32).collect(); + let counts = vec![1i64; groups.len()]; + state_batches.push(RecordBatch::try_new( + Arc::clone(&state_schema), + vec![ + Arc::new(Int32Array::from(groups)), + Arc::new(Int64Array::from(counts)), + ], + )?); + } + } + let state_input = TestMemoryExec::try_new_exec( + &[state_batches], + Arc::clone(&state_schema), + None, + )?; + let final_agg = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggr_expr, + vec![None], + state_input, + Arc::clone(&schema), + )?); + + let batches = + crate::collect(Arc::clone(&final_agg) as Arc, task_ctx) + .await?; + let mut rows = vec![]; + for batch in &batches { + let groups = batch.column(0).as_primitive::(); + let counts = batch + .column(1) + .as_primitive::(); + rows.extend( + groups + .values() + .iter() + .copied() + .zip(counts.values().iter().copied()), + ); + } + rows.sort_unstable(); + + let metrics = final_agg.metrics().expect("final aggregate has metrics"); + let bucket_splits = metrics + .sum_by_name("bucket_splits") + .map(|value| value.as_usize()) + .unwrap_or(0); + BUCKET_COMPACTIONS.with(|compactions| { + compactions.set( + metrics + .sum_by_name("bucket_compactions") + .map(|value| value.as_usize()) + .unwrap_or(0), + ) + }); + Ok((rows, bucket_splits, metrics.spill_count().unwrap_or(0))) + } + + /// Runs the partial hash aggregation of `SELECT group_col, COUNT(value_col) + /// .. GROUP BY group_col` over `keys`, and returns the count of every group + /// summed over all the state rows emitted for it, the number of state rows, + /// and the `table_flush_count` metric. + async fn run_partial_hash_aggregate( + keys: Vec, + bucket_threshold: usize, + ) -> Result<(BTreeMap, usize, usize)> { + use datafusion_common::ScalarValue; + + let batch_size = 1024; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + let batches = keys + .chunks(batch_size) + .map(|keys| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys.to_vec())), + Arc::new(Int64Array::from(vec![1i64; keys.len()])), + ], + ) + }) + .collect::, _>>()?; + + let task_ctx = TaskContext::default(); + let session_config = task_ctx + .session_config() + .clone() + .set( + "datafusion.execution.batch_size", + &ScalarValue::UInt64(Some(batch_size as u64)), + ) + .set( + "datafusion.execution.hash_aggregate_bucket_threshold", + &ScalarValue::UInt64(Some(bucket_threshold as u64)), + ) + // keep the skip aggregation probe out of the way + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ); + let task_ctx = Arc::new(task_ctx.with_session_config(session_config)); + + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![( + col("group_col", &schema)?, + "group_col".to_string(), + )]), + aggr_expr, + vec![None], + input, + Arc::clone(&schema), + )?); + + let output = + crate::collect(Arc::clone(&partial) as Arc, task_ctx) + .await?; + let mut counts = BTreeMap::new(); + let mut state_rows = 0; + for batch in &output { + state_rows += batch.num_rows(); + let groups = batch.column(0).as_primitive::(); + let states = batch + .column(1) + .as_primitive::(); + for (group, count) in groups.values().iter().zip(states.values()) { + *counts.entry(*group).or_insert(0) += count; + } + } + let flushes = partial + .metrics() + .expect("partial aggregate has metrics") + .sum_by_name("table_flush_count") + .map(|value| value.as_usize()) + .unwrap_or(0); + Ok((counts, state_rows, flushes)) + } + + #[tokio::test] + async fn partial_hash_aggregate_flushes_groups_that_do_not_recur() -> Result<()> { + // Clustered keys: the 3 rows of a group are adjacent + let keys: Vec = (0..60_000).map(|row| row / 3).collect(); + + let (expected, state_rows, flushes) = + run_partial_hash_aggregate(keys.clone(), 0).await?; + assert_eq!((state_rows, flushes), (20_000, 0)); + assert!(expected.values().all(|&count| count == 3)); + + let (counts, state_rows, flushes) = + run_partial_hash_aggregate(keys, 2_000).await?; + assert_eq!(counts, expected); + assert!(flushes >= 9, "the table was flushed throughout: {flushes}"); + // A group is only emitted twice when a flush falls between its rows + assert!(state_rows <= 20_000 + flushes); + Ok(()) + } + + #[tokio::test] + async fn partial_hash_aggregate_stops_flushing_groups_that_recur() -> Result<()> { + // The same 30000 keys come around five times: every flushed group + // returns, but only 15 flushes later. + let keys: Vec = (0..150_000).map(|row| row % 30_000).collect(); + let (expected, _, _) = run_partial_hash_aggregate(keys.clone(), 0).await?; + + let (counts, state_rows, flushes) = + run_partial_hash_aggregate(keys, 2_000).await?; + assert_eq!(counts, expected); + assert!(expected.values().all(|&count| count == 5)); + // The first round is flushed; the flush that sees its groups again is the last + assert!((15..=17).contains(&flushes), "flushing stopped: {flushes}"); + assert!( + state_rows <= 64_000, + "later rounds are reduced: {state_rows}" + ); + Ok(()) + } + + thread_local! { + /// `bucket_compactions` metric of the last [`run_final_hash_aggregate`] + static BUCKET_COMPACTIONS: std::cell::Cell = const { std::cell::Cell::new(0) }; + } + + #[tokio::test] + async fn final_hash_aggregate_compacts_repeated_groups() -> Result<()> { + // Every group arrives 40 times, so buffering the input as is would + // hold 40 rows per group. + let (rows, splits, spills) = + run_final_hash_aggregate(20_000, 40, 1_000, None).await?; + assert_eq!(rows.len(), 20_000); + assert!(rows.iter().all(|&(_, count)| count == 40)); + assert!(splits >= 1); + assert_eq!(spills, 0); + let compactions = BUCKET_COMPACTIONS.with(|compactions| compactions.get()); + assert!(compactions > 0, "buckets were compacted"); + Ok(()) + } + + #[tokio::test] + async fn final_hash_aggregate_buckets_match_single_table() -> Result<()> { + let (expected, splits, _) = run_final_hash_aggregate(50_000, 3, 0, None).await?; + assert_eq!(splits, 0); + assert_eq!(expected.len(), 50_000); + assert!(expected.iter().all(|&(_, count)| count == 3)); + + // One split: 64 buckets of ~780 groups stay below the threshold + let (rows, splits, spills) = + run_final_hash_aggregate(50_000, 3, 10_000, None).await?; + assert_eq!(rows, expected); + assert_eq!(splits, 1); + assert_eq!(spills, 0); + + // A table that does not reach a quarter of the threshold is left alone + let (rows, splits, _) = + run_final_hash_aggregate(50_000, 3, 200_004, None).await?; + assert_eq!(rows, expected); + assert_eq!(splits, 0); + Ok(()) + } + + #[test] + fn starts_buckets_early_only_when_groups_do_not_repeat() { + // One row per group: a quarter of the threshold is enough + assert!(!starts_buckets(249, 249, 1_000)); + assert!(starts_buckets(250, 250, 1_000)); + // Groups repeat: only the full threshold + assert!(!starts_buckets(999, 10_000, 1_000)); + assert!(starts_buckets(1_000, 10_000, 1_000)); + // A threshold below the divisor + assert!(starts_buckets(1, 1, 2)); + } + + #[tokio::test] + async fn final_hash_aggregate_starts_buckets_early_for_unique_groups() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(50_000, 1, 0, None).await?; + + // Every group arrives once, so the table moves into buckets at a + // quarter of the threshold, which the 50,000 groups never reach. + // Buckets of ~780 groups are not split again. + let (rows, splits, _) = run_final_hash_aggregate(50_000, 1, 60_000, None).await?; + assert_eq!(rows, expected); + assert_eq!(splits, 1); + Ok(()) + } + + #[tokio::test] + async fn final_hash_aggregate_splits_large_buckets_again() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(50_000, 2, 0, None).await?; + + // Buckets of ~780 groups exceed the threshold and are split once more + let (rows, splits, _) = run_final_hash_aggregate(50_000, 2, 100, None).await?; + assert_eq!(rows, expected); + assert!(splits > 1, "buckets were split again, got {splits} splits"); + Ok(()) + } + + #[tokio::test] + async fn final_hash_aggregate_spills_buckets_under_memory_limit() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(200_000, 3, 0, None).await?; + + let (rows, splits, spills) = + run_final_hash_aggregate(200_000, 3, 10_000, Some(3 * 1024 * 1024)).await?; + assert_eq!(rows, expected); + assert!(splits >= 1); + assert!(spills > 0, "buckets were spilled"); + Ok(()) + } + /// Builds a partial hash aggregate stream over a single input batch of /// `num_groups` distinct groups, running under `memory_limit` bytes. /// diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6552bd95e06a5..d63e060dedea8 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -212,6 +212,8 @@ use topk::heap::is_supported_heap_type; mod aggregate_hash_table; mod aggregate_stream; +mod bucketed_aggregation; +mod final_buckets; pub mod group_values; mod grouped_hash_stream; mod grouped_topk_stream; diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 887be41855996..e778c74ac9efa 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -207,6 +207,13 @@ impl AggregateSpill { }) } + /// The spill manager of this stream, so that other parts of it — such as + /// the hash buckets of a final aggregation — spill through the same + /// manager and report to the same metrics. + pub(super) fn spill_manager(&self) -> &SpillManager { + &self.spill_manager + } + pub(super) fn has_spills(&self) -> bool { !self.spills.is_empty() } diff --git a/datafusion/sqllogictest/test_files/aggregate_bucketed.slt b/datafusion/sqllogictest/test_files/aggregate_bucketed.slt new file mode 100644 index 0000000000000..1be1a2ca243db --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_bucketed.slt @@ -0,0 +1,197 @@ +# 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. + +# Final hash aggregation splits its groups into hash buckets once the table +# reaches `hash_aggregate_bucket_threshold` groups. Every query runs with the +# option off and with a tiny threshold; the results must be identical. + +# The test context runs with 4 target partitions, so the aggregation is two-stage. + +statement ok +set datafusion.execution.batch_size = 1024; + +statement ok +create table t as select value as v from generate_series(1, 200000); + +# --- hash_aggregate_bucket_threshold = 0 + +statement ok +set datafusion.execution.hash_aggregate_bucket_threshold = 0; + +# integer key, several accumulators +query IIIII +select count(*), sum(c), sum(s), min(mn), max(mx) from (select v % 10007 as k, count(*) c, sum(v) s, min(v) mn, max(v) mx from t group by k); +---- +10007 200000 20000100000 1 200000 + +# avg keeps sum and count as state +query IR +select count(*), round(sum(a), 3) from (select v % 10007 as k, avg(v) a from t group by k); +---- +10007 1000705003.5 + +# string key +query II +select count(*), sum(c) from (select concat('key-', cast(v % 9973 as varchar)) as k, count(*) c from t group by k); +---- +9973 200000 + +# string view key +query II +select count(*), sum(c) from (select arrow_cast(concat('a-long-enough-string-view-key-', cast(v % 9973 as varchar)), 'Utf8View') as k, count(*) c from t group by k); +---- +9973 200000 + +# several key columns of mixed types +query III +select count(*), sum(c), sum(s) from (select v % 101 as k1, concat('k', cast(v % 103 as varchar)) as k2, v % 7 = 0 as k3, count(*) c, sum(v) s from t group by k1, k2, k3); +---- +20806 200000 20000100000 + +# nullable key +query II +select count(*), sum(c) from (select case when v % 11 = 0 then null else v % 5003 end as k, count(*) c from t group by k); +---- +5004 200000 + +# distinct has no accumulators +query I +select count(*) from (select distinct v % 20011 as k, v % 3 as j from t); +---- +60033 + +# accumulators with list state +query II +select count(*), sum(d) from (select v % 5003 as k, count(distinct v % 17) d from t group by k); +---- +5003 85051 + +# median +query IR +select count(*), sum(m) from (select v % 5003 as k, median(v) m from t group by k); +---- +5003 500302501.5 + +# array_agg +query II +select count(*), sum(array_length(a)) from (select v % 5003 as k, array_agg(v) a from t group by k); +---- +5003 200000 + +# every row its own group, largest groups first +query III +select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) order by k desc limit 3; +---- +200000 1 200000 +199999 1 199999 +199998 1 199998 + +# --- hash_aggregate_bucket_threshold = 100 + +statement ok +set datafusion.execution.hash_aggregate_bucket_threshold = 100; + +# Each of the four final partitions moved its groups into buckets once, and +# the partial aggregations flushed their tables +query TT +explain analyze select count(*) from (select v % 10007 as k, count(*) c from t group by k); +---- + +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v % Int64(10007)@0 as t.v % Int64(10007)], aggr=[], metrics=[bucket_splits=4,] + +08)--------------AggregateExec: mode=Partial, gby=[v@0 % 10007 as t.v % Int64(10007)], aggr=[], metrics=[table_flush_count=] + + +# integer key, several accumulators +query IIIII +select count(*), sum(c), sum(s), min(mn), max(mx) from (select v % 10007 as k, count(*) c, sum(v) s, min(v) mn, max(v) mx from t group by k); +---- +10007 200000 20000100000 1 200000 + +# avg keeps sum and count as state +query IR +select count(*), round(sum(a), 3) from (select v % 10007 as k, avg(v) a from t group by k); +---- +10007 1000705003.5 + +# The next three queries have a string key, so they keep a single table: only +# fixed-width state is bucketed. Their results must not change. + +# string key +query II +select count(*), sum(c) from (select concat('key-', cast(v % 9973 as varchar)) as k, count(*) c from t group by k); +---- +9973 200000 + +# string view key +query II +select count(*), sum(c) from (select arrow_cast(concat('a-long-enough-string-view-key-', cast(v % 9973 as varchar)), 'Utf8View') as k, count(*) c from t group by k); +---- +9973 200000 + +# several key columns of mixed types +query III +select count(*), sum(c), sum(s) from (select v % 101 as k1, concat('k', cast(v % 103 as varchar)) as k2, v % 7 = 0 as k3, count(*) c, sum(v) s from t group by k1, k2, k3); +---- +20806 200000 20000100000 + +# nullable key +query II +select count(*), sum(c) from (select case when v % 11 = 0 then null else v % 5003 end as k, count(*) c from t group by k); +---- +5004 200000 + +# distinct has no accumulators +query I +select count(*) from (select distinct v % 20011 as k, v % 3 as j from t); +---- +60033 + +# accumulators with list state +query II +select count(*), sum(d) from (select v % 5003 as k, count(distinct v % 17) d from t group by k); +---- +5003 85051 + +# median +query IR +select count(*), sum(m) from (select v % 5003 as k, median(v) m from t group by k); +---- +5003 500302501.5 + +# array_agg +query II +select count(*), sum(array_length(a)) from (select v % 5003 as k, array_agg(v) a from t group by k); +---- +5003 200000 + +# every row its own group, largest groups first +query III +select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) order by k desc limit 3; +---- +200000 1 200000 +199999 1 199999 +199998 1 199998 + +statement ok +reset datafusion.execution.hash_aggregate_bucket_threshold; + +statement ok +reset datafusion.execution.batch_size; + +statement ok +drop table t; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index f6633a165df18..89ce72171957a 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -222,6 +222,7 @@ datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true datafusion.execution.enforce_batch_size_in_joins false +datafusion.execution.hash_aggregate_bucket_threshold 0 datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false datafusion.execution.listing_table_factory_infer_partitions true @@ -384,6 +385,7 @@ datafusion.execution.enable_migration_aggregate true Whether aggregation uses th datafusion.execution.enable_nlj_coordinated_fallback true Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. +datafusion.execution.hash_aggregate_bucket_threshold 0 (experimental) Number of groups above which a hash aggregation stops growing a single hash table, so that its tables stay small enough to be cache friendly. A partial aggregation then emits the state of its table and starts over, as long as the emitted groups do not come back. A final aggregation splits the groups seen so far and all further input into hash buckets, which are aggregated one after another and can be spilled and released independently; it does so at a quarter of this number when its input holds about one row per group. Aggregations of millions of groups per partition run faster and with less memory. Only aggregations whose group keys and aggregate state are all fixed-width are affected: a string or other variable-length value costs about as much to move into its bucket as the smaller table saves. Moving rows has a cost of its own, so aggregations that end at a few times this number of groups can run a few percent slower, and input that repeats its groups can use more memory. Set to 0 to disable. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches datafusion.execution.listing_table_factory_infer_partitions true Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ba8ea52bb0850..b382af49ba249 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -141,6 +141,7 @@ The following configuration settings are available: | datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. | | datafusion.execution.skip_partial_aggregation_probe_ratio_threshold | 0.8 | Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input | | datafusion.execution.skip_partial_aggregation_probe_rows_threshold | 100000 | Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode | +| datafusion.execution.hash_aggregate_bucket_threshold | 0 | (experimental) Number of groups above which a hash aggregation stops growing a single hash table, so that its tables stay small enough to be cache friendly. A partial aggregation then emits the state of its table and starts over, as long as the emitted groups do not come back. A final aggregation splits the groups seen so far and all further input into hash buckets, which are aggregated one after another and can be spilled and released independently; it does so at a quarter of this number when its input holds about one row per group. Aggregations of millions of groups per partition run faster and with less memory. Only aggregations whose group keys and aggregate state are all fixed-width are affected: a string or other variable-length value costs about as much to move into its bucket as the smaller table saves. Moving rows has a cost of its own, so aggregations that end at a few times this number of groups can run a few percent slower, and input that repeats its groups can use more memory. Set to 0 to disable. | | datafusion.execution.use_row_number_estimates_to_optimize_partitioning | false | Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. | | datafusion.execution.enforce_batch_size_in_joins | false | Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. | | datafusion.execution.objectstore_writer_buffer_size | 10485760 | Size (bytes) of data buffer DataFusion uses when writing output files. This affects the size of the data chunks that are uploaded to remote object stores (e.g. AWS S3). If very large (>= 100 GiB) output files are being written, it may be necessary to increase this size to avoid errors from the remote end point. |