This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24858-9a1e1d589b1aeef482bd26d035182dee4980ed43 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 20475d6e34ad9f8bbe020b9da84cb54ebc5cd2a0 Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Wed Sep 9 21:03:19 2026 +0000 fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter (#24858) `GroupsAccumulatorAdapter` never charges the capacity of its scratch `indices` vector to the `MemoryPool`. An aggregate holds megabytes that the pool does not see, so a memory limit does not stop it. ## Reproduction This needs only `datafusion-cli`. There is no patch, no custom allocator and no data file. ```sql -- repro.sql SET datafusion.execution.target_partitions = 1; SET datafusion.execution.batch_size = 8192; EXPLAIN ANALYZE SELECT v / 8192 AS g, covar_samp(v, v) AS c FROM generate_series(0, 1048575) AS t(v) GROUP BY v / 8192; ``` ``` datafusion-cli -m 1M -f repro.sql ``` `covar_samp` has no specialized `GroupsAccumulator`, so it runs through `GroupsAccumulatorAdapter`. The query makes 128 groups. Each group gets one full 8192-row batch. The scratch vectors hold `128 * 8192 * 4` bytes, which is 4 MiB against a 1 MiB limit. Merge base `da89c7c85b`. The aggregate runs past the limit and does not spill: ``` AggregateExec: mode=Single, gby=[v@1 / 8192 as t.v / Int64(8192)], aggr=[covar_samp(t.v,t.v)], metrics=[output_rows=128, elapsed_compute=14.48ms, output_bytes=2.0 KB, output_batches=1, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, ...] ``` This branch. The aggregate sees the same bytes and spills: ``` AggregateExec: mode=Single, gby=[v@1 / 8192 as t.v / Int64(8192)], aggr=[covar_samp(t.v,t.v)], metrics=[output_rows=128, elapsed_compute=15.02ms, output_bytes=2.0 KB, output_batches=2, spill_count=5, spilled_bytes=9.2 KB, spilled_rows=128, ...] ``` | | merge base `da89c7c85b` | this branch | | --- | --- | --- | | `spill_count` | 0 | 5 | | `spilled_bytes` | 0.0 B | 9.2 KB | | `spilled_rows` | 0 | 128 | The query returns the same 128 rows on both builds. The run takes under a second. Both numbers repeat exactly across runs. `target_partitions = 1` makes the effect visible. The planner then folds the aggregate into one `AggregateMode::Single` node, which spills. An `AggregateMode::Partial` node uses `OutOfMemoryMode::EmitEarly` and sheds the bytes instead. ## Which issue does this PR close? No existing issue. I found this when I investigated a production out of memory. I can file an issue if you want it in the changelog. ## Rationale for this change The adapter keeps a running total in `allocation_bytes`. It measures `AccumulatorState::size()` before and after the accumulator work, then charges the difference. The scratch vector grows in the per-row push loop. That loop runs before the adapter measures `sizes_pre`. The `indices.clear()` call after the work keeps the capacity. Both measurements therefore see the same capacity, the difference is always zero, and the adapter never charges the capacity. `evaluate` and `state` have the opposite error. Both call `free_allocation(state.size())` and release a capacity that the adapter never charged. `allocation_bytes` thus falls to zero across the partial emits. The size of the hole is `groups * rows_per_batch * 4` bytes. ### How large the error is An instrumented allocator measured these numbers, so the CLI cannot reproduce them. A counting `GlobalAlloc` gives the heap that the query holds. A peak-recording `MemoryPool` gives the reported bytes. | groups | heap held | reported, base | error | reported, this branch | error | | --- | --- | --- | --- | --- | --- | | 512 | 17,197,804 | 145,408 | 99.15% | 16,922,624 | 1.60% | | 4,096 | 135,061,228 | 704,512 | 99.48% | 134,922,240 | 0.10% | The peak heap agrees between the two builds to within 8 bytes. The memory use does not change. Only the reported number moves. ## What changes are included in this PR? A new private field `indices_allocation_bytes` records the capacity that the adapter already charged. Each batch totals the current capacity in the loop that already visits every group, then charges only the growth. An emit removes the capacity of the emitted state from that total. This adds no `size()` call and no per-row work. It adds one `usize` addition per group per batch to an existing loop. The invariant is `allocation_bytes == sum(state.size()) + states.allocated_size()`. Two end-to-end tests in `datafusion/core/tests/memory_limit/mod.rs` cover it through the memory pool rather than through the adapter's internals. Both run the `covar_samp` query above with `target_partitions = 1`, so the aggregate runs in `Single` mode and spills under memory pressure. | test | batches per group | limit | asserts | on `da89c7c85b` | charging per batch | | --- | --- | --- | --- | --- | --- | | `aggregate_adapter_spills_on_retained_indices` | 1 | 1 MiB | spills | fails, `spill_count = 0` | passes | | `aggregate_adapter_charges_retained_indices_once` | 4 | 8 MiB | does not spill | passes | fails, spills | The second test guards the other direction: the capacity is 4 MiB however many batches a group receives, so charging it once per batch would report 16 MiB and spill under the 8 MiB limit. I verified both columns by running the tests against the merge base and against a variant of this branch that charges the full capacity every batch. ## Are there any user-facing changes? No public API change and no change to query results. Only the accounting arithmetic changes. A memory-limited aggregate now reports its true size to the `MemoryPool`. It can therefore spill, or fail where it cannot spill, in cases where it previously ran past its limit. ## A note on metrics The number this PR corrects is the aggregate's pool reservation, and nothing exposes it. With the default `enable_migration_aggregate = true` a `Single` mode `GROUP BY` runs on `SingleHashAggregateStream`, which records no memory metric at all. Only the legacy `GroupedHashAggregateStream` keeps a `peak_mem_used` gauge, and neither `EXPLAIN ANALYZE` nor `EXPLAIN ANALYZE VERBOSE` prints it. The reproduction and the tests therefore observe the reservation indirectly, as `spill_count` under a fixed limit. If the migrated streams exposed a peak memory gauge, a reviewer could see this bug with no memory limit at all, and the tests could assert on the reported bytes directly. --------- Co-authored-by: Yongting You <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]> --- datafusion/core/tests/memory_limit/mod.rs | 101 +++++++++++++++++++++ .../src/aggregate/groups_accumulator.rs | 32 ++++++- 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index f873f95601..b72eaaa28c 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -946,6 +946,107 @@ async fn test_spill_file_compressed_with_lz4_frame() -> Result<()> { Ok(()) } + +/// Number of groups the `covar_samp` queries below produce. +const ADAPTER_GROUPS: i64 = 128; +/// Rows per input batch for the `covar_samp` queries below. +const ADAPTER_BATCH_SIZE: i64 = 8192; +/// Bytes of scratch row indices `GroupsAccumulatorAdapter` retains for the +/// `covar_samp` queries below: one `u32` per row of the largest batch each +/// group has ever received, kept for the lifetime of the group. +const ADAPTER_RETAINED_BYTES: usize = + (ADAPTER_GROUPS * ADAPTER_BATCH_SIZE) as usize * size_of::<u32>(); + +/// Runs a `GROUP BY` `covar_samp` query under `memory_limit` and returns the +/// executed plan so the caller can inspect its metrics. +/// +/// `covar_samp` has no native [`GroupsAccumulator`], so its per-group state is +/// held by `GroupsAccumulatorAdapter`. The adapter keeps one scratch `Vec<u32>` +/// 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. +/// +/// The query hands each of the [`ADAPTER_GROUPS`] groups `batches_per_group` +/// consecutive full batches of [`ADAPTER_BATCH_SIZE`] rows, so the adapter +/// retains [`ADAPTER_RETAINED_BYTES`] (4 MiB) of scratch capacity however many +/// batches each group receives. Everything else the aggregate holds is two +/// orders of magnitude smaller. +/// +/// `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 +async fn run_adapter_query( + memory_limit: usize, + batches_per_group: i64, +) -> Result<Arc<dyn ExecutionPlan>> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(ADAPTER_BATCH_SIZE as usize); + let ctx = SessionContext::new_with_config_rt(config, runtime); + + let rows_per_group = ADAPTER_BATCH_SIZE * batches_per_group; + let sql = format!( + "SELECT v / {rows_per_group} AS g, covar_samp(v, v) AS c \ + FROM generate_series(0, {}) AS t(v) \ + GROUP BY v / {rows_per_group}", + ADAPTER_GROUPS * rows_per_group - 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, ADAPTER_GROUPS as usize); + + Ok(plan) +} + +/// The scratch capacity `GroupsAccumulatorAdapter` retains is four times the +/// memory limit, so the aggregate must spill. Without the capacity everything +/// the aggregate reports is far under the limit, and the query runs to +/// completion without ever asking the pool for what it is really using. +#[tokio::test] +async fn aggregate_adapter_spills_on_retained_indices() -> Result<()> { + let plan = run_adapter_query(ADAPTER_RETAINED_BYTES / 4, 1).await?; + + let spill_count = plan_spill_count(plan.as_ref()); + assert!( + spill_count > 0, + "the aggregate retains {ADAPTER_RETAINED_BYTES} bytes of scratch indices \ + against a limit of a quarter of that, so it must spill, \ + but spill_count was {spill_count}" + ); + + Ok(()) +} + +/// The retained scratch capacity is charged once, not once per batch. Every +/// group receives four batches, so charging the capacity per batch would +/// report four times the retained bytes, and the memory limit of twice the +/// retained bytes fits the aggregate only if it is charged once. +#[tokio::test] +async fn aggregate_adapter_charges_retained_indices_once() -> Result<()> { + let plan = run_adapter_query(ADAPTER_RETAINED_BYTES * 2, 4).await?; + + let spill_count = plan_spill_count(plan.as_ref()); + assert_eq!( + spill_count, 0, + "the aggregate retains {ADAPTER_RETAINED_BYTES} bytes of scratch indices \ + under a limit of twice that, so it must not spill" + ); + + 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 6704b068ac..ad01fbce5b 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -102,6 +102,10 @@ 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`]. + indices_allocation_bytes: usize, } struct AccumulatorState { @@ -139,6 +143,7 @@ impl GroupsAccumulatorAdapter { factory: Box::new(factory), states: vec![], allocation_bytes: 0, + indices_allocation_bytes: 0, } } @@ -221,8 +226,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 +243,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 +300,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 +353,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { let results: Vec<ScalarValue> = states .into_iter() .map(|mut state| { - self.free_allocation(state.size()); + self.free_state_allocation(&state); state.accumulator.evaluate() }) .collect::<Result<_>>()?; @@ -375,7 +403,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { let mut results: Vec<Vec<ScalarValue>> = 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() { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
