comphead commented on code in PR #26562: URL: https://github.com/apache/datafusion/pull/26562#discussion_r4364432639
########## datafusion/physical-plan/src/aggregates/memory_call_sequence_tests.rs: ########## @@ -0,0 +1,886 @@ +// 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. + +//! Snapshot tests that lock down the exact sequence of memory pool calls made +//! by the grouped aggregate streams, for the memory accounting paths that the +//! other tests do not reach: spill headroom and replay, refusals with +//! temporary files disabled, refused output materialization, and holding +//! completed groups in the clustered streams. The early emit of the +//! unclustered Partial and PartialReduce streams is covered by their own tests. +//! +//! Each test runs one single-partition `AggregateExec` under a +//! [`RecordingPool`] and snapshots every pool call of the aggregate consumers +//! (the stream and the merge / replay siblings that share its consumer), then +//! the query outcome. Limits were chosen by experiment so that each test +//! really takes its path, which the test also asserts through metrics or the +//! output batch sizes. Byte counts follow the current accounting of the hash +//! tables and arrays, so a change there shows up as a snapshot diff. Some of +//! them also depend on the target (see [`assert_snapshot_per_group_width`]). + +use std::fmt; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility, +}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use hashbrown::HashTable; +use insta::assert_snapshot; + +use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy}; +use crate::ExecutionPlan; +use crate::common::collect; +use crate::metrics::MetricsSet; +use crate::test::TestMemoryExec; + +/// A [`GreedyMemoryPool`] that records every call made on behalf of an +/// aggregate consumer, as `"<consumer name> <op> <bytes>"`. +#[derive(Debug)] +struct RecordingPool { + inner: GreedyMemoryPool, + calls: Mutex<Vec<(String, String)>>, +} + +impl RecordingPool { + fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + calls: Mutex::new(vec![]), + } + } + + fn record(&self, consumer: &MemoryConsumer, call: String) { + let name = consumer.name(); + if name.contains("Aggregate") { + let line = format!("{name} {call}"); + self.calls.lock().unwrap().push((name.to_string(), line)); + } + } + + /// The recorded calls grouped by consumer, in first-seen consumer order. + fn calls(&self) -> String { Review Comment: Each migrated stream registers a single consumer, and the spill merge reuses it through `reservation.new_empty()`. So every snapshot has one name, and the regrouping in `calls()` and the `(String, String)` pairs never change the output. Recording into a `Mutex<Vec<String>>` in call order would be simpler. I would also drop the `contains("Aggregate")` filter. As far as I can tell, no other consumer takes part in these runs. Without the filter, a consumer that #26563 adds under a different name shows up in the snapshot instead of being filtered out. ########## datafusion/physical-plan/src/aggregates/memory_call_sequence_tests.rs: ########## @@ -0,0 +1,886 @@ +// 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. + +//! Snapshot tests that lock down the exact sequence of memory pool calls made +//! by the grouped aggregate streams, for the memory accounting paths that the +//! other tests do not reach: spill headroom and replay, refusals with +//! temporary files disabled, refused output materialization, and holding +//! completed groups in the clustered streams. The early emit of the +//! unclustered Partial and PartialReduce streams is covered by their own tests. +//! +//! Each test runs one single-partition `AggregateExec` under a +//! [`RecordingPool`] and snapshots every pool call of the aggregate consumers +//! (the stream and the merge / replay siblings that share its consumer), then +//! the query outcome. Limits were chosen by experiment so that each test +//! really takes its path, which the test also asserts through metrics or the +//! output batch sizes. Byte counts follow the current accounting of the hash +//! tables and arrays, so a change there shows up as a snapshot diff. Some of +//! them also depend on the target (see [`assert_snapshot_per_group_width`]). + +use std::fmt; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility, +}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use hashbrown::HashTable; +use insta::assert_snapshot; + +use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy}; +use crate::ExecutionPlan; +use crate::common::collect; +use crate::metrics::MetricsSet; +use crate::test::TestMemoryExec; + +/// A [`GreedyMemoryPool`] that records every call made on behalf of an +/// aggregate consumer, as `"<consumer name> <op> <bytes>"`. +#[derive(Debug)] +struct RecordingPool { + inner: GreedyMemoryPool, + calls: Mutex<Vec<(String, String)>>, +} + +impl RecordingPool { + fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + calls: Mutex::new(vec![]), + } + } + + fn record(&self, consumer: &MemoryConsumer, call: String) { + let name = consumer.name(); + if name.contains("Aggregate") { + let line = format!("{name} {call}"); + self.calls.lock().unwrap().push((name.to_string(), line)); + } + } + + /// The recorded calls grouped by consumer, in first-seen consumer order. + fn calls(&self) -> String { + let calls = self.calls.lock().unwrap(); + let mut names: Vec<&str> = vec![]; + for (name, _) in calls.iter() { + if !names.contains(&name.as_str()) { + names.push(name); + } + } + let mut out = String::new(); + for name in names { + for (_, line) in calls.iter().filter(|(n, _)| n == name) { + out.push_str(line); + out.push('\n'); + } + } + out + } +} + +impl fmt::Display for RecordingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "RecordingPool({})", self.inner) + } +} + +impl MemoryPool for RecordingPool { + fn name(&self) -> &str { + "RecordingPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + self.record( + consumer, + format!("register can_spill={}", consumer.can_spill()), + ); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + self.record(consumer, "unregister".to_string()); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(reservation.consumer(), format!("grow {additional}")); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.record(reservation.consumer(), format!("shrink {shrink}")); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + let result = self.inner.try_grow(reservation, additional); + let suffix = if result.is_ok() { "" } else { " -> refused" }; + self.record( + reservation.consumer(), + format!("try_grow {additional}{suffix}"), + ); + result + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +/// The single aggregate of a [`Case`]. +#[derive(Clone, Copy)] +enum Agg { + /// `SUM(v)`, whose output is no larger than its accounted state. + Sum, + /// [`PaddedOutputUdaf`], whose output is far larger than its accounted state. + PaddedOutput, +} + +/// One aggregate run: `agg` grouped by `group_by` over raw `(a, b, v)` rows. +struct Case { + mode: AggregateMode, + agg: Agg, + /// Group columns, a prefix of `["a", "b"]`. + group_by: &'static [&'static str], + /// Declare the input sorted by `a`, enabling group clustering. + sorted_by_a: bool, + /// The `(a, b)` keys of each input batch; every row has `v = 1`. + batches: Vec<Vec<(i64, i64)>>, + batch_size: usize, + memory_limit: usize, + spill_enabled: bool, +} + +/// What a [`Case`] did. +struct Run { + /// Recorded pool calls followed by the query outcome. + snapshot: String, + clustering: GroupClusteringMode, + metrics: MetricsSet, + output_batch_rows: Vec<usize>, +} + +impl Run { + fn metric(&self, name: &str) -> usize { + self.metrics + .sum_by_name(name) + .map(|m| m.as_usize()) + .unwrap_or(0) + } + + fn spill_count(&self) -> usize { + self.metrics.spill_count().unwrap_or(0) + } + + fn max_output_rows(&self) -> usize { + self.output_batch_rows.iter().copied().max().unwrap_or(0) + } +} + +fn raw_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])) +} + +/// Distinct keys `(start..start + n, 0)`. +fn distinct_a(start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|a| (a, 0)).collect() +} + +/// Keys `(a, start..start + n)` with one `a` value. +fn one_a(a: i64, start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|b| (a, b)).collect() +} + +impl Case { + async fn run(self) -> Result<Run> { + let raw = raw_schema(); + let group_by = PhysicalGroupBy::new_single( + self.group_by + .iter() + .map(|name| Ok((col(name, &raw)?, name.to_string()))) + .collect::<Result<_>>()?, + ); + let udaf = match self.agg { + Agg::Sum => sum_udaf(), + Agg::PaddedOutput => Arc::new(AggregateUDF::from(PaddedOutputUdaf::new())), + }; + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("v", &raw)?]) + .schema(Arc::clone(&raw)) + .alias("s") + .build()?, + )]; + + // Stages after the first read partial states: the group columns followed + // by the `SUM` state, which for Int64 is the running sum itself. + let reads_partial_state = matches!( + self.mode, + AggregateMode::Final + | AggregateMode::FinalPartitioned + | AggregateMode::PartialReduce + ); + let input_schema = if reads_partial_state { + Arc::new(super::create_schema( + &raw, + &group_by, + &aggregates, + AggregateMode::Partial, + )?) + } else { + Arc::clone(&raw) + }; + let batches = self + .batches + .iter() + .map(|keys| { + let a: ArrayRef = + Arc::new(keys.iter().map(|k| k.0).collect::<Int64Array>()); + let b: ArrayRef = + Arc::new(keys.iter().map(|k| k.1).collect::<Int64Array>()); + let v: ArrayRef = Arc::new(Int64Array::from(vec![1; keys.len()])); + let columns = if reads_partial_state { + let mut columns = vec![a, b]; + columns.truncate(self.group_by.len()); + columns.push(match self.agg { + Agg::Sum => v, + Agg::PaddedOutput => { + Arc::new(StringArray::from(vec!["x"; keys.len()])) + } + }); + columns + } else { + vec![a, b, v] + }; + Ok(RecordBatch::try_new(Arc::clone(&input_schema), columns)?) + }) + .collect::<Result<Vec<_>>>()?; + + let mut input = + TestMemoryExec::try_new(&[batches], Arc::clone(&input_schema), None)?; + if self.sorted_by_a { + input = input.try_with_sort_information(vec![ + LexOrdering::new([PhysicalSortExpr::new_default(col( + "a", + &input_schema, + )?)]) + .unwrap(), + ])?; + } + let aggregate = AggregateExec::try_new( + self.mode, + group_by, + aggregates, + vec![None], + Arc::new(input), + raw, + )?; + + let pool = Arc::new(RecordingPool::new(self.memory_limit)); + let mut runtime = + RuntimeEnvBuilder::new().with_memory_pool(Arc::clone(&pool) as _); + if !self.spill_enabled { + runtime = runtime.with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ); + } + let context = Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(self.batch_size) + .set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + ), + ) + .with_runtime(runtime.build_arc()?), + ); + + let result = match aggregate.execute(0, Arc::clone(&context)) { + Ok(stream) => collect(stream).await, + Err(e) => Err(e), + }; + let output_batch_rows: Vec<usize> = result + .iter() + .flatten() + .map(|batch| batch.num_rows()) + .collect(); + let outcome = match &result { + Ok(batches) => format!( + "Ok({} rows in {} batches)", + output_batch_rows.iter().sum::<usize>(), + batches.len() + ), + Err(e) => format!("Err({})", e.strip_backtrace()), + }; + drop(result); + drop(context); + assert_eq!(pool.reserved(), 0, "every reservation is released"); + + Ok(Run { + snapshot: format!("{}=> {outcome}", pool.calls()), + clustering: aggregate.group_clustering_mode().clone(), + metrics: aggregate.metrics().unwrap(), + output_batch_rows, + }) + } +} + +/// Bytes of every [`PaddedOutputUdaf`] output and state value. +const PAD: usize = 256; + +/// An aggregate that accounts almost no memory per group but materializes a +/// `PAD`-byte string per group as its state and its output, so materializing +/// the output can need far more memory than the table held. +#[derive(Debug, PartialEq, Eq, Hash)] +struct PaddedOutputUdaf { + signature: Signature, +} + +impl PaddedOutputUdaf { + fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int64], Volatility::Immutable), + } + } +} + +impl AggregateUDFImpl for PaddedOutputUdaf { + fn name(&self) -> &str { + "padded_output" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Utf8) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format!("{}[pad]", args.name), + DataType::Utf8, + true, + ))]) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(PaddedOutputAccumulator)) + } +} + +#[derive(Debug)] +struct PaddedOutputAccumulator; + +impl Accumulator for PaddedOutputAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + Ok(ScalarValue::Utf8(Some("x".repeat(PAD)))) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![ScalarValue::Utf8(Some("x".repeat(PAD)))]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } +} + +fn case( + mode: AggregateMode, + group_by: &'static [&'static str], + sorted_by_a: bool, + batches: Vec<Vec<(i64, i64)>>, +) -> Case { + Case { + mode, + agg: Agg::Sum, + group_by, + sorted_by_a, + batches, + batch_size: 64, + memory_limit: usize::MAX, + spill_enabled: true, + } +} + +/// Asserts `snapshot` against `snapshots/<module>__<name>@ctrl<width>.snap`. +/// +/// A group hash table's [`HashTable::allocation_size`] includes one control +/// byte per bucket plus hashbrown's SIMD group width, which is 16 with SSE2 +/// (x86_64) and 8 on aarch64. Whether a pool call includes the table depends +/// on the stream's state, so the difference cannot be normalized away, and +/// snapshots whose byte counts differ are kept once per width. +fn assert_snapshot_per_group_width(name: &str, snapshot: &str) { + // Capacity 1 allocates 4 buckets of 16 bytes, then 4 + width control bytes. + let table = HashTable::<[u64; 2]>::with_capacity(1); + let width = table.allocation_size() - 4 * 16 - 4; + insta::with_settings!({ snapshot_suffix => format!("ctrl{width}") }, { + assert_snapshot!(name, snapshot); + }); +} + +/// Every key appears twice, in different spill runs. +fn hash_input_repeated() -> Vec<Vec<(i64, i64)>> { + vec![ + distinct_a(0, 256), + distinct_a(256, 256), + distinct_a(512, 256), + distinct_a(0, 256), + distinct_a(256, 256), + distinct_a(512, 256), + ] +} + +/// `GROUP BY a, b` input sorted by `a`: one `a` per batch, so each batch +/// completes the groups of the previous one. +fn clustered_partial_input() -> Vec<Vec<(i64, i64)>> { + (0..4).map(|a| one_a(a, 0, 256)).collect() +} + +/// `GROUP BY a, b` input sorted by `a` whose first `a` spans several batches. +fn clustered_partial_input_long_run() -> Vec<Vec<(i64, i64)>> { + vec![ + one_a(0, 0, 256), + one_a(0, 256, 256), + one_a(0, 0, 256), + one_a(0, 512, 256), + one_a(1, 0, 64), + ] +} + +// --------------------------------------------------------------------------- +// Final, Single (GroupClusteringMode::None): spill or fail +// --------------------------------------------------------------------------- + +/// Final: two spills, then the merge and the replay stream run as siblings of +/// the stream's reservation (same consumer name). +#[tokio::test] +async fn final_spills_twice_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 39_000, + ..case(AggregateMode::Final, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::None); + assert_eq!(run.spill_count(), 2); + assert_snapshot_per_group_width("final_spills_twice_then_replays", &run.snapshot); + Ok(()) +} + +/// Final with temporary files disabled: the refusal is returned as an error. +#[tokio::test] +async fn final_without_spilling_returns_refusal() -> Result<()> { + let run = Case { + memory_limit: 6_000, + spill_enabled: false, + ..case(AggregateMode::Final, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.spill_count(), 0); + assert!(run.snapshot.contains("temporary files are not enabled")); + assert_snapshot_per_group_width( + "final_without_spilling_returns_refusal", + &run.snapshot, + ); + Ok(()) +} + +/// Single: two spills, then replay. +#[tokio::test] +async fn single_spills_twice_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 39_000, + ..case(AggregateMode::Single, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::None); + assert_eq!(run.spill_count(), 2); + assert_snapshot_per_group_width("single_spills_twice_then_replays", &run.snapshot); + Ok(()) +} + +/// Single with temporary files disabled: the refusal is returned as an error. +#[tokio::test] +async fn single_without_spilling_returns_refusal() -> Result<()> { + let run = Case { + memory_limit: 6_000, + spill_enabled: false, + ..case(AggregateMode::Single, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.spill_count(), 0); + assert!(run.snapshot.contains("temporary files are not enabled")); + assert_snapshot_per_group_width( + "single_without_spilling_returns_refusal", + &run.snapshot, + ); + Ok(()) +} + +/// The table fits, but growing the reservation to the materialized output +/// (`try_resize(memory_size())` on the first of several output batches) is +/// refused. Final returns the error, Single hands the output off unreserved +/// (`reserved = false`, no further calls), and Partial ignores it (and retries +/// on every batch). +/// +/// No built-in aggregate materializes an output larger than its accounted +/// state, so this uses [`PaddedOutputUdaf`]. +#[tokio::test] +async fn output_materialization_refused() -> Result<()> { + let padded = |mode| Case { + agg: Agg::PaddedOutput, + memory_limit: 50_000, + ..case(mode, &["a"], false, vec![distinct_a(0, 256)]) + }; + + let run = padded(AggregateMode::Final).run().await?; + assert_eq!(run.spill_count(), 0); + assert_snapshot_per_group_width( + "output_materialization_refused_final", + &run.snapshot, + ); + + let run = padded(AggregateMode::Single).run().await?; + assert_eq!(run.spill_count(), 0); + assert_eq!(run.output_batch_rows, vec![64; 4]); + assert_snapshot_per_group_width( + "output_materialization_refused_single", + &run.snapshot, + ); + + let run = padded(AggregateMode::Partial).run().await?; + assert_eq!(run.metric("early_emit_count"), 0); + assert_eq!(run.output_batch_rows, vec![64; 4]); + assert_snapshot_per_group_width( + "output_materialization_refused_partial", + &run.snapshot, + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Clustered streams: GROUP BY a, b over input sorted by a +// (GroupClusteringMode::Partial), and GROUP BY a (GroupClusteringMode::Full) +// --------------------------------------------------------------------------- + +fn clustered(mode: AggregateMode, memory_limit: usize) -> Case { + Case { + memory_limit, + ..case(mode, &["a", "b"], true, clustered_partial_input()) + } +} + +/// ClusteredPartial: the table fits, but the completed groups do not fit next +/// to it, so they are emitted whole. +#[tokio::test] +async fn clustered_partial_emits_completed_groups_whole() -> Result<()> { + let run = clustered(AggregateMode::Partial, 32_000).run().await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.max_output_rows(), 256); + assert_snapshot!(run.snapshot, @r" + ClusteredPartialAggregateStream[0] register can_spill=true + ClusteredPartialAggregateStream[0] try_grow 600 + ClusteredPartialAggregateStream[0] try_grow 17280 + ClusteredPartialAggregateStream[0] try_grow 13312 + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] shrink 18616 + ClusteredPartialAggregateStream[0] shrink 12576 + ClusteredPartialAggregateStream[0] unregister + => Ok(1024 rows in 7 batches) + "); + Ok(()) +} + +/// ClusteredFinal: completed groups held while sliced, then emitted whole +/// under a lower limit. +#[tokio::test] +async fn clustered_final_completed_groups() -> Result<()> { Review Comment: `clustered_final_completed_groups`/`clustered_single_completed_groups`, `clustered_final_spills_then_replays`/`clustered_single_spills_then_replays` and `final_spills_twice_then_replays`/`single_spills_twice_then_replays` record identical sequences apart from the stream name. I diffed them. You could run both modes in one test, as `output_materialization_refused` already does. Assert that the Single snapshot equals the Final one once the stream name is replaced, then snapshot once. That removes roughly 100 lines and two snapshot files. It also turns "Final and Single account the same way" into an explicit assertion, which is the property #26563 relies on. ########## datafusion/physical-plan/src/aggregates/memory_call_sequence_tests.rs: ########## @@ -0,0 +1,886 @@ +// 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. + +//! Snapshot tests that lock down the exact sequence of memory pool calls made +//! by the grouped aggregate streams, for the memory accounting paths that the +//! other tests do not reach: spill headroom and replay, refusals with +//! temporary files disabled, refused output materialization, and holding +//! completed groups in the clustered streams. The early emit of the +//! unclustered Partial and PartialReduce streams is covered by their own tests. +//! +//! Each test runs one single-partition `AggregateExec` under a +//! [`RecordingPool`] and snapshots every pool call of the aggregate consumers +//! (the stream and the merge / replay siblings that share its consumer), then +//! the query outcome. Limits were chosen by experiment so that each test +//! really takes its path, which the test also asserts through metrics or the +//! output batch sizes. Byte counts follow the current accounting of the hash +//! tables and arrays, so a change there shows up as a snapshot diff. Some of +//! them also depend on the target (see [`assert_snapshot_per_group_width`]). + +use std::fmt; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility, +}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use hashbrown::HashTable; +use insta::assert_snapshot; + +use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy}; +use crate::ExecutionPlan; +use crate::common::collect; +use crate::metrics::MetricsSet; +use crate::test::TestMemoryExec; + +/// A [`GreedyMemoryPool`] that records every call made on behalf of an +/// aggregate consumer, as `"<consumer name> <op> <bytes>"`. +#[derive(Debug)] +struct RecordingPool { + inner: GreedyMemoryPool, + calls: Mutex<Vec<(String, String)>>, +} + +impl RecordingPool { + fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + calls: Mutex::new(vec![]), + } + } + + fn record(&self, consumer: &MemoryConsumer, call: String) { + let name = consumer.name(); + if name.contains("Aggregate") { + let line = format!("{name} {call}"); + self.calls.lock().unwrap().push((name.to_string(), line)); + } + } + + /// The recorded calls grouped by consumer, in first-seen consumer order. + fn calls(&self) -> String { + let calls = self.calls.lock().unwrap(); + let mut names: Vec<&str> = vec![]; + for (name, _) in calls.iter() { + if !names.contains(&name.as_str()) { + names.push(name); + } + } + let mut out = String::new(); + for name in names { + for (_, line) in calls.iter().filter(|(n, _)| n == name) { + out.push_str(line); + out.push('\n'); + } + } + out + } +} + +impl fmt::Display for RecordingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "RecordingPool({})", self.inner) + } +} + +impl MemoryPool for RecordingPool { + fn name(&self) -> &str { + "RecordingPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + self.record( + consumer, + format!("register can_spill={}", consumer.can_spill()), + ); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + self.record(consumer, "unregister".to_string()); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(reservation.consumer(), format!("grow {additional}")); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.record(reservation.consumer(), format!("shrink {shrink}")); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + let result = self.inner.try_grow(reservation, additional); + let suffix = if result.is_ok() { "" } else { " -> refused" }; + self.record( + reservation.consumer(), + format!("try_grow {additional}{suffix}"), + ); + result + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +/// The single aggregate of a [`Case`]. +#[derive(Clone, Copy)] +enum Agg { + /// `SUM(v)`, whose output is no larger than its accounted state. + Sum, + /// [`PaddedOutputUdaf`], whose output is far larger than its accounted state. + PaddedOutput, +} + +/// One aggregate run: `agg` grouped by `group_by` over raw `(a, b, v)` rows. +struct Case { + mode: AggregateMode, + agg: Agg, + /// Group columns, a prefix of `["a", "b"]`. + group_by: &'static [&'static str], + /// Declare the input sorted by `a`, enabling group clustering. + sorted_by_a: bool, + /// The `(a, b)` keys of each input batch; every row has `v = 1`. + batches: Vec<Vec<(i64, i64)>>, + batch_size: usize, + memory_limit: usize, + spill_enabled: bool, +} + +/// What a [`Case`] did. +struct Run { + /// Recorded pool calls followed by the query outcome. + snapshot: String, + clustering: GroupClusteringMode, + metrics: MetricsSet, + output_batch_rows: Vec<usize>, +} + +impl Run { + fn metric(&self, name: &str) -> usize { + self.metrics + .sum_by_name(name) + .map(|m| m.as_usize()) + .unwrap_or(0) + } + + fn spill_count(&self) -> usize { + self.metrics.spill_count().unwrap_or(0) + } + + fn max_output_rows(&self) -> usize { + self.output_batch_rows.iter().copied().max().unwrap_or(0) + } +} + +fn raw_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])) +} + +/// Distinct keys `(start..start + n, 0)`. +fn distinct_a(start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|a| (a, 0)).collect() +} + +/// Keys `(a, start..start + n)` with one `a` value. +fn one_a(a: i64, start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|b| (a, b)).collect() +} + +impl Case { + async fn run(self) -> Result<Run> { + let raw = raw_schema(); + let group_by = PhysicalGroupBy::new_single( + self.group_by + .iter() + .map(|name| Ok((col(name, &raw)?, name.to_string()))) + .collect::<Result<_>>()?, + ); + let udaf = match self.agg { + Agg::Sum => sum_udaf(), + Agg::PaddedOutput => Arc::new(AggregateUDF::from(PaddedOutputUdaf::new())), + }; + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("v", &raw)?]) + .schema(Arc::clone(&raw)) + .alias("s") + .build()?, + )]; + + // Stages after the first read partial states: the group columns followed + // by the `SUM` state, which for Int64 is the running sum itself. + let reads_partial_state = matches!( + self.mode, + AggregateMode::Final + | AggregateMode::FinalPartitioned + | AggregateMode::PartialReduce + ); + let input_schema = if reads_partial_state { + Arc::new(super::create_schema( + &raw, + &group_by, + &aggregates, + AggregateMode::Partial, + )?) + } else { + Arc::clone(&raw) + }; + let batches = self + .batches + .iter() + .map(|keys| { + let a: ArrayRef = + Arc::new(keys.iter().map(|k| k.0).collect::<Int64Array>()); + let b: ArrayRef = + Arc::new(keys.iter().map(|k| k.1).collect::<Int64Array>()); + let v: ArrayRef = Arc::new(Int64Array::from(vec![1; keys.len()])); + let columns = if reads_partial_state { + let mut columns = vec![a, b]; + columns.truncate(self.group_by.len()); + columns.push(match self.agg { + Agg::Sum => v, + Agg::PaddedOutput => { + Arc::new(StringArray::from(vec!["x"; keys.len()])) + } + }); + columns + } else { + vec![a, b, v] + }; + Ok(RecordBatch::try_new(Arc::clone(&input_schema), columns)?) + }) + .collect::<Result<Vec<_>>>()?; + + let mut input = + TestMemoryExec::try_new(&[batches], Arc::clone(&input_schema), None)?; + if self.sorted_by_a { + input = input.try_with_sort_information(vec![ + LexOrdering::new([PhysicalSortExpr::new_default(col( + "a", + &input_schema, + )?)]) + .unwrap(), + ])?; + } + let aggregate = AggregateExec::try_new( + self.mode, + group_by, + aggregates, + vec![None], + Arc::new(input), + raw, + )?; + + let pool = Arc::new(RecordingPool::new(self.memory_limit)); + let mut runtime = + RuntimeEnvBuilder::new().with_memory_pool(Arc::clone(&pool) as _); + if !self.spill_enabled { + runtime = runtime.with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ); + } + let context = Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(self.batch_size) + .set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + ), + ) + .with_runtime(runtime.build_arc()?), + ); + + let result = match aggregate.execute(0, Arc::clone(&context)) { + Ok(stream) => collect(stream).await, + Err(e) => Err(e), + }; + let output_batch_rows: Vec<usize> = result + .iter() + .flatten() + .map(|batch| batch.num_rows()) + .collect(); + let outcome = match &result { + Ok(batches) => format!( + "Ok({} rows in {} batches)", + output_batch_rows.iter().sum::<usize>(), + batches.len() + ), + Err(e) => format!("Err({})", e.strip_backtrace()), + }; + drop(result); + drop(context); + assert_eq!(pool.reserved(), 0, "every reservation is released"); + + Ok(Run { + snapshot: format!("{}=> {outcome}", pool.calls()), + clustering: aggregate.group_clustering_mode().clone(), + metrics: aggregate.metrics().unwrap(), + output_batch_rows, + }) + } +} + +/// Bytes of every [`PaddedOutputUdaf`] output and state value. +const PAD: usize = 256; + +/// An aggregate that accounts almost no memory per group but materializes a +/// `PAD`-byte string per group as its state and its output, so materializing +/// the output can need far more memory than the table held. +#[derive(Debug, PartialEq, Eq, Hash)] +struct PaddedOutputUdaf { + signature: Signature, +} + +impl PaddedOutputUdaf { + fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int64], Volatility::Immutable), + } + } +} + +impl AggregateUDFImpl for PaddedOutputUdaf { + fn name(&self) -> &str { + "padded_output" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Utf8) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format!("{}[pad]", args.name), + DataType::Utf8, + true, + ))]) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(PaddedOutputAccumulator)) + } +} + +#[derive(Debug)] +struct PaddedOutputAccumulator; + +impl Accumulator for PaddedOutputAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + Ok(ScalarValue::Utf8(Some("x".repeat(PAD)))) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![ScalarValue::Utf8(Some("x".repeat(PAD)))]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } +} + +fn case( + mode: AggregateMode, + group_by: &'static [&'static str], + sorted_by_a: bool, + batches: Vec<Vec<(i64, i64)>>, +) -> Case { + Case { + mode, + agg: Agg::Sum, + group_by, + sorted_by_a, + batches, + batch_size: 64, + memory_limit: usize::MAX, + spill_enabled: true, + } +} + +/// Asserts `snapshot` against `snapshots/<module>__<name>@ctrl<width>.snap`. +/// +/// A group hash table's [`HashTable::allocation_size`] includes one control +/// byte per bucket plus hashbrown's SIMD group width, which is 16 with SSE2 +/// (x86_64) and 8 on aarch64. Whether a pool call includes the table depends +/// on the stream's state, so the difference cannot be normalized away, and +/// snapshots whose byte counts differ are kept once per width. +fn assert_snapshot_per_group_width(name: &str, snapshot: &str) { + // Capacity 1 allocates 4 buckets of 16 bytes, then 4 + width control bytes. + let table = HashTable::<[u64; 2]>::with_capacity(1); + let width = table.allocation_size() - 4 * 16 - 4; + insta::with_settings!({ snapshot_suffix => format!("ctrl{width}") }, { + assert_snapshot!(name, snapshot); + }); +} + +/// Every key appears twice, in different spill runs. +fn hash_input_repeated() -> Vec<Vec<(i64, i64)>> { + vec![ + distinct_a(0, 256), + distinct_a(256, 256), + distinct_a(512, 256), + distinct_a(0, 256), + distinct_a(256, 256), + distinct_a(512, 256), + ] +} + +/// `GROUP BY a, b` input sorted by `a`: one `a` per batch, so each batch +/// completes the groups of the previous one. +fn clustered_partial_input() -> Vec<Vec<(i64, i64)>> { + (0..4).map(|a| one_a(a, 0, 256)).collect() +} + +/// `GROUP BY a, b` input sorted by `a` whose first `a` spans several batches. +fn clustered_partial_input_long_run() -> Vec<Vec<(i64, i64)>> { + vec![ + one_a(0, 0, 256), + one_a(0, 256, 256), + one_a(0, 0, 256), + one_a(0, 512, 256), + one_a(1, 0, 64), + ] +} + +// --------------------------------------------------------------------------- +// Final, Single (GroupClusteringMode::None): spill or fail +// --------------------------------------------------------------------------- + +/// Final: two spills, then the merge and the replay stream run as siblings of +/// the stream's reservation (same consumer name). +#[tokio::test] +async fn final_spills_twice_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 39_000, + ..case(AggregateMode::Final, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::None); + assert_eq!(run.spill_count(), 2); + assert_snapshot_per_group_width("final_spills_twice_then_replays", &run.snapshot); + Ok(()) +} + +/// Final with temporary files disabled: the refusal is returned as an error. +#[tokio::test] +async fn final_without_spilling_returns_refusal() -> Result<()> { + let run = Case { + memory_limit: 6_000, + spill_enabled: false, + ..case(AggregateMode::Final, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.spill_count(), 0); + assert!(run.snapshot.contains("temporary files are not enabled")); + assert_snapshot_per_group_width( + "final_without_spilling_returns_refusal", + &run.snapshot, + ); + Ok(()) +} + +/// Single: two spills, then replay. +#[tokio::test] +async fn single_spills_twice_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 39_000, + ..case(AggregateMode::Single, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::None); + assert_eq!(run.spill_count(), 2); + assert_snapshot_per_group_width("single_spills_twice_then_replays", &run.snapshot); + Ok(()) +} + +/// Single with temporary files disabled: the refusal is returned as an error. +#[tokio::test] +async fn single_without_spilling_returns_refusal() -> Result<()> { + let run = Case { + memory_limit: 6_000, + spill_enabled: false, + ..case(AggregateMode::Single, &["a"], false, hash_input_repeated()) + } + .run() + .await?; + assert_eq!(run.spill_count(), 0); + assert!(run.snapshot.contains("temporary files are not enabled")); + assert_snapshot_per_group_width( + "single_without_spilling_returns_refusal", + &run.snapshot, + ); + Ok(()) +} + +/// The table fits, but growing the reservation to the materialized output +/// (`try_resize(memory_size())` on the first of several output batches) is +/// refused. Final returns the error, Single hands the output off unreserved +/// (`reserved = false`, no further calls), and Partial ignores it (and retries +/// on every batch). +/// +/// No built-in aggregate materializes an output larger than its accounted +/// state, so this uses [`PaddedOutputUdaf`]. +#[tokio::test] +async fn output_materialization_refused() -> Result<()> { + let padded = |mode| Case { + agg: Agg::PaddedOutput, + memory_limit: 50_000, + ..case(mode, &["a"], false, vec![distinct_a(0, 256)]) + }; + + let run = padded(AggregateMode::Final).run().await?; + assert_eq!(run.spill_count(), 0); + assert_snapshot_per_group_width( + "output_materialization_refused_final", + &run.snapshot, + ); + + let run = padded(AggregateMode::Single).run().await?; + assert_eq!(run.spill_count(), 0); + assert_eq!(run.output_batch_rows, vec![64; 4]); + assert_snapshot_per_group_width( + "output_materialization_refused_single", + &run.snapshot, + ); + + let run = padded(AggregateMode::Partial).run().await?; + assert_eq!(run.metric("early_emit_count"), 0); + assert_eq!(run.output_batch_rows, vec![64; 4]); + assert_snapshot_per_group_width( + "output_materialization_refused_partial", + &run.snapshot, + ); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Clustered streams: GROUP BY a, b over input sorted by a +// (GroupClusteringMode::Partial), and GROUP BY a (GroupClusteringMode::Full) +// --------------------------------------------------------------------------- + +fn clustered(mode: AggregateMode, memory_limit: usize) -> Case { + Case { + memory_limit, + ..case(mode, &["a", "b"], true, clustered_partial_input()) + } +} + +/// ClusteredPartial: the table fits, but the completed groups do not fit next +/// to it, so they are emitted whole. +#[tokio::test] +async fn clustered_partial_emits_completed_groups_whole() -> Result<()> { + let run = clustered(AggregateMode::Partial, 32_000).run().await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.max_output_rows(), 256); + assert_snapshot!(run.snapshot, @r" + ClusteredPartialAggregateStream[0] register can_spill=true + ClusteredPartialAggregateStream[0] try_grow 600 + ClusteredPartialAggregateStream[0] try_grow 17280 + ClusteredPartialAggregateStream[0] try_grow 13312 + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] try_grow 6432 -> refused + ClusteredPartialAggregateStream[0] shrink 18616 + ClusteredPartialAggregateStream[0] shrink 12576 + ClusteredPartialAggregateStream[0] unregister + => Ok(1024 rows in 7 batches) + "); + Ok(()) +} + +/// ClusteredFinal: completed groups held while sliced, then emitted whole +/// under a lower limit. +#[tokio::test] +async fn clustered_final_completed_groups() -> Result<()> { + let run = clustered(AggregateMode::Final, 39_000).run().await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.spill_count(), 0); + assert_eq!(run.max_output_rows(), 64); + assert_snapshot!(run.snapshot, @r" + ClusteredFinalAggregateStream[0] register can_spill=true + ClusteredFinalAggregateStream[0] try_grow 600 + ClusteredFinalAggregateStream[0] try_grow 18304 + ClusteredFinalAggregateStream[0] try_grow 14336 + ClusteredFinalAggregateStream[0] try_grow 5408 + ClusteredFinalAggregateStream[0] shrink 6432 + ClusteredFinalAggregateStream[0] try_grow 1024 + ClusteredFinalAggregateStream[0] try_grow 5408 + ClusteredFinalAggregateStream[0] shrink 6432 + ClusteredFinalAggregateStream[0] try_grow 1024 + ClusteredFinalAggregateStream[0] try_grow 5408 + ClusteredFinalAggregateStream[0] shrink 6432 + ClusteredFinalAggregateStream[0] shrink 19640 + ClusteredFinalAggregateStream[0] shrink 12576 + ClusteredFinalAggregateStream[0] unregister + => Ok(1024 rows in 16 batches) + "); + + let run = clustered(AggregateMode::Final, 34_000).run().await?; + assert_eq!(run.spill_count(), 0); + assert_eq!(run.max_output_rows(), 256); + assert_snapshot!(run.snapshot, @r" + ClusteredFinalAggregateStream[0] register can_spill=true + ClusteredFinalAggregateStream[0] try_grow 600 + ClusteredFinalAggregateStream[0] try_grow 18304 + ClusteredFinalAggregateStream[0] try_grow 14336 + ClusteredFinalAggregateStream[0] try_grow 5408 -> refused + ClusteredFinalAggregateStream[0] shrink 1024 + ClusteredFinalAggregateStream[0] try_grow 1024 + ClusteredFinalAggregateStream[0] try_grow 5408 -> refused + ClusteredFinalAggregateStream[0] shrink 1024 + ClusteredFinalAggregateStream[0] try_grow 1024 + ClusteredFinalAggregateStream[0] try_grow 5408 -> refused + ClusteredFinalAggregateStream[0] shrink 1024 + ClusteredFinalAggregateStream[0] shrink 19640 + ClusteredFinalAggregateStream[0] shrink 12576 + ClusteredFinalAggregateStream[0] unregister + => Ok(1024 rows in 7 batches) + "); + Ok(()) +} + +/// ClusteredFinal: one long run of `a = 0` overflows the table, which spills; +/// the spills are then merged and replayed. +#[tokio::test] +async fn clustered_final_spills_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 34_000, + ..case( + AggregateMode::Final, + &["a", "b"], + true, + clustered_partial_input_long_run(), + ) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.spill_count(), 2); + assert_snapshot!(run.snapshot, @r" + ClusteredFinalAggregateStream[0] register can_spill=true + ClusteredFinalAggregateStream[0] try_grow 600 + ClusteredFinalAggregateStream[0] try_grow 18304 + ClusteredFinalAggregateStream[0] try_grow 14336 + ClusteredFinalAggregateStream[0] try_grow 13312 -> refused + ClusteredFinalAggregateStream[0] shrink 32640 + ClusteredFinalAggregateStream[0] try_grow 4480 + ClusteredFinalAggregateStream[0] shrink 5080 + ClusteredFinalAggregateStream[0] try_grow 272 + ClusteredFinalAggregateStream[0] try_grow 49152 -> refused + ClusteredFinalAggregateStream[0] try_grow 24576 + ClusteredFinalAggregateStream[0] try_grow 6144 + ClusteredFinalAggregateStream[0] shrink 15360 + ClusteredFinalAggregateStream[0] try_grow 10240 + ClusteredFinalAggregateStream[0] try_grow 312 + ClusteredFinalAggregateStream[0] shrink 6432 + ClusteredFinalAggregateStream[0] try_grow 10224 + ClusteredFinalAggregateStream[0] try_grow 312 + ClusteredFinalAggregateStream[0] shrink 10536 + ClusteredFinalAggregateStream[0] try_grow 10224 + ClusteredFinalAggregateStream[0] try_grow 312 + ClusteredFinalAggregateStream[0] shrink 10536 + ClusteredFinalAggregateStream[0] try_grow 2544 + ClusteredFinalAggregateStream[0] try_grow 312 + ClusteredFinalAggregateStream[0] shrink 2856 + ClusteredFinalAggregateStream[0] shrink 15360 + ClusteredFinalAggregateStream[0] shrink 4080 + ClusteredFinalAggregateStream[0] shrink 312 + ClusteredFinalAggregateStream[0] unregister + => Ok(832 rows in 5 batches) + "); + Ok(()) +} + +/// ClusteredSingle: completed groups held while sliced, then emitted whole +/// under a lower limit. +#[tokio::test] +async fn clustered_single_completed_groups() -> Result<()> { + let run = clustered(AggregateMode::Single, 39_000).run().await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.spill_count(), 0); + assert_eq!(run.max_output_rows(), 64); + assert_snapshot!(run.snapshot, @r" + ClusteredSingleAggregateStream[0] register can_spill=true + ClusteredSingleAggregateStream[0] try_grow 600 + ClusteredSingleAggregateStream[0] try_grow 18304 + ClusteredSingleAggregateStream[0] try_grow 14336 + ClusteredSingleAggregateStream[0] try_grow 5408 + ClusteredSingleAggregateStream[0] shrink 6432 + ClusteredSingleAggregateStream[0] try_grow 1024 + ClusteredSingleAggregateStream[0] try_grow 5408 + ClusteredSingleAggregateStream[0] shrink 6432 + ClusteredSingleAggregateStream[0] try_grow 1024 + ClusteredSingleAggregateStream[0] try_grow 5408 + ClusteredSingleAggregateStream[0] shrink 6432 + ClusteredSingleAggregateStream[0] shrink 19640 + ClusteredSingleAggregateStream[0] shrink 12576 + ClusteredSingleAggregateStream[0] unregister + => Ok(1024 rows in 16 batches) + "); + + let run = clustered(AggregateMode::Single, 34_000).run().await?; + assert_eq!(run.spill_count(), 0); + assert_eq!(run.max_output_rows(), 256); + assert_snapshot!(run.snapshot, @r" + ClusteredSingleAggregateStream[0] register can_spill=true + ClusteredSingleAggregateStream[0] try_grow 600 + ClusteredSingleAggregateStream[0] try_grow 18304 + ClusteredSingleAggregateStream[0] try_grow 14336 + ClusteredSingleAggregateStream[0] try_grow 5408 -> refused + ClusteredSingleAggregateStream[0] shrink 1024 + ClusteredSingleAggregateStream[0] try_grow 1024 + ClusteredSingleAggregateStream[0] try_grow 5408 -> refused + ClusteredSingleAggregateStream[0] shrink 1024 + ClusteredSingleAggregateStream[0] try_grow 1024 + ClusteredSingleAggregateStream[0] try_grow 5408 -> refused + ClusteredSingleAggregateStream[0] shrink 1024 + ClusteredSingleAggregateStream[0] shrink 19640 + ClusteredSingleAggregateStream[0] shrink 12576 + ClusteredSingleAggregateStream[0] unregister + => Ok(1024 rows in 7 batches) + "); + Ok(()) +} + +/// ClusteredSingle: spills, then replays. +#[tokio::test] +async fn clustered_single_spills_then_replays() -> Result<()> { + let run = Case { + batch_size: 256, + memory_limit: 34_000, + ..case( + AggregateMode::Single, + &["a", "b"], + true, + clustered_partial_input_long_run(), + ) + } + .run() + .await?; + assert_eq!(run.clustering, GroupClusteringMode::Partial(vec![0])); + assert_eq!(run.spill_count(), 2); + assert_snapshot!(run.snapshot, @r" + ClusteredSingleAggregateStream[0] register can_spill=true + ClusteredSingleAggregateStream[0] try_grow 600 + ClusteredSingleAggregateStream[0] try_grow 18304 + ClusteredSingleAggregateStream[0] try_grow 14336 + ClusteredSingleAggregateStream[0] try_grow 13312 -> refused + ClusteredSingleAggregateStream[0] shrink 32640 + ClusteredSingleAggregateStream[0] try_grow 4480 + ClusteredSingleAggregateStream[0] shrink 5080 + ClusteredSingleAggregateStream[0] try_grow 272 + ClusteredSingleAggregateStream[0] try_grow 49152 -> refused + ClusteredSingleAggregateStream[0] try_grow 24576 + ClusteredSingleAggregateStream[0] try_grow 6144 + ClusteredSingleAggregateStream[0] shrink 15360 + ClusteredSingleAggregateStream[0] try_grow 10240 + ClusteredSingleAggregateStream[0] try_grow 312 + ClusteredSingleAggregateStream[0] shrink 6432 + ClusteredSingleAggregateStream[0] try_grow 10224 + ClusteredSingleAggregateStream[0] try_grow 312 + ClusteredSingleAggregateStream[0] shrink 10536 + ClusteredSingleAggregateStream[0] try_grow 10224 + ClusteredSingleAggregateStream[0] try_grow 312 + ClusteredSingleAggregateStream[0] shrink 10536 + ClusteredSingleAggregateStream[0] try_grow 2544 + ClusteredSingleAggregateStream[0] try_grow 312 + ClusteredSingleAggregateStream[0] shrink 2856 + ClusteredSingleAggregateStream[0] shrink 15360 + ClusteredSingleAggregateStream[0] shrink 4080 + ClusteredSingleAggregateStream[0] shrink 312 + ClusteredSingleAggregateStream[0] unregister + => Ok(832 rows in 5 batches) + "); + Ok(()) +} + +/// Full clustering never spills or emits early: a limit too small for even +/// one group returns the refusal, in every clustered stream. +#[tokio::test] +async fn full_clustering_refusal_is_an_error() -> Result<()> { Review Comment: This branch predates #26576, which is now on main. There `ClusteredPartialAggregateStream` under full clustering emits the completed groups on a refused reservation, and returns the error only when no group is complete. So "never ... emits early" no longer holds for Partial. The Partial case still errors only because its input has a single group that never completes. The Final and Single cases pin the behavior that #26628 reports as a bug. Could you rebase, describe this as current behavior, and reference #26628? Then the fix there reads as an intended snapshot change rather than a regression. The same applies to `output_materialization_refused`, where Final and Single react in opposite ways to the same refusal. If one of them is a known gap, saying so in the doc would help. ########## datafusion/physical-plan/src/aggregates/memory_call_sequence_tests.rs: ########## @@ -0,0 +1,886 @@ +// 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. + +//! Snapshot tests that lock down the exact sequence of memory pool calls made +//! by the grouped aggregate streams, for the memory accounting paths that the +//! other tests do not reach: spill headroom and replay, refusals with +//! temporary files disabled, refused output materialization, and holding +//! completed groups in the clustered streams. The early emit of the +//! unclustered Partial and PartialReduce streams is covered by their own tests. +//! +//! Each test runs one single-partition `AggregateExec` under a +//! [`RecordingPool`] and snapshots every pool call of the aggregate consumers +//! (the stream and the merge / replay siblings that share its consumer), then +//! the query outcome. Limits were chosen by experiment so that each test +//! really takes its path, which the test also asserts through metrics or the +//! output batch sizes. Byte counts follow the current accounting of the hash +//! tables and arrays, so a change there shows up as a snapshot diff. Some of +//! them also depend on the target (see [`assert_snapshot_per_group_width`]). + +use std::fmt; +use std::sync::{Arc, Mutex}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility, +}; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use hashbrown::HashTable; +use insta::assert_snapshot; + +use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy}; +use crate::ExecutionPlan; +use crate::common::collect; +use crate::metrics::MetricsSet; +use crate::test::TestMemoryExec; + +/// A [`GreedyMemoryPool`] that records every call made on behalf of an +/// aggregate consumer, as `"<consumer name> <op> <bytes>"`. +#[derive(Debug)] +struct RecordingPool { + inner: GreedyMemoryPool, + calls: Mutex<Vec<(String, String)>>, +} + +impl RecordingPool { + fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + calls: Mutex::new(vec![]), + } + } + + fn record(&self, consumer: &MemoryConsumer, call: String) { + let name = consumer.name(); + if name.contains("Aggregate") { + let line = format!("{name} {call}"); + self.calls.lock().unwrap().push((name.to_string(), line)); + } + } + + /// The recorded calls grouped by consumer, in first-seen consumer order. + fn calls(&self) -> String { + let calls = self.calls.lock().unwrap(); + let mut names: Vec<&str> = vec![]; + for (name, _) in calls.iter() { + if !names.contains(&name.as_str()) { + names.push(name); + } + } + let mut out = String::new(); + for name in names { + for (_, line) in calls.iter().filter(|(n, _)| n == name) { + out.push_str(line); + out.push('\n'); + } + } + out + } +} + +impl fmt::Display for RecordingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "RecordingPool({})", self.inner) + } +} + +impl MemoryPool for RecordingPool { + fn name(&self) -> &str { + "RecordingPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + self.record( + consumer, + format!("register can_spill={}", consumer.can_spill()), + ); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + self.record(consumer, "unregister".to_string()); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(reservation.consumer(), format!("grow {additional}")); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.record(reservation.consumer(), format!("shrink {shrink}")); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + let result = self.inner.try_grow(reservation, additional); + let suffix = if result.is_ok() { "" } else { " -> refused" }; + self.record( + reservation.consumer(), + format!("try_grow {additional}{suffix}"), + ); + result + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +/// The single aggregate of a [`Case`]. +#[derive(Clone, Copy)] +enum Agg { + /// `SUM(v)`, whose output is no larger than its accounted state. + Sum, + /// [`PaddedOutputUdaf`], whose output is far larger than its accounted state. + PaddedOutput, +} + +/// One aggregate run: `agg` grouped by `group_by` over raw `(a, b, v)` rows. +struct Case { + mode: AggregateMode, + agg: Agg, + /// Group columns, a prefix of `["a", "b"]`. + group_by: &'static [&'static str], + /// Declare the input sorted by `a`, enabling group clustering. + sorted_by_a: bool, + /// The `(a, b)` keys of each input batch; every row has `v = 1`. + batches: Vec<Vec<(i64, i64)>>, + batch_size: usize, + memory_limit: usize, + spill_enabled: bool, +} + +/// What a [`Case`] did. +struct Run { + /// Recorded pool calls followed by the query outcome. + snapshot: String, + clustering: GroupClusteringMode, + metrics: MetricsSet, + output_batch_rows: Vec<usize>, +} + +impl Run { + fn metric(&self, name: &str) -> usize { + self.metrics + .sum_by_name(name) + .map(|m| m.as_usize()) + .unwrap_or(0) + } + + fn spill_count(&self) -> usize { + self.metrics.spill_count().unwrap_or(0) + } + + fn max_output_rows(&self) -> usize { + self.output_batch_rows.iter().copied().max().unwrap_or(0) + } +} + +fn raw_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])) +} + +/// Distinct keys `(start..start + n, 0)`. +fn distinct_a(start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|a| (a, 0)).collect() +} + +/// Keys `(a, start..start + n)` with one `a` value. +fn one_a(a: i64, start: i64, n: usize) -> Vec<(i64, i64)> { + (start..start + n as i64).map(|b| (a, b)).collect() +} + +impl Case { + async fn run(self) -> Result<Run> { + let raw = raw_schema(); + let group_by = PhysicalGroupBy::new_single( + self.group_by + .iter() + .map(|name| Ok((col(name, &raw)?, name.to_string()))) + .collect::<Result<_>>()?, + ); + let udaf = match self.agg { + Agg::Sum => sum_udaf(), + Agg::PaddedOutput => Arc::new(AggregateUDF::from(PaddedOutputUdaf::new())), + }; + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("v", &raw)?]) + .schema(Arc::clone(&raw)) + .alias("s") + .build()?, + )]; + + // Stages after the first read partial states: the group columns followed + // by the `SUM` state, which for Int64 is the running sum itself. + let reads_partial_state = matches!( + self.mode, + AggregateMode::Final + | AggregateMode::FinalPartitioned + | AggregateMode::PartialReduce + ); Review Comment: This can reuse `self.mode.input_mode() == AggregateInputMode::Partial`. No case uses `FinalPartitioned` or `PartialReduce` today. ########## datafusion/physical-plan/src/aggregates/snapshots/datafusion_physical_plan__aggregates__memory_call_sequence_tests__final_spills_twice_then_repl...@ctrl16.snap: ########## @@ -0,0 +1,45 @@ +--- +source: datafusion/physical-plan/src/aggregates/memory_call_sequence_tests.rs Review Comment: The per-hashbrown-width split is what forces these into files. One way to make every snapshot inline is an insta filter that replaces the byte counts with a placeholder. Exact bytes would then stay only where the byte value is the property under test. By my reading of the 14 changes in the description, only the dropped 4 bytes per group of sort-index headroom needs exact bytes. The others show up as a missing or extra call, a flipped `refused` or `can_spill`, or a different outcome. The headroom could get one targeted assert instead. This also avoids a cost of exact bytes. Any accounting change (the #23393 series, a hashbrown or arrow bump) would churn up to 18 files, and a contributor on aarch64 cannot regenerate the `ctrl16` ones locally. ########## datafusion/physical-plan/src/aggregates/mod.rs: ########## @@ -221,6 +221,8 @@ pub mod group_values; mod grouped_hash_stream; mod grouped_topk_stream; mod hash_stream; +#[cfg(test)] +mod memory_call_sequence_tests; Review Comment: Moving to `tests/` looks feasible. The only private item the file uses is `super::create_schema`, and the Final input schema can come from a Partial `AggregateExec::try_new(..).schema()` instead. `TestMemoryExec`, `group_clustering_mode()` and `common::collect` are already reachable from outside the crate. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
