diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf4..e9f28262980a1 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -786,6 +786,65 @@ async fn test_spill_file_compressed_with_lz4_frame() -> Result<()> { Ok(()) } + +/// `covar_samp` has no native [`GroupsAccumulator`], so its per-group state is +/// held by `GroupsAccumulatorAdapter`. The adapter keeps one scratch +/// `Vec` of row indices per group, grown to the largest number of rows +/// that group has ever taken from a single input batch and retained (cleared, +/// but not deallocated) for the lifetime of the group. +/// +/// This query hands the aggregate one batch of 8192 rows per group, so 128 +/// groups retain `128 * 8192 * 4` bytes = 4 MiB of scratch capacity, four +/// times the 1 MiB limit. Everything else the aggregate holds is two orders of +/// magnitude smaller: without the scratch capacity the reported size peaks +/// around 85 KB, twelve times under the limit, and the query runs to +/// completion without ever asking the pool for what it is really using. +/// +/// `target_partitions = 1` puts the aggregate in `Single` mode, which spills +/// under memory pressure instead of emitting groups early, so the accounting +/// is observable as a spill. +/// +/// [`GroupsAccumulator`]: datafusion_expr::GroupsAccumulator +#[tokio::test] +async fn aggregate_adapter_spills_on_retained_indices() -> Result<()> { + const GROUPS: i64 = 128; + const BATCH_SIZE: i64 = 8192; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(BATCH_SIZE as usize); + let ctx = SessionContext::new_with_config_rt(config, runtime); + + let sql = format!( + "SELECT v / {BATCH_SIZE} AS g, covar_samp(v, v) AS c \ + FROM generate_series(0, {}) AS t(v) \ + GROUP BY v / {BATCH_SIZE}", + GROUPS * BATCH_SIZE - 1 + ); + + let plan = ctx.sql(&sql).await?.create_physical_plan().await?; + let batches = collect_batches(Arc::clone(&plan), ctx.task_ctx()).await?; + + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(rows, GROUPS as usize); + + let spill_count = plan_spill_count(plan.as_ref()); + assert!( + spill_count > 0, + "the aggregate retains 4 MiB of scratch indices against a 1 MiB limit, \ + so it must spill, but spill_count was {spill_count}" + ); + + Ok(()) +} + /// Run the query with the specified memory limit, /// and verifies the expected errors are returned #[derive(Clone, Debug)] diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 6704b068acf0b..9bdd81544e94e 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -102,6 +102,17 @@ pub struct GroupsAccumulatorAdapter { /// bottleneck in earlier implementations when there were many /// distinct groups. allocation_bytes: usize, + + /// The portion of [`Self::allocation_bytes`] that is the scratch + /// [`AccumulatorState::indices`] capacity held by [`Self::states`]. + /// + /// The scratch vectors are cleared, but not deallocated, at the end of every + /// batch, so their capacity is retained for the lifetime of the group. The + /// pre/post deltas taken around [`Accumulator`] work therefore see the same + /// capacity on both sides and can never charge it. This field records what + /// has already been charged so each batch charges only the growth since the + /// previous one. + indices_allocation_bytes: usize, } struct AccumulatorState { @@ -139,6 +150,7 @@ impl GroupsAccumulatorAdapter { factory: Box::new(factory), states: vec![], allocation_bytes: 0, + indices_allocation_bytes: 0, } } @@ -221,8 +233,12 @@ impl GroupsAccumulatorAdapter { let mut offsets = vec![0]; let mut offset_so_far = 0; + let mut indices_allocation_bytes = 0; for (group_index, state) in self.states.iter_mut().enumerate() { let indices = &state.indices; + // this pass already visits every group, so totalling the scratch + // capacity here costs a field read rather than a `size()` call + indices_allocation_bytes += indices.allocated_size(); if indices.is_empty() { continue; } @@ -234,6 +250,13 @@ impl GroupsAccumulatorAdapter { } let batch_indices = batch_indices.into(); + // The push loop above is the only place `indices` grows. Charge the + // growth since the previous batch here: the pre/post deltas below + // observe the identical capacity on both sides, because `f` does not + // touch `indices` and the `clear()` after it retains the capacity. + self.adjust_allocation(self.indices_allocation_bytes, indices_allocation_bytes); + self.indices_allocation_bytes = indices_allocation_bytes; + // reorder the values and opt_filter by batch_indices so that // all values for each group are contiguous, then invoke the // accumulator once per group with values @@ -284,6 +307,18 @@ impl GroupsAccumulatorAdapter { self.allocation_bytes = self.allocation_bytes.saturating_sub(size) } + /// Release the allocation held by a state that is being emitted. + /// + /// [`AccumulatorState::size`] covers the scratch `indices` capacity, so + /// this also drops it from [`Self::indices_allocation_bytes`] to keep that + /// running total equal to the capacity still held by [`Self::states`]. + fn free_state_allocation(&mut self, state: &AccumulatorState) { + self.free_allocation(state.size()); + self.indices_allocation_bytes = self + .indices_allocation_bytes + .saturating_sub(state.indices.allocated_size()); + } + /// Adjusts the allocation for something that started with /// start_size and now has new_size avoiding overflow /// @@ -325,7 +360,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { let results: Vec = states .into_iter() .map(|mut state| { - self.free_allocation(state.size()); + self.free_state_allocation(&state); state.accumulator.evaluate() }) .collect::>()?; @@ -375,7 +410,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { let mut results: Vec> = vec![]; for mut state in states { - self.free_allocation(state.size()); + self.free_state_allocation(&state); let accumulator_state = state.accumulator.state()?; results.resize_with(accumulator_state.len(), Vec::new); for (idx, state_val) in accumulator_state.into_iter().enumerate() { @@ -576,4 +611,136 @@ mod tests { assert!(!accumulator.supports_state_preserving()); Ok(()) } + + /// Accumulator whose `size()` is constant, so that the only thing that can + /// move the adapter's reported size is the adapter's own bookkeeping. + #[derive(Debug, Default)] + struct CountingAccumulator { + count: i64, + } + + impl Accumulator for CountingAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.count += values[0].len() as i64; + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(Some(self.count))) + } + + fn size(&self) -> usize { + size_of::() + } + + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(Some(self.count))]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let counts = states[0].as_primitive::(); + self.count += counts.iter().flatten().sum::(); + Ok(()) + } + } + + fn test_adapter() -> GroupsAccumulatorAdapter { + GroupsAccumulatorAdapter::new(|| Ok(Box::new(CountingAccumulator::default()))) + } + + /// The size the adapter should be reporting, computed directly from the + /// states it is holding rather than from its running total. + fn retained_size(adapter: &GroupsAccumulatorAdapter) -> usize { + adapter + .states + .iter() + .map(|state| state.size()) + .sum::() + + adapter.states.allocated_size() + } + + /// Builds a batch where group 0 takes `hot_rows` rows and group 1 takes one. + fn skewed_batch(hot_rows: usize) -> (Vec, Vec) { + let mut group_indices = vec![0; hot_rows]; + group_indices.push(1); + let values: ArrayRef = Arc::new(Int64Array::from(vec![1; group_indices.len()])); + (vec![values], group_indices) + } + + fn update(adapter: &mut GroupsAccumulatorAdapter, hot_rows: usize) { + let (values, group_indices) = skewed_batch(hot_rows); + adapter + .update_batch(&values, &group_indices, None, 2) + .unwrap(); + } + + #[test] + fn accounts_for_retained_indices_capacity() { + const HOT_ROWS: usize = 8192; + + let mut adapter = test_adapter(); + update(&mut adapter, HOT_ROWS); + + // The scratch `indices` vector for group 0 grew to hold `HOT_ROWS` + // `u32`s and keeps that capacity after `clear()`, so it has to show up + // in the reported size. + assert_eq!(adapter.size(), retained_size(&adapter)); + assert!( + adapter.size() >= HOT_ROWS * size_of::(), + "reported {} bytes, which cannot cover {} retained indices", + adapter.size(), + HOT_ROWS + ); + + // A second batch of the same shape reuses the existing capacity, so the + // reported size must not move: the capacity is charged once, not once + // per batch. + let after_first_batch = adapter.size(); + update(&mut adapter, HOT_ROWS); + assert_eq!(adapter.size(), after_first_batch); + assert_eq!(adapter.size(), retained_size(&adapter)); + + // Growing the hot group further charges only the additional capacity. + update(&mut adapter, HOT_ROWS * 4); + assert!(adapter.size() > after_first_batch); + assert_eq!(adapter.size(), retained_size(&adapter)); + + // A smaller batch afterwards keeps the capacity, and keeps charging it. + let after_growth = adapter.size(); + update(&mut adapter, 1); + assert_eq!(adapter.size(), after_growth); + assert_eq!(adapter.size(), retained_size(&adapter)); + } + + #[test] + fn releases_retained_indices_capacity_on_emit() { + const HOT_ROWS: usize = 8192; + + let mut adapter = test_adapter(); + update(&mut adapter, HOT_ROWS); + + // Emitting the hot group must release its retained capacity, and must + // not release capacity that belongs to the groups left behind. + adapter.state(EmitTo::First(1)).unwrap(); + assert_eq!(adapter.size(), retained_size(&adapter)); + + update(&mut adapter, HOT_ROWS); + assert_eq!(adapter.size(), retained_size(&adapter)); + + adapter.evaluate(EmitTo::All).unwrap(); + assert_eq!(adapter.size(), 0); + assert_eq!(adapter.size(), retained_size(&adapter)); + } + + #[test] + fn merge_batch_accounts_for_retained_indices_capacity() { + const HOT_ROWS: usize = 8192; + + let mut adapter = test_adapter(); + let (values, group_indices) = skewed_batch(HOT_ROWS); + adapter.merge_batch(&values, &group_indices, 2).unwrap(); + + assert_eq!(adapter.size(), retained_size(&adapter)); + assert!(adapter.size() >= HOT_ROWS * size_of::()); + } }