From 14346eba70d74cc5b90b5f7871a5bf3fedb78205 Mon Sep 17 00:00:00 2001 From: Rachel Int Date: Tue, 29 Sep 2026 02:50:51 +0800 Subject: [PATCH] feat: prototype frozen final aggregation buckets --- .../aggregate_frozen/RESULTS.md | 74 ++++ .../aggregate_frozen/aggregate_frozen.suite | 7 + .../aggregate_frozen/benchmarks/q01.benchmark | 11 + .../aggregate_frozen/benchmarks/q02.benchmark | 8 + .../aggregate_frozen/benchmarks/q03.benchmark | 8 + .../aggregate_frozen/benchmarks/q04.benchmark | 11 + .../aggregate_frozen/benchmarks/q05.benchmark | 8 + .../aggregate_frozen/benchmarks/q06.benchmark | 8 + datafusion/common/src/config.rs | 6 + .../aggregates/aggregate_hash_table/common.rs | 9 +- .../src/aggregates/bucketed_aggregation.rs | 8 + .../src/aggregates/hash_stream.rs | 403 +++++++++++++++++- .../test_files/aggregate_bucketed.slt | 29 ++ .../test_files/information_schema.slt | 2 + 14 files changed, 580 insertions(+), 12 deletions(-) create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/RESULTS.md create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/aggregate_frozen.suite create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q01.benchmark create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q02.benchmark create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q03.benchmark create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q04.benchmark create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q05.benchmark create mode 100644 benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q06.benchmark diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/RESULTS.md b/benchmarks/sql_benchmarks/aggregate_frozen/RESULTS.md new file mode 100644 index 0000000000000..3b1b83e82422c --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/RESULTS.md @@ -0,0 +1,74 @@ + + +# Frozen final aggregation experiment + +This experiment starts from DataFusion PR #25724 at `f9df89e2b1077deaae9be106682c9f1e4d0d61d2`. The benchmark data is generated by `range()`; no ClickBench `hits.parquet` or `hits_partitioned` dataset was available in the local workspace. These results do not establish an outcome for ClickBench Q11, Q12, Q14, Q18, Q33, or Q34. The ClickBench runner maps query number `N` to `queries/qN.sql`, starting at `q0.sql`. + +The experiment uses `hash_aggregate_bucket_threshold=4000`. Its optional `hash_aggregate_bucket_adaptive=true` retains the frozen table and probes subsequent rows against its original keys. Hits merge into that table, misses enter buckets, and the frozen state enters the buckets at EOF. This last merge preserves correctness after bypass sends later occurrences of frozen keys to buckets. Each partition samples 65,536 post-freeze rows and bypasses the probe when fewer than 25% hit. The sample is one-shot; it does not reactivate probing after a distribution change. + +The current DataFusion `GroupValues` interface has no lookup-only operation. This implementation therefore materializes a second encoded-key index with `RowConverter` and `HashSet`, then calls the original aggregation method for hit rows. Unlike the upstream [ClickHouse design](https://github.com/ClickHouse/ClickHouse/blob/master/src/Interpreters/AdaptiveAggregation.h) and its [producer state](https://github.com/ClickHouse/ClickHouse/blob/master/src/Interpreters/AdaptiveAggregationImpl.h), it does not reuse the table or routing hash for the first probe. The additional encoding and gathering work is part of the measured result. + +## Reproduce the SQL comparison + +Build with `cargo build -p datafusion-benchmarks --release --bin benchmark_runner`. From the repository root, run the following commands with the same host and no other benchmark workload running: + +```bash +DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_THRESHOLD=0 \ + target/release/benchmark_runner aggregate_frozen --iterations 10 --partitions 12 +DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_THRESHOLD=4000 \ + target/release/benchmark_runner aggregate_frozen --iterations 10 --partitions 12 +DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_THRESHOLD=4000 \ +DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_ADAPTIVE=true \ + target/release/benchmark_runner aggregate_frozen --iterations 10 --partitions 12 +``` + +The six SQL files each consume 2.4 million generated rows. Q01/Q04 have 80% hot keys and 20% unique keys; Q02/Q05 have five rows per key; Q03/Q06 are near unique. Q04-Q06 use approximately 80-character string prefixes. All six plans were checked to contain `AggregateExec: mode=FinalPartitioned`, and `EXPLAIN ANALYZE` showed 12 final bucket splits for each bucket-enabled query. + +### Optimized build, median of iterations 1-9, milliseconds + +| Query | Bucketing off | Original buckets | Frozen experiment | Frozen / original | +| --------------------------- | ------------: | ---------------: | ----------------: | ----------------: | +| Q01 hot integer | 6.9 | 10.6 | 11.9 | 1.12x | +| Q02 uniform integer | 12.8 | 13.8 | 16.7 | 1.21x | +| Q03 near unique integer | 16.2 | 14.5 | 17.1 | 1.18x | +| Q04 hot long string | 29.6 | 35.1 | 41.2 | 1.17x | +| Q05 uniform long string | 57.5 | 64.3 | 70.3 | 1.09x | +| Q06 near unique long string | 64.4 | 72.3 | 77.4 | 1.07x | + +The original buckets beat the single table only for Q03 in this suite. The frozen experiment was slower than the original buckets on all six queries. These queries are short, and their source and output costs differ from ClickBench; the comparison is evidence about this implementation on these distributions, not about the real ClickBench queries. + +### Final operator metrics from one separate `EXPLAIN ANALYZE` run + +| Query | Frozen probes | Hits | Bypassed rows | Staged rows original -> frozen | Routing CPU original -> frozen | EOF tail original -> frozen | +| ----- | ------------: | ----: | ------------: | -----------------------------: | -----------------------------: | --------------------------: | +| Q01 | 500.9k | 99.6k | 0 | 500.9k -> 401.3k | 31 -> 209 ms | 130 -> 102 ms | +| Q02 | 786.4k | 98.3k | 1.52M | 2.30M -> 2.20M | 151 -> 439 ms | 410 -> 382 ms | +| Q03 | 786.4k | 0 | 1.52M | 2.30M -> 2.30M | 156 -> 413 ms | 570 -> 576 ms | +| Q04 | 500.9k | 99.6k | 0 | 500.9k -> 401.3k | 132 -> 761 ms | 376 -> 326 ms | +| Q05 | 786.4k | 98.3k | 1.52M | 2.30M -> 2.20M | 528 -> 1620 ms | 903 -> 849 ms | +| Q06 | 786.4k | 0 | 1.52M | 2.30M -> 2.30M | 512 -> 1520 ms | 1550 -> 1450 ms | + +The final input is already reduced by the partial stage. Q02/Q05 therefore remove only about 4% of staged rows before bypass, and Q03/Q06 remove none. Q01/Q04 remove about 20%, but still spend more CPU in the extra lookup and row gathering. `bucket_routing_time` is summed CPU time across final partitions. `bucket_tail_time` sums each partition's wall time from input EOF through output completion; it is not query-wide wall latency. These metrics were collected with the development CLI build in a separate run and must not be added to or directly compared with the optimized benchmark milliseconds. + +The ignored Rust diagnostic `final_hash_aggregate_frozen_table_diagnostics` can print exact counters with `DATAFUSION_BUCKET_DIAGNOSTICS=1 cargo test -p datafusion-physical-plan final_hash_aggregate_frozen_table_diagnostics --lib -- --ignored --nocapture`. Its ordered repeated-key case sampled only 6,144 hits in 65,536 rows (9.375%) and bypassed before later repeats arrived. That demonstrates a possible sampling phase bias; it does not predict the ClickBench hit rate. + +## Assessment + +This version should remain an experiment. Before proposing it for merge, a direct lookup-only `GroupValues` path and reuse of the routing hash would remove two measured costs; a later measurement could then isolate the benefit of retaining frozen state. The one-shot sample also needs testing against changing key distributions. Real ClickBench Q11/Q12/Q14 and Q18/Q33/Q34, with Jay's dataset and threshold of 262,144, remain the decision point. diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/aggregate_frozen.suite b/benchmarks/sql_benchmarks/aggregate_frozen/aggregate_frozen.suite new file mode 100644 index 0000000000000..360dd437564e9 --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/aggregate_frozen.suite @@ -0,0 +1,7 @@ +description = "Synthetic final aggregation workloads for the frozen-table bucket experiment" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_THRESHOLD=4000 cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- aggregate_frozen --partitions 12 --iterations 3" +description = "Run the original bucket path; use threshold 0 for no bucketing or set DATAFUSION_EXECUTION_HASH_AGGREGATE_BUCKET_ADAPTIVE=true for the experiment." diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..e4333ce2cf33e --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q01.benchmark @@ -0,0 +1,11 @@ +name Q01_hot_integer +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +-- 80% of rows share 1,000 hot keys; remaining rows have unique keys. +SELECT k, COUNT(*) FROM ( + SELECT CASE WHEN value % 10 < 8 THEN value % 1000 ELSE value END AS k + FROM range(2400000) +) GROUP BY k; diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..6b0c7e6646241 --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q02.benchmark @@ -0,0 +1,8 @@ +name Q02_uniform_integer +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +-- Five rows per key, evenly spread over the input. +SELECT value % 480000 AS k, COUNT(*) FROM range(2400000) GROUP BY k; diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..e217f96b0f5d3 --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q03.benchmark @@ -0,0 +1,8 @@ +name Q03_near_unique_integer +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +-- Most keys occur once; the first 50,000 repeat at the end. +SELECT value % 2200000 AS k, COUNT(*) FROM range(2400000) GROUP BY k; diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..c5a71b07f6760 --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q04.benchmark @@ -0,0 +1,11 @@ +name Q04_hot_long_string +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +SELECT concat('xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-', CAST(k AS VARCHAR)), COUNT(*) +FROM ( + SELECT CASE WHEN value % 10 < 8 THEN value % 1000 ELSE value END AS k + FROM range(2400000) +) GROUP BY 1; diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..eba18a767b4e2 --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q05.benchmark @@ -0,0 +1,8 @@ +name Q05_uniform_long_string +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +SELECT concat('xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-', CAST(value % 480000 AS VARCHAR)), COUNT(*) +FROM range(2400000) GROUP BY 1; diff --git a/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..7762c947e257a --- /dev/null +++ b/benchmarks/sql_benchmarks/aggregate_frozen/benchmarks/q06.benchmark @@ -0,0 +1,8 @@ +name Q06_near_unique_long_string +group aggregate_frozen + +expect_plan mode=FinalPartitioned + +run +SELECT concat('xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-', CAST(value % 2200000 AS VARCHAR)), COUNT(*) +FROM range(2400000) GROUP BY 1; diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 2381d5603f73c..6a32446e381f1 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1141,6 +1141,12 @@ config_namespace! { /// Set to 0 to disable. pub hash_aggregate_bucket_threshold: usize, default = 0 + /// Experimental: retain the final aggregation's threshold-sized table + /// for keys already seen, staging only misses in hash buckets. After + /// 65,536 post-freeze rows, bypass its probe if fewer than 25% hit. + /// Requires `hash_aggregate_bucket_threshold` to be nonzero. + pub hash_aggregate_bucket_adaptive: bool, default = 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. 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 ebbf14a9cf73c..278000d4bed0d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -28,7 +28,7 @@ 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::{AggregateMetrics, EmitTo, GroupsAccumulator}; +use datafusion_expr::{AggregateMetrics, EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_physical_expr::GroupsAccumulatorAdapter; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use log::debug; @@ -416,6 +416,13 @@ impl AggregateHashTable { self.state.building().group_values.len() } + /// Returns the current group keys in group-id order without draining state. + pub(in crate::aggregates) fn building_group_keys(&mut self) -> Result> { + let values = &mut self.state.building_mut().group_values; + let group_count = values.len(); + values.values_preserving(GroupSelection::all(group_count)) + } + /// Takes every intermediate aggregate state and resets the table so it can /// continue accumulating raw input. /// diff --git a/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs b/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs index cb7b9a387bd35..623b253cf5024 100644 --- a/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs +++ b/datafusion/physical-plan/src/aggregates/bucketed_aggregation.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use arrow::compute::BatchCoalescer; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::instant::Instant; use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::{TryEmitter, async_try_stream}; @@ -66,6 +67,8 @@ pub(super) struct BucketedAggregation { bucket_splits: metrics::Count, /// Number of times the rows of a bucket were replaced by their aggregated state bucket_compactions: metrics::Count, + /// CPU time merging state after input has reached the buckets. + bucket_merge_time: metrics::Time, } /// What [`BucketedAggregation::output_stream`] works with while it runs. @@ -92,6 +95,8 @@ impl BucketedAggregation { MetricBuilder::new(&agg.metrics).counter("bucket_splits", partition); let bucket_compactions = MetricBuilder::new(&agg.metrics).counter("bucket_compactions", partition); + let bucket_merge_time = + MetricBuilder::new(&agg.metrics).subset_time("bucket_merge_time", partition); Self { threshold, agg, @@ -102,6 +107,7 @@ impl BucketedAggregation { spill_manager, bucket_splits, bucket_compactions, + bucket_merge_time, } } @@ -314,7 +320,9 @@ impl BucketedAggregation { continue; } + let started = Instant::now(); hash_table.aggregate_batch(&batch)?; + self.bucket_merge_time.add_elapsed(started); table_rows += batch.num_rows(); let can_split = next_level < MAX_BUCKET_LEVELS; diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index df86f5192eefa..5fc525187c71d 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -20,15 +20,20 @@ //! See comments in [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] //! for details. -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::mem::size_of; use std::sync::Arc; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; +use arrow::array::{ArrayRef, PrimitiveArray}; +use arrow::compute::take_arrays; +use arrow::datatypes::{DataType, SchemaRef, UInt32Type}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use arrow::row::{RowConverter, SortField}; use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::instant::Instant; use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, + internal_err, }; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; @@ -208,6 +213,159 @@ pub(crate) struct FinalHashAggregateStream { spill_context: Option>, /// `None` unless `hash_aggregate_bucket_threshold` is set and applies. bucketing: Option>, + /// Enables the frozen-table experiment for this execution. + adaptive_buckets: bool, + adaptive_metrics: AdaptiveBucketMetrics, +} + +struct AdaptiveBucketMetrics { + probes: metrics::Count, + hits: metrics::Count, + misses: metrics::Count, + bypass_rows: metrics::Count, + staged_rows: metrics::Count, + staged_bytes: metrics::Count, + routing_time: metrics::Time, + aggregation_time: metrics::Time, + tail_time: metrics::Time, +} + +impl AdaptiveBucketMetrics { + fn new(agg: &AggregateExec, partition: usize) -> Self { + let counter = |name| MetricBuilder::new(&agg.metrics).counter(name, partition); + let timer = |name| MetricBuilder::new(&agg.metrics).subset_time(name, partition); + Self { + probes: counter("frozen_probes"), + hits: counter("frozen_hits"), + misses: counter("frozen_misses"), + bypass_rows: counter("frozen_bypass_rows"), + staged_rows: counter("bucket_staged_rows"), + staged_bytes: counter("bucket_staged_bytes"), + routing_time: timer("bucket_routing_time"), + aggregation_time: timer("bucket_aggregation_time"), + tail_time: timer("bucket_tail_time"), + } + } +} + +/// Keys present when a final aggregation first enters bucketing. +/// The table retains their aggregate state; this index only classifies rows. +struct FrozenGroups { + converter: RowConverter, + keys: HashSet>, + group_columns: usize, + key_bytes: usize, + sampled_rows: usize, + sampled_hits: usize, + bypass: bool, +} + +impl FrozenGroups { + fn supports(keys: &[ArrayRef]) -> bool { + !keys.is_empty() + && keys.iter().all(|array| { + matches!( + array.data_type(), + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + | DataType::Date32 + | DataType::Date64 + ) + }) + } + + fn new(keys: &[ArrayRef]) -> Result { + let fields = keys + .iter() + .map(|array| SortField::new(array.data_type().clone())) + .collect::>(); + let converter = RowConverter::new(fields)?; + let rows = converter.convert_columns(keys)?; + let mut stored = HashSet::with_capacity(rows.num_rows()); + let mut key_bytes = 0; + for row in rows.iter() { + key_bytes += row.data().len(); + stored.insert(row.data().to_vec()); + } + Ok(Self { + converter, + keys: stored, + group_columns: keys.len(), + key_bytes, + sampled_rows: 0, + sampled_hits: 0, + bypass: false, + }) + } + + fn memory_size(&self) -> usize { + self.key_bytes + self.keys.capacity() * (size_of::>() + 8) + } + + fn partition(&mut self, batch: &RecordBatch) -> Result<(Vec, Vec, usize)> { + let mut hits = Vec::new(); + let mut misses = Vec::new(); + if self.bypass { + misses.extend(0..batch.num_rows() as u32); + return Ok((hits, misses, 0)); + } + let rows = self + .converter + .convert_columns(&batch.columns()[..self.group_columns])?; + let mut probes = 0; + for (index, row) in rows.iter().enumerate() { + let hit = self.keys.contains(row.data()); + probes += 1; + if hit { + hits.push(index as u32); + } else { + misses.push(index as u32); + } + if self.sampled_rows < 65_536 { + self.sampled_rows += 1; + self.sampled_hits += usize::from(hit); + if self.sampled_rows == 65_536 { + self.bypass = self.sampled_hits * 4 < self.sampled_rows; + } + } + if self.bypass { + misses.extend((index + 1..batch.num_rows()).map(|i| i as u32)); + break; + } + } + Ok((hits, misses, probes)) + } +} + +fn take_state_rows( + batch: &RecordBatch, + indices: Vec, +) -> Result> { + if indices.is_empty() { + return Ok(None); + } + let row_count = indices.len(); + let indices: PrimitiveArray = indices.into(); + let columns = take_arrays(batch.columns(), &indices, None)?; + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + Ok(Some(RecordBatch::try_new_with_options( + batch.schema(), + columns, + &options, + )?)) } #[derive(PartialEq)] @@ -825,6 +983,11 @@ impl FinalHashAggregateStream { .execution .hash_aggregate_bucket_threshold; // A soft limit stops reading input early, which bucketing cannot do. + let adaptive_buckets = context + .session_config() + .options() + .execution + .hash_aggregate_bucket_adaptive; let bucketing = (bucket_threshold > 0 && group_values_soft_limit.is_none() && BucketedAggregation::supports_state( @@ -854,6 +1017,8 @@ impl FinalHashAggregateStream { hash_table: Some(hash_table), spill_context, bucketing, + adaptive_buckets, + adaptive_metrics: AdaptiveBucketMetrics::new(agg, partition), }) } @@ -878,6 +1043,7 @@ impl FinalHashAggregateStream { let buckets = self .consume_input(&mut hash_table, &mut spill_context) .await?; + let input_eof = Instant::now(); self.close_input(); if let (Some(buckets), Some(bucketing)) = (buckets, self.bucketing.clone()) { @@ -895,6 +1061,7 @@ impl FinalHashAggregateStream { while let Some(batch) = output.next().await.transpose()? { emitter.emit(batch).await; } + self.adaptive_metrics.tail_time.add_elapsed(input_eof); return Ok(()); } @@ -958,6 +1125,7 @@ impl FinalHashAggregateStream { ) -> Result> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let mut buckets: Option = None; + let mut frozen: Option = None; let mut compaction_table = None; // Rows aggregated by `hash_table` let mut table_rows = 0usize; @@ -968,9 +1136,70 @@ impl FinalHashAggregateStream { if let (Some(buckets), Some(bucketing)) = (buckets.as_mut(), self.bucketing.as_ref()) { - buckets.route(&batch)?; + if let Some(frozen) = frozen.as_mut() { + let started = Instant::now(); + let (hits, misses, probes) = frozen.partition(&batch)?; + self.adaptive_metrics.routing_time.add_elapsed(started); + self.adaptive_metrics.probes.add(probes); + self.adaptive_metrics.hits.add(hits.len()); + self.adaptive_metrics.misses.add(probes - hits.len()); + self.adaptive_metrics + .bypass_rows + .add(batch.num_rows() - probes); + if let Some(hit_batch) = take_state_rows(&batch, hits)? { + let started = Instant::now(); + let group_count = hash_table.building_group_count(); + hash_table.aggregate_batch(&hit_batch)?; + if hash_table.building_group_count() != group_count { + return internal_err!( + "Frozen aggregate table inserted a new group while merging hits" + ); + } + self.adaptive_metrics.aggregation_time.add_elapsed(started); + } + if let Some(miss_batch) = take_state_rows(&batch, misses)? { + self.adaptive_metrics.staged_rows.add(miss_batch.num_rows()); + self.adaptive_metrics + .staged_bytes + .add(miss_batch.get_array_memory_size()); + let started = Instant::now(); + buckets.route(&miss_batch)?; + self.adaptive_metrics.routing_time.add_elapsed(started); + } + } else { + self.adaptive_metrics.staged_rows.add(batch.num_rows()); + self.adaptive_metrics + .staged_bytes + .add(batch.get_array_memory_size()); + let started = Instant::now(); + buckets.route(&batch)?; + self.adaptive_metrics.routing_time.add_elapsed(started); + } + let started = Instant::now(); bucketing.compact(buckets, &mut compaction_table)?; - bucketing.reserve(&self.reservation, 0, buckets)?; + self.adaptive_metrics.aggregation_time.add_elapsed(started); + let retained_bytes = frozen + .as_ref() + .map_or(0, |frozen| hash_table.memory_size() + frozen.memory_size()); + if let Err(error) = + bucketing.reserve(&self.reservation, retained_bytes, buckets) + { + match error { + DataFusionError::ResourcesExhausted(_) if frozen.is_some() => { + if let Some(state) = hash_table.take_state_batch()? { + buckets.route(&state)?; + } + drop(frozen.take()); + bucketing.compact(buckets, &mut compaction_table)?; + bucketing.reserve( + &self.reservation, + hash_table.memory_size(), + buckets, + )?; + } + error => return Err(error), + } + } continue; } @@ -999,13 +1228,41 @@ impl FinalHashAggregateStream { { 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( + let candidate = if self.adaptive_buckets { + let keys = hash_table.building_group_keys()?; + FrozenGroups::supports(&keys) + .then(|| FrozenGroups::new(&keys)) + .transpose()? + } else { + None + }; + let mut new_buckets = if candidate.is_some() { + bucketing.split(0, None, kept)? + } else { + bucketing.split(0, hash_table.take_state_batch()?, kept)? + }; + let retained_bytes = hash_table.memory_size() + + candidate.as_ref().map_or(0, FrozenGroups::memory_size); + match bucketing.reserve( &self.reservation, - hash_table.memory_size(), + retained_bytes, &mut new_buckets, - )?; + ) { + Ok(()) => frozen = candidate, + Err(DataFusionError::ResourcesExhausted(_)) + if candidate.is_some() => + { + if let Some(state) = hash_table.take_state_batch()? { + new_buckets.route(&state)?; + } + bucketing.reserve( + &self.reservation, + hash_table.memory_size(), + &mut new_buckets, + )?; + } + Err(error) => return Err(error), + } buckets = Some(new_buckets); continue; } @@ -1062,6 +1319,20 @@ impl FinalHashAggregateStream { } } + if frozen.is_some() { + let buckets = buckets.as_mut().expect("frozen table has buckets"); + let bucketing = self.bucketing.as_ref().expect("frozen table has bucketing"); + if let Some(state) = hash_table.take_state_batch()? { + let started = Instant::now(); + buckets.route(&state)?; + self.adaptive_metrics.routing_time.add_elapsed(started); + let started = Instant::now(); + bucketing.compact(buckets, &mut compaction_table)?; + self.adaptive_metrics.aggregation_time.add_elapsed(started); + } + drop(frozen.take()); + bucketing.reserve(&self.reservation, hash_table.memory_size(), buckets)?; + } Ok(buckets) } @@ -1490,6 +1761,24 @@ mod tests { bucket_threshold: usize, memory_limit: Option, ) -> Result<(Vec<(i32, i64)>, usize, usize)> { + let (rows, splits, spills, _, _) = run_final_hash_aggregate_adaptive( + num_groups, + num_partitions, + bucket_threshold, + memory_limit, + false, + ) + .await?; + Ok((rows, splits, spills)) + } + + async fn run_final_hash_aggregate_adaptive( + num_groups: usize, + num_partitions: usize, + bucket_threshold: usize, + memory_limit: Option, + adaptive: bool, + ) -> Result<(Vec<(i32, i64)>, usize, usize, usize, usize)> { use datafusion_common::ScalarValue; let batch_size = 1024; @@ -1513,6 +1802,10 @@ mod tests { .set( "datafusion.execution.hash_aggregate_bucket_threshold", &ScalarValue::UInt64(Some(bucket_threshold as u64)), + ) + .set( + "datafusion.execution.hash_aggregate_bucket_adaptive", + &ScalarValue::Boolean(Some(adaptive)), ); let task_ctx = Arc::new(task_ctx.with_session_config(session_config)); @@ -1599,7 +1892,34 @@ mod tests { .unwrap_or(0), ) }); - Ok((rows, bucket_splits, metrics.spill_count().unwrap_or(0))) + let metric_count = |name| { + metrics + .sum_by_name(name) + .map(|value| value.as_usize()) + .unwrap_or(0) + }; + if std::env::var_os("DATAFUSION_BUCKET_DIAGNOSTICS").is_some() { + println!( + "adaptive={adaptive} groups={num_groups} repeats={num_partitions} splits={bucket_splits} probes={} hits={} misses={} bypass={} staged_rows={} staged_bytes={} routing_ns={} aggregation_ns={} bucket_merge_ns={} tail_ns={}", + metric_count("frozen_probes"), + metric_count("frozen_hits"), + metric_count("frozen_misses"), + metric_count("frozen_bypass_rows"), + metric_count("bucket_staged_rows"), + metric_count("bucket_staged_bytes"), + metric_count("bucket_routing_time"), + metric_count("bucket_aggregation_time"), + metric_count("bucket_merge_time"), + metric_count("bucket_tail_time"), + ); + } + Ok(( + rows, + bucket_splits, + metrics.spill_count().unwrap_or(0), + metric_count("frozen_hits"), + metric_count("frozen_bypass_rows"), + )) } /// Runs the partial hash aggregation of `SELECT group_col, COUNT(value_col) @@ -1773,6 +2093,50 @@ mod tests { Ok(()) } + #[tokio::test] + async fn final_hash_aggregate_frozen_table_merges_bypassed_keys() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(100_000, 2, 0, None).await?; + let (rows, splits, spills, hits, bypass_rows) = + run_final_hash_aggregate_adaptive(100_000, 2, 4_000, None, true).await?; + assert_eq!(rows, expected); + assert!(splits > 0); + assert_eq!(spills, 0); + assert!(bypass_rows > 0, "low hit rate bypasses the frozen table"); + assert!(hits < 65_536, "the initial sample should have few hits"); + Ok(()) + } + + #[tokio::test] + async fn final_hash_aggregate_frozen_table_merges_hits_and_misses() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(10_000, 20, 0, None).await?; + let (rows, splits, _, hits, _) = + run_final_hash_aggregate_adaptive(10_000, 20, 4_000, None, true).await?; + assert_eq!(rows, expected); + assert!(splits > 0); + assert!(hits > 0, "frozen keys were merged during input"); + Ok(()) + } + + #[tokio::test] + #[ignore = "diagnostic comparison; run with DATAFUSION_BUCKET_DIAGNOSTICS=1 --nocapture"] + async fn final_hash_aggregate_frozen_table_diagnostics() -> Result<()> { + for (groups, repeats) in [(10_000, 20), (100_000, 6), (200_000, 1)] { + for (threshold, adaptive) in [(0, false), (4_000, false), (4_000, true)] { + let started = Instant::now(); + let (rows, _, _, _, _) = run_final_hash_aggregate_adaptive( + groups, repeats, threshold, None, adaptive, + ) + .await?; + println!( + "diagnostic groups={groups} repeats={repeats} threshold={threshold} adaptive={adaptive} total_ms={:.2} output_groups={}", + started.elapsed().as_secs_f64() * 1000.0, + rows.len(), + ); + } + } + Ok(()) + } + #[test] fn starts_buckets_early_only_when_groups_do_not_repeat() { // One row per group: a quarter of the threshold is enough @@ -1821,6 +2185,23 @@ mod tests { Ok(()) } + #[tokio::test] + async fn final_hash_aggregate_frozen_table_spills_buckets() -> Result<()> { + let (expected, _, _) = run_final_hash_aggregate(200_000, 3, 0, None).await?; + let (rows, splits, spills, _, _) = run_final_hash_aggregate_adaptive( + 200_000, + 3, + 10_000, + Some(3 * 1024 * 1024), + true, + ) + .await?; + assert_eq!(rows, expected); + assert!(splits > 0); + assert!(spills > 0); + 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/sqllogictest/test_files/aggregate_bucketed.slt b/datafusion/sqllogictest/test_files/aggregate_bucketed.slt index e555c44de52d1..6264b900efc6a 100644 --- a/datafusion/sqllogictest/test_files/aggregate_bucketed.slt +++ b/datafusion/sqllogictest/test_files/aggregate_bucketed.slt @@ -184,6 +184,35 @@ select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) orde 199999 1 199999 199998 1 199998 +statement ok +set datafusion.execution.hash_aggregate_bucket_adaptive = true; + +# frozen integer groups merge aggregate state +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 + +# long string view keys exercise the frozen lookup +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 + +# nulls and mixed key columns keep their grouping semantics +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 + +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 + +statement ok +reset datafusion.execution.hash_aggregate_bucket_adaptive; + statement ok reset datafusion.execution.hash_aggregate_bucket_threshold; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 2c9e3182fbcc7..36c7b8facb82c 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_adaptive false datafusion.execution.hash_aggregate_bucket_threshold 0 datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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_adaptive false Experimental: retain the final aggregation's threshold-sized table for keys already seen, staging only misses in hash buckets. After 65,536 post-freeze rows, bypass its probe if fewer than 25% hit. Requires `hash_aggregate_bucket_threshold` to be nonzero. 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. Moving rows into buckets has a cost of its own, so aggregations that end at a few times this number of groups, string keys in particular, 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