From 550121214535ee4d778bcf74e0d3c20220f8044c Mon Sep 17 00:00:00 2001 From: naman-modi Date: Wed, 29 Jul 2026 21:02:30 +0530 Subject: [PATCH 1/3] test(aggregation-fuzzer): exercise spilling under a bounded memory pool --- .../aggregation_fuzzer/context_generator.rs | 272 ++++++++++++++++-- .../fuzz_cases/aggregation_fuzzer/fuzzer.rs | 32 ++- 2 files changed, 272 insertions(+), 32 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs index 3579c6af844bb..3f4cb9baa33fd 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::{cmp, sync::Arc}; use datafusion::{ @@ -24,11 +26,87 @@ use datafusion::{ use datafusion_catalog::TableProvider; use datafusion_common::ScalarValue; use datafusion_common::{error::Result, utils::get_available_parallelism}; +use datafusion_execution::memory_pool::{ + FairSpillPool, MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::col; use rand::{Rng, rng}; use crate::fuzz_cases::aggregation_fuzzer::data_generator::Dataset; +/// Bounded pool sizes, each a fraction of the aggregate's peak reservation: +/// 2x, 1/2, 2/5, and 1/3 of the peak. 2x is generous and should not spill (the +/// light end); the smaller ones force a spill, deeper as the pool shrinks. Each +/// entry is `(numerator, denominator)`. +const SPILL_POOL_PEAK_FRACTIONS: [(usize, usize); 4] = [(2, 1), (1, 2), (2, 5), (1, 3)]; + +/// Batch size cap for spilling contexts, so the spill's sort reservation stays +/// small enough to fit under the pool. +const SPILL_BATCH_SIZE_CAP: usize = 256; + +/// Unbounded pool that records the peak total reservation. +/// +/// The baseline runs under it to measure the aggregate's real footprint, which +/// sizes the spill pools (see `generate`). Unlike `TrackConsumersPool`, the peak +/// survives consumer deregistration, so it is readable after the query. +#[derive(Debug, Default)] +pub(crate) struct PeakRecordingPool { + inner: UnboundedMemoryPool, + peak: AtomicUsize, +} + +impl PeakRecordingPool { + /// Peak total reservation seen so far, in bytes. + pub(crate) fn peak(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + fn record_peak(&self) { + self.peak + .fetch_max(self.inner.reserved(), Ordering::Relaxed); + } +} + +impl Display for PeakRecordingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "PeakRecordingPool") + } +} + +impl MemoryPool for PeakRecordingPool { + fn name(&self) -> &str { + "PeakRecordingPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record_peak(); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record_peak(); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } +} + /// SessionContext generator /// /// During testing, `generate_baseline` will be called firstly to generate a standard [`SessionContext`], @@ -43,8 +121,7 @@ use crate::fuzz_cases::aggregation_fuzzer::data_generator::Dataset; /// - `skip_partial parameters` /// - `enable_migration_aggregate` /// - hint `sorted` or not -/// - `spilling` or not (TODO, I think a special `MemoryPool` may be needed -/// to support this) +/// - memory limit (bounded pool to force spilling, or unbounded) pub struct SessionContextGenerator { /// Current testing dataset dataset: Arc, @@ -86,8 +163,13 @@ impl SessionContextGenerator { } impl SessionContextGenerator { - /// Generate the `SessionContext` for the baseline run - pub fn generate_baseline(&self) -> Result { + /// Generate the `SessionContext` for the baseline run. + /// + /// Runs under a `PeakRecordingPool` so the aggregate's peak memory can be + /// read after the query (see `generate`, which sizes the spill pools from it). + pub fn generate_baseline( + &self, + ) -> Result<(SessionContextWithParams, Arc)> { let schema = self.dataset.batches[0].schema(); let batches = self.dataset.batches.clone(); let provider = MemTable::try_new(schema, vec![batches])?; @@ -99,21 +181,30 @@ impl SessionContextGenerator { let skip_partial_params = SkipPartialParams::ensure_not_trigger(); let enable_migration_aggregate = false; + // Records peak usage; unbounded, so the baseline (the oracle) never spills. + let tracker = Arc::new(PeakRecordingPool::default()); + let builder = GeneratedSessionContextBuilder { batch_size, target_partitions, skip_partial_params, enable_migration_aggregate, sort_hint: false, + memory_pool: Some(Arc::clone(&tracker) as Arc), + // Baseline: never spill. + memory_limit: None, table_name: self.table_name.clone(), table_provider: Arc::new(provider), }; - builder.build() + Ok((builder.build()?, tracker)) } - /// Randomly generate session context - pub fn generate(&self) -> Result { + /// Randomly generate a session context. + /// + /// `agg_peak` is the aggregate's peak reservation from the baseline, used to + /// size the spilling pool. + pub fn generate(&self, agg_peak: usize) -> Result { let mut rng = rng(); let schema = self.dataset.batches[0].schema(); let batches = self.dataset.batches.clone(); @@ -125,10 +216,29 @@ impl SessionContextGenerator { // - `skip_partial`, trigger or not trigger currently for simplicity // - `enable_migration_aggregate`, true or false // - `sorted`, if found a sorted dataset, will or will not push down this information - // - `spilling`(TODO) - let batch_size = rng.random_range(1..=self.max_batch_size); + // - memory limit, `None` (unbounded) or a bounded pool that forces spilling + + // Decide spilling first; it constrains batch_size and partitions below. + let spilling = agg_peak > 0 && rng.random_bool(0.5); + + // Cap batch_size when spilling. The spill reserves ~2x the emitted batch + // to sort it, so a large batch needs a pool bigger than the aggregate's + // peak and never spills. Large batches stay covered by the unbounded rounds. + let batch_size = if spilling { + rng.random_range(1..=cmp::min(self.max_batch_size, SPILL_BATCH_SIZE_CAP)) + } else { + rng.random_range(1..=self.max_batch_size) + }; - let target_partitions = rng.random_range(1..=self.max_target_partitions); + // Single partition when spilling. `FairSpillPool` splits across the + // per-partition aggregate consumers, so with many partitions each share + // is too small and hits `ResourcesExhausted` before it can spill. Multi + // partition stays covered by the unbounded rounds. + let target_partitions = if spilling { + 1 + } else { + rng.random_range(1..=self.max_target_partitions) + }; let skip_partial_params_idx = rng.random_range(0..self.candidate_skip_partial_params.len()); @@ -137,6 +247,19 @@ impl SessionContextGenerator { let enable_migration_aggregate = rng.random_bool(0.5); + // Pool at a randomly picked fraction of the measured peak, so pressure + // ranges from light (above peak, no spill needed) to heavy across + // contexts. + let memory_limit = if spilling { + let (num, den) = SPILL_POOL_PEAK_FRACTIONS + [rng.random_range(0..SPILL_POOL_PEAK_FRACTIONS.len())]; + Some(cmp::max(1, agg_peak.saturating_mul(num) / den)) + } else { + None + }; + let memory_pool = memory_limit + .map(|bytes| Arc::new(FairSpillPool::new(bytes)) as Arc); + let (provider, sort_hint) = if rng.random_bool(0.5) && !self.dataset.sort_keys.is_empty() { // Sort keys exist and random to push down @@ -157,6 +280,8 @@ impl SessionContextGenerator { sort_hint, skip_partial_params, enable_migration_aggregate, + memory_pool, + memory_limit, table_name: self.table_name.clone(), table_provider: Arc::new(provider), }; @@ -181,6 +306,10 @@ struct GeneratedSessionContextBuilder { sort_hint: bool, skip_partial_params: SkipPartialParams, enable_migration_aggregate: bool, + /// Pool to install, or `None` for the default unbounded pool. + memory_pool: Option>, + /// Bounded pool size in bytes, recorded in the failure report; `None` if unbounded. + memory_limit: Option, table_name: String, table_provider: Arc, } @@ -210,7 +339,17 @@ impl GeneratedSessionContextBuilder { self.enable_migration_aggregate, ); - let ctx = SessionContext::new_with_config(session_config); + // Bounded pool forces the spill path; `None` keeps the default unbounded + // pool (behavior-preserving). + let ctx = match self.memory_pool { + Some(pool) => { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + SessionContext::new_with_config_rt(session_config, runtime) + } + None => SessionContext::new_with_config(session_config), + }; ctx.register_table(self.table_name, self.table_provider)?; let params = SessionContextParams { @@ -219,6 +358,7 @@ impl GeneratedSessionContextBuilder { sort_hint: self.sort_hint, skip_partial_params: self.skip_partial_params, enable_migration_aggregate: self.enable_migration_aggregate, + memory_limit: self.memory_limit, }; Ok(SessionContextWithParams { ctx, params }) @@ -234,6 +374,14 @@ pub struct SessionContextParams { sort_hint: bool, skip_partial_params: SkipPartialParams, enable_migration_aggregate: bool, + memory_limit: Option, // Bounded pool size in bytes for this run, or `None` for unbounded. +} + +impl SessionContextParams { + /// Bounded pool size for this run, or `None` if the pool was unbounded. + pub(crate) fn memory_limit(&self) -> Option { + self.memory_limit + } } /// Partial skipping parameters @@ -266,8 +414,10 @@ impl SkipPartialParams { #[cfg(test)] mod test { - use arrow::array::{RecordBatch, StringArray, UInt32Array}; + use arrow::array::{RecordBatch, StringArray, UInt32Array, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema}; + use arrow::util::pretty::pretty_format_batches; + use datafusion_common::DataFusionError; use crate::fuzz_cases::aggregation_fuzzer::check_equality_of_batches; @@ -325,14 +475,72 @@ mod test { let ctx_generator = SessionContextGenerator::new(Arc::new(dataset), "fuzz_table"); let query = "select b, count(a) from fuzz_table group by b"; - let baseline_wrapped_ctx = ctx_generator.generate_baseline().unwrap(); + let (baseline_wrapped_ctx, tracker) = ctx_generator.generate_baseline().unwrap(); + + // Run the baseline first to record the peak, then size the spilling + // contexts from it. + let base_result = baseline_wrapped_ctx + .ctx + .sql(query) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let agg_peak = tracker.peak(); + let mut random_wrapped_ctxs = Vec::with_capacity(8); for _ in 0..8 { - let ctx = ctx_generator.generate().unwrap(); + let ctx = ctx_generator.generate(agg_peak).unwrap(); random_wrapped_ctxs.push(ctx); } - let base_result = baseline_wrapped_ctx + for wrapped_ctx in random_wrapped_ctxs { + let memory_limit = wrapped_ctx.params.memory_limit(); + match wrapped_ctx.ctx.sql(query).await.unwrap().collect().await { + Ok(random_result) => { + check_equality_of_batches(&base_result, &random_result).unwrap(); + } + // A bounded pool may be too tight to spill on this tiny dataset; + // that is an acceptable skip, matching the fuzzer's behavior. + Err(e) + if memory_limit.is_some() + && matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => {} + Err(e) => panic!("unexpected error running generated context: {e}"), + } + } + } + + /// Guards against the spill knob silently becoming a no-op: a bounded context + /// on high-cardinality data must actually reach the spill path. Without this, + /// a regression that stopped forcing spills would still pass every other test. + #[tokio::test] + async fn test_generated_context_spills() { + // High cardinality: every row is its own group, so the aggregate builds a + // large state and a fraction-of-peak pool is forced to spill. No sort keys, + // so this stays on the hash aggregate (spilling) path. + let num_rows = 20_000; + let k: UInt64Array = (0..num_rows as u64).map(Some).collect(); + let schema = Schema::new(vec![Field::new("k", DataType::UInt64, true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(k)]).unwrap(); + + let mut batches = Vec::new(); + for start in (0..batch.num_rows()).step_by(1024) { + let len = cmp::min(1024, batch.num_rows() - start); + batches.push(batch.slice(start, len)); + } + + let dataset = Dataset::new(batches, vec![]); + let ctx_generator = SessionContextGenerator::new(Arc::new(dataset), "fuzz_table"); + let query = "select k, count(*) from fuzz_table group by k"; + + // Run the baseline to record the peak that sizes the spill pools. + let (baseline, tracker) = ctx_generator.generate_baseline().unwrap(); + baseline .ctx .sql(query) .await @@ -340,17 +548,29 @@ mod test { .collect() .await .unwrap(); - - for wrapped_ctx in random_wrapped_ctxs { - let random_result = wrapped_ctx - .ctx - .sql(query) - .await - .unwrap() - .collect() - .await - .unwrap(); - check_equality_of_batches(&base_result, &random_result).unwrap(); + let agg_peak = tracker.peak(); + assert!(agg_peak > 0, "baseline should have reserved memory"); + + // Generate contexts until a bounded one actually spills. Spilling is a + // random knob, so loop; on this data a bounded context spills nearly every + // time, and exhausted rounds are skipped just like the fuzzer does. + let mut spilled = false; + for _ in 0..50 { + let wrapped = ctx_generator.generate(agg_peak).unwrap(); + if wrapped.params.memory_limit().is_none() { + continue; + } + let explain = format!("EXPLAIN ANALYZE {query}"); + let Ok(rows) = wrapped.ctx.sql(&explain).await.unwrap().collect().await + else { + continue; // pool too tight this round, skip like the fuzzer + }; + let plan = pretty_format_batches(&rows).unwrap().to_string(); + if plan.contains("spill_count=") && !plan.contains("spill_count=0,") { + spilled = true; + break; + } } + assert!(spilled, "a bounded context should have spilled to disk"); } } diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs index 430762b1c28db..c9ca3be791bac 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use arrow::array::RecordBatch; use arrow::util::pretty::pretty_format_batches; -use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_common_runtime::JoinSet; use rand::{Rng, rng}; @@ -214,17 +214,19 @@ impl AggregationFuzzer { SessionContextGenerator::new(dataset_ref.clone(), &self.table_name); // Generate the baseline context, and get the baseline result firstly - let baseline_ctx_with_params = ctx_generator + let (baseline_ctx_with_params, tracker) = ctx_generator .generate_baseline() .expect("should succeed to generate baseline session context"); let baseline_result = run_sql(&sql, &baseline_ctx_with_params.ctx) .await .expect("should succeed to run baseline sql"); let baseline_result = Arc::new(baseline_result); + // Baseline peak (dominated by the aggregate) sizes the spill pools. + let agg_peak = tracker.peak(); // Generate test tasks for _ in 0..CTX_GEN_ROUNDS { let ctx_with_params = ctx_generator - .generate() + .generate(agg_peak) .expect("should succeed to generate session context"); let task = AggregationFuzzTestTask { dataset_ref: dataset_ref.clone(), @@ -271,9 +273,27 @@ struct AggregationFuzzTestTask { impl AggregationFuzzTestTask { async fn run(&self) -> Result<()> { - let task_result = run_sql(&self.sql, &self.ctx_with_params.ctx) - .await - .map_err(|e| e.context(self.context_error_report()))?; + let task_result = match run_sql(&self.sql, &self.ctx_with_params.ctx).await { + Ok(result) => result, + // A bounded pool too tight to even spill fails with + // `ResourcesExhausted`. Not a correctness bug, so skip and log it + // (no silent drop) instead of failing the run. + Err(e) + if self.ctx_with_params.params.memory_limit().is_some() + && matches!( + e.find_root(), + DataFusionError::ResourcesExhausted(_) + ) => + { + println!( + "skipping bounded run, pool too tight to spill \ + (memory_limit={:?}): {e}", + self.ctx_with_params.params.memory_limit() + ); + return Ok(()); + } + Err(e) => return Err(e.context(self.context_error_report())), + }; self.check_result(&task_result, &self.expected_result) } From 47e3f65ec05bbf5985d4c617c16e587e7b9f7907 Mon Sep 17 00:00:00 2001 From: naman-modi Date: Wed, 29 Jul 2026 21:07:54 +0530 Subject: [PATCH 2/3] nit: updated code comments --- .../fuzz_cases/aggregation_fuzzer/context_generator.rs | 10 ++++------ .../core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs | 5 ++--- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs index 3f4cb9baa33fd..dc62bffaa1525 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs @@ -221,18 +221,16 @@ impl SessionContextGenerator { // Decide spilling first; it constrains batch_size and partitions below. let spilling = agg_peak > 0 && rng.random_bool(0.5); - // Cap batch_size when spilling. The spill reserves ~2x the emitted batch - // to sort it, so a large batch needs a pool bigger than the aggregate's - // peak and never spills. Large batches stay covered by the unbounded rounds. + // Cap batch_size when spilling. The spill reserves ~2x the emitted batch to sort it, so a large batch + // needs a pool bigger than the aggregate's peak and never spills. Large batches stay covered by the unbounded rounds. let batch_size = if spilling { rng.random_range(1..=cmp::min(self.max_batch_size, SPILL_BATCH_SIZE_CAP)) } else { rng.random_range(1..=self.max_batch_size) }; - // Single partition when spilling. `FairSpillPool` splits across the - // per-partition aggregate consumers, so with many partitions each share - // is too small and hits `ResourcesExhausted` before it can spill. Multi + // Single partition when spilling. `FairSpillPool` splits across the per-partition aggregate consumers, + // so with many partitions each share is too small and hits `ResourcesExhausted` before it can spill. Multi // partition stays covered by the unbounded rounds. let target_partitions = if spilling { 1 diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs index c9ca3be791bac..abdc4d2724b86 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs @@ -275,9 +275,8 @@ impl AggregationFuzzTestTask { async fn run(&self) -> Result<()> { let task_result = match run_sql(&self.sql, &self.ctx_with_params.ctx).await { Ok(result) => result, - // A bounded pool too tight to even spill fails with - // `ResourcesExhausted`. Not a correctness bug, so skip and log it - // (no silent drop) instead of failing the run. + // A bounded pool too tight to even spill fails with `ResourcesExhausted`. Not a correctness bug, so skip + // and log it (no silent drop) instead of failing the run. Err(e) if self.ctx_with_params.params.memory_limit().is_some() && matches!( From 2fc1e4861e9dcf6ef9d670f5e3abc5cf1e3d83c8 Mon Sep 17 00:00:00 2001 From: naman-modi Date: Wed, 29 Jul 2026 23:29:53 +0530 Subject: [PATCH 3/3] chore: remove PeakRecordingPool and size spill pools from the peak_mem_used metric --- .../aggregation_fuzzer/context_generator.rs | 125 ++++-------------- .../fuzz_cases/aggregation_fuzzer/fuzzer.rs | 18 +-- .../fuzz_cases/aggregation_fuzzer/mod.rs | 40 ++++++ 3 files changed, 73 insertions(+), 110 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs index dc62bffaa1525..553a1a000827f 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::fmt::{self, Display}; -use std::sync::atomic::{AtomicUsize, Ordering}; use std::{cmp, sync::Arc}; use datafusion::{ @@ -26,9 +24,7 @@ use datafusion::{ use datafusion_catalog::TableProvider; use datafusion_common::ScalarValue; use datafusion_common::{error::Result, utils::get_available_parallelism}; -use datafusion_execution::memory_pool::{ - FairSpillPool, MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, -}; +use datafusion_execution::memory_pool::{FairSpillPool, MemoryPool}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::col; use rand::{Rng, rng}; @@ -45,68 +41,6 @@ const SPILL_POOL_PEAK_FRACTIONS: [(usize, usize); 4] = [(2, 1), (1, 2), (2, 5), /// small enough to fit under the pool. const SPILL_BATCH_SIZE_CAP: usize = 256; -/// Unbounded pool that records the peak total reservation. -/// -/// The baseline runs under it to measure the aggregate's real footprint, which -/// sizes the spill pools (see `generate`). Unlike `TrackConsumersPool`, the peak -/// survives consumer deregistration, so it is readable after the query. -#[derive(Debug, Default)] -pub(crate) struct PeakRecordingPool { - inner: UnboundedMemoryPool, - peak: AtomicUsize, -} - -impl PeakRecordingPool { - /// Peak total reservation seen so far, in bytes. - pub(crate) fn peak(&self) -> usize { - self.peak.load(Ordering::Relaxed) - } - - fn record_peak(&self) { - self.peak - .fetch_max(self.inner.reserved(), Ordering::Relaxed); - } -} - -impl Display for PeakRecordingPool { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "PeakRecordingPool") - } -} - -impl MemoryPool for PeakRecordingPool { - fn name(&self) -> &str { - "PeakRecordingPool" - } - - fn register(&self, consumer: &MemoryConsumer) { - self.inner.register(consumer); - } - - fn unregister(&self, consumer: &MemoryConsumer) { - self.inner.unregister(consumer); - } - - fn grow(&self, reservation: &MemoryReservation, additional: usize) { - self.inner.grow(reservation, additional); - self.record_peak(); - } - - fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { - self.inner.shrink(reservation, shrink); - } - - fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { - self.inner.try_grow(reservation, additional)?; - self.record_peak(); - Ok(()) - } - - fn reserved(&self) -> usize { - self.inner.reserved() - } -} - /// SessionContext generator /// /// During testing, `generate_baseline` will be called firstly to generate a standard [`SessionContext`], @@ -165,11 +99,10 @@ impl SessionContextGenerator { impl SessionContextGenerator { /// Generate the `SessionContext` for the baseline run. /// - /// Runs under a `PeakRecordingPool` so the aggregate's peak memory can be - /// read after the query (see `generate`, which sizes the spill pools from it). - pub fn generate_baseline( - &self, - ) -> Result<(SessionContextWithParams, Arc)> { + /// Runs under the default unbounded pool, so it never spills (it is the + /// oracle). The caller reads the aggregate's peak from the plan metrics of + /// this run (see `run_sql_capturing_peak`) to size the spilling pools. + pub fn generate_baseline(&self) -> Result { let schema = self.dataset.batches[0].schema(); let batches = self.dataset.batches.clone(); let provider = MemTable::try_new(schema, vec![batches])?; @@ -181,23 +114,19 @@ impl SessionContextGenerator { let skip_partial_params = SkipPartialParams::ensure_not_trigger(); let enable_migration_aggregate = false; - // Records peak usage; unbounded, so the baseline (the oracle) never spills. - let tracker = Arc::new(PeakRecordingPool::default()); - let builder = GeneratedSessionContextBuilder { batch_size, target_partitions, skip_partial_params, enable_migration_aggregate, sort_hint: false, - memory_pool: Some(Arc::clone(&tracker) as Arc), - // Baseline: never spill. + memory_pool: None, memory_limit: None, table_name: self.table_name.clone(), table_provider: Arc::new(provider), }; - Ok((builder.build()?, tracker)) + builder.build() } /// Randomly generate a session context. @@ -229,6 +158,9 @@ impl SessionContextGenerator { rng.random_range(1..=self.max_batch_size) }; + // Single partition when spilling. `FairSpillPool` splits across the per-partition aggregate consumers, + // so with many partitions each share is too small and hits `ResourcesExhausted` before it can spill. Multi + // partition stays covered by the unbounded rounds. // Single partition when spilling. `FairSpillPool` splits across the per-partition aggregate consumers, // so with many partitions each share is too small and hits `ResourcesExhausted` before it can spill. Multi // partition stays covered by the unbounded rounds. @@ -417,7 +349,9 @@ mod test { use arrow::util::pretty::pretty_format_batches; use datafusion_common::DataFusionError; - use crate::fuzz_cases::aggregation_fuzzer::check_equality_of_batches; + use crate::fuzz_cases::aggregation_fuzzer::{ + check_equality_of_batches, run_sql_capturing_peak, + }; use super::*; @@ -473,20 +407,14 @@ mod test { let ctx_generator = SessionContextGenerator::new(Arc::new(dataset), "fuzz_table"); let query = "select b, count(a) from fuzz_table group by b"; - let (baseline_wrapped_ctx, tracker) = ctx_generator.generate_baseline().unwrap(); + let baseline_wrapped_ctx = ctx_generator.generate_baseline().unwrap(); - // Run the baseline first to record the peak, then size the spilling + // Run the baseline first to capture the peak, then size the spilling // contexts from it. - let base_result = baseline_wrapped_ctx - .ctx - .sql(query) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let agg_peak = tracker.peak(); + let (base_result, agg_peak) = + run_sql_capturing_peak(query, &baseline_wrapped_ctx.ctx) + .await + .unwrap(); let mut random_wrapped_ctxs = Vec::with_capacity(8); for _ in 0..8 { @@ -536,17 +464,10 @@ mod test { let ctx_generator = SessionContextGenerator::new(Arc::new(dataset), "fuzz_table"); let query = "select k, count(*) from fuzz_table group by k"; - // Run the baseline to record the peak that sizes the spill pools. - let (baseline, tracker) = ctx_generator.generate_baseline().unwrap(); - baseline - .ctx - .sql(query) - .await - .unwrap() - .collect() - .await - .unwrap(); - let agg_peak = tracker.peak(); + // Run the baseline to capture the peak that sizes the spill pools. + let baseline = ctx_generator.generate_baseline().unwrap(); + let (_baseline_result, agg_peak) = + run_sql_capturing_peak(query, &baseline.ctx).await.unwrap(); assert!(agg_peak > 0, "baseline should have reserved memory"); // Generate contexts until a bounded one actually spills. Spilling is a diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs index abdc4d2724b86..b361228f7fb60 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/fuzzer.rs @@ -28,7 +28,7 @@ use crate::fuzz_cases::aggregation_fuzzer::{ check_equality_of_batches, context_generator::{SessionContextGenerator, SessionContextWithParams}, data_generator::{Dataset, DatasetGenerator, DatasetGeneratorConfig}, - run_sql, + run_sql, run_sql_capturing_peak, }; /// Rounds to call `generate` of [`SessionContextGenerator`] @@ -213,16 +213,18 @@ impl AggregationFuzzer { let ctx_generator = SessionContextGenerator::new(dataset_ref.clone(), &self.table_name); - // Generate the baseline context, and get the baseline result firstly - let (baseline_ctx_with_params, tracker) = ctx_generator + // Generate the baseline context, and get the baseline result firstly. + // The baseline run also gives the aggregate's peak memory, which sizes + // the spill pools for the randomized contexts. + let baseline_ctx_with_params = ctx_generator .generate_baseline() .expect("should succeed to generate baseline session context"); - let baseline_result = run_sql(&sql, &baseline_ctx_with_params.ctx) - .await - .expect("should succeed to run baseline sql"); + + let (baseline_result, agg_peak) = + run_sql_capturing_peak(&sql, &baseline_ctx_with_params.ctx) + .await + .expect("should succeed to run baseline sql"); let baseline_result = Arc::new(baseline_result); - // Baseline peak (dominated by the aggregate) sizes the spill pools. - let agg_peak = tracker.peak(); // Generate test tasks for _ in 0..CTX_GEN_ROUNDS { let ctx_with_params = ctx_generator diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs index e7ce557d2267d..50dbad6b2a3c8 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs @@ -35,10 +35,13 @@ //! different configuration parameters to control the execution path of aggregate //! queries. +use std::sync::Arc; + use arrow::array::RecordBatch; use arrow::util::pretty::pretty_format_batches; use datafusion::prelude::SessionContext; use datafusion_common::error::Result; +use datafusion_physical_plan::{ExecutionPlan, collect}; mod context_generator; mod data_generator; @@ -89,3 +92,40 @@ pub(crate) fn check_equality_of_batches( pub(crate) async fn run_sql(sql: &str, ctx: &SessionContext) -> Result> { ctx.sql(sql).await?.collect().await } + +/// Run `sql` and return the result along with the aggregate's peak memory +/// reservation, read from the `peak_mem_used` metric the aggregate publishes. +/// The peak sizes the spilling pools for the randomized contexts. +/// +/// We read the metric rather than the memory pool because the pool cannot report +/// a peak: it only tracks current usage, which drops back to zero as reservations +/// are freed, so it reads as ~0 once the query finishes. +pub(crate) async fn run_sql_capturing_peak( + sql: &str, + ctx: &SessionContext, +) -> Result<(Vec, usize)> { + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + let result = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + + Ok((result, peak_mem_used(plan.as_ref()))) +} + +/// Sum the `peak_mem_used` metric over the whole plan tree. Metrics are per +/// node, so this walks children and adds up the aggregate nodes. Returns 0 if +/// no node reports it (e.g. a query with no grouping). +/// +/// If this metric name drifts, the peak silently becomes 0 and nothing spills; +/// `test_generated_context_spills` guards against that. +fn peak_mem_used(plan: &dyn ExecutionPlan) -> usize { + let here = plan + .metrics() + .and_then(|m| m.sum_by_name("peak_mem_used")) + .map(|v| v.as_usize()) + .unwrap_or(0); + + here + plan + .children() + .iter() + .map(|child| peak_mem_used(child.as_ref())) + .sum::() +}