This is an automated email from the ASF dual-hosted git repository.
adriangb pushed a commit to branch claude/groups-accumulator-indices-accounting
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to
refs/heads/claude/groups-accumulator-indices-accounting by this push:
new 65b3f7e975 test: cover the adapter's retained indices end to end
instead of via internals
65b3f7e975 is described below
commit 65b3f7e9752997d86e626064669e9ef4ebf525dd
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Wed Sep 9 08:51:40 2026 -0500
test: cover the adapter's retained indices end to end instead of via
internals
Review feedback asked for the accounting to be tested against end
behaviour rather than through the adapter's private fields. Drop the
three unit tests that recomputed `allocation_bytes` from the states and
extend the memory-limit test instead:
* `aggregate_adapter_spills_on_retained_indices` keeps the original
check: 4 MiB of retained scratch indices against a 1 MiB limit must
spill. It fails on main with `spill_count = 0`.
* `aggregate_adapter_charges_retained_indices_once` hands every group
four batches under a limit of twice the retained bytes and asserts
that the query does not spill. Charging the capacity once per batch
would report four times the retained bytes and fail this test.
Both share one query runner so the two limits and batch counts are the
only differences between them.
Co-Authored-By: Claude Fable 5.1 <[email protected]>
---
datafusion/core/tests/memory_limit/mod.rs | 88 ++++++++++----
.../src/aggregate/groups_accumulator.rs | 132 ---------------------
2 files changed, 65 insertions(+), 155 deletions(-)
diff --git a/datafusion/core/tests/memory_limit/mod.rs
b/datafusion/core/tests/memory_limit/mod.rs
index e9f2826298..9f68aaabc1 100644
--- a/datafusion/core/tests/memory_limit/mod.rs
+++ b/datafusion/core/tests/memory_limit/mod.rs
@@ -787,31 +787,42 @@ 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.
+/// 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.
///
-/// 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.
+/// 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
-#[tokio::test]
-async fn aggregate_adapter_spills_on_retained_indices() -> Result<()> {
- const GROUPS: i64 = 128;
- const BATCH_SIZE: i64 = 8192;
-
+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(1024 * 1024)))
+ .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit)))
.with_disk_manager_builder(
DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
)
@@ -819,27 +830,58 @@ async fn aggregate_adapter_spills_on_retained_indices()
-> Result<()> {
let config = SessionConfig::new()
.with_target_partitions(1)
- .with_batch_size(BATCH_SIZE as usize);
+ .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 / {BATCH_SIZE} AS g, covar_samp(v, v) AS c \
+ "SELECT v / {rows_per_group} AS g, covar_samp(v, v) AS c \
FROM generate_series(0, {}) AS t(v) \
- GROUP BY v / {BATCH_SIZE}",
- GROUPS * BATCH_SIZE - 1
+ 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, GROUPS as usize);
+ 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 4 MiB of scratch indices against a 1 MiB limit,
\
- so it must spill, but spill_count was {spill_count}"
+ "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(())
diff --git
a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
index b9b6a66702..ad01fbce5b 100644
--- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
+++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
@@ -604,136 +604,4 @@ 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<ScalarValue> {
- Ok(ScalarValue::Int64(Some(self.count)))
- }
-
- fn size(&self) -> usize {
- size_of::<Self>()
- }
-
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- Ok(vec![ScalarValue::Int64(Some(self.count))])
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- let counts = states[0].as_primitive::<Int64Type>();
- self.count += counts.iter().flatten().sum::<i64>();
- 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::<usize>()
- + 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<ArrayRef>, Vec<usize>) {
- 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::<u32>(),
- "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::<u32>());
- }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
