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-24889-423df6f4fc6a6eac6d0d1457166a166dafe8030d in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0015a7521c5f6936a7e359b34dc36554870c297e Author: Raz Luvaton <[email protected]> AuthorDate: Thu Sep 10 07:35:46 2026 +0000 fix: duplicate groups after spilling in legacy hash aggregation with a fallback group key (type not having dedicated impl) (#24889) Found by the new aggregate fuzz tests: - https://github.com/apache/datafusion/pull/24881 Claude: ## Which issue does this PR close? N/A ## Rationale for this change `GROUP BY` on a single nested column (`Struct`, `Map`) returns duplicate groups when the legacy `GroupedHashAggregateStream` spills: the same key comes out as several rows, with the aggregate values split between them. It needs the legacy stream (`datafusion.execution.enable_migration_aggregate = false`), a single nested group key, and enough memory pressure to spill in a `Final` or `Single` stage. Results are silently wrong rather than an error. The migrated streams are not affected: after spilling they hand the merged input to `OrderedFinalAggregateStream`, which builds a fresh group values collector with `GroupOrdering::Full`. After spilling, the legacy stream re-aggregates the merged spill files with `GroupOrderingFull`, which requires group ids in first-seen order along the sorted input. The stream recreates its group values collector for that phase to guarantee the order, but only when there is more than one group column, assuming a single column always uses a sequential single-column collector. A single nested column has no specialized single-column collector and is served by `GroupValuesColumn` through a row-backed column, whose vectorized interning assigns new ids out of input order under hash collisions. In a merged batch of 28 sorted rows the ids came out as 0 to 4, then 9 to 13, then 5 to 8. `GroupOrderingFull` then treated a group that was still arriving as complete and emitted it, and the next batch reopened it as a new group. ## What changes are included in this PR? `GroupedHashAggregateStream` now always recreates the group values collector when it switches to merging spill files, instead of only for multi-column keys. ## What is the testing strategy for this PR? New integration test `memory_limit::nested_key_spill_keeps_groups_unique`: a 200k-row table grouped by a struct of a list and an integer, with null and empty lists, null numbers and null structs mixed in, aggregated with six aggregates including `count(distinct)` under an 8 MB `FairSpillPool` with a 64-row batch size, compared against the same query with unlimited memory. It runs both the legacy stream and the migrated streams. Without the fix the legacy run fails deterministically with 72 rows instead of 71, keys split into two rows whose counts add up to the reference. The migrated run passes with and without the fix and is kept as a regression guard. ## Are there any user-facing changes? No API changes. Queries that hit this path now return correct results. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- datafusion/core/tests/memory_limit/mod.rs | 132 ++++++++++++++++++++- .../src/aggregates/grouped_hash_stream.rs | 11 +- 2 files changed, 134 insertions(+), 9 deletions(-) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index f873f95601..1636b676ba 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -27,11 +27,12 @@ mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; use arrow::array::{ - ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray, - StringViewArray, + ArrayRef, DictionaryArray, Int32Array, Int64Array, Int64Builder, ListBuilder, + RecordBatch, StringArray, StringViewArray, StructArray, }; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; -use arrow::datatypes::{Int32Type, SchemaRef}; +use arrow::datatypes::{Fields, Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; use datafusion::assert_batches_eq; use datafusion::config::SpillCompression; @@ -46,6 +47,7 @@ use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_catalog::Session; use datafusion_catalog::streaming::StreamingTable; +use datafusion_common::test_util::batches_to_sort_string; use datafusion_common::{Result, assert_contains}; use datafusion_execution::TaskContext; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; @@ -136,7 +138,6 @@ async fn group_by_hash() { #[cfg(not(feature = "force_hash_collisions"))] mod count_distinct_spill { use super::*; - use arrow::array::Int64Array; use datafusion::assert_batches_sorted_eq; /// `count(distinct)` over integers under a memory limit. @@ -223,6 +224,129 @@ mod count_distinct_spill { } } +/// `GROUP BY` on a single nested key under a memory limit, on both the legacy +/// `GroupedHashAggregateStream` and the migrated streams. The legacy stream used +/// to emit duplicate groups after spilling. +#[tokio::test] +async fn nested_key_spill_keeps_groups_unique() { + const NESTED_KEY_ROWS: usize = 200_000; + const NESTED_KEY_GROUPS: i64 = 16; + const NESTED_KEY_BATCH_ROWS: usize = 8_192; + + /// Small enough that the final stages must spill their `count(distinct)` + /// state, large enough that the migrated final stream can hold one merged + /// batch under a `FairSpillPool` shared by four partitions. + const NESTED_KEY_MEMORY_LIMIT: usize = 8 * 1024 * 1024; + + fn nested_key_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("list", DataType::new_list(DataType::Int64, true), true), + Field::new("num", DataType::Int64, true), + ]) + } + + fn nested_key_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new_struct("st", nested_key_struct_fields(), true), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..NESTED_KEY_ROWS) + .step_by(NESTED_KEY_BATCH_ROWS) + .map(|start| { + let rows = start..(start + NESTED_KEY_BATCH_ROWS).min(NESTED_KEY_ROWS); + let mut list = ListBuilder::new(Int64Builder::new()); + let mut num = Vec::with_capacity(rows.len()); + let mut valid = Vec::with_capacity(rows.len()); + for row in rows.clone() { + let group = row as i64 % NESTED_KEY_GROUPS; + match row % 37 { + 0 => list.append_null(), + 1 => list.append(true), + _ => { + list.values().append_value(group); + list.values().append_value(group + 1); + list.append(true); + } + } + num.push((row % 41 != 0).then_some(group)); + valid.push(row % 43 != 0); + } + let st = StructArray::new( + nested_key_struct_fields(), + vec![Arc::new(list.finish()), Arc::new(Int64Array::from(num))], + Some(NullBuffer::from(valid)), + ); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(st), + Arc::new(Int64Array::from_iter_values( + rows.map(|row| row as i64), + )), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() + } + + async fn run_nested_key_query(memory_limit: Option<usize>, legacy: bool) -> String { + let mut runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder(DiskManagerBuilder::default()); + if let Some(limit) = memory_limit { + runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(limit))); + } + let config = SessionConfig::new() + .with_target_partitions(4) + // small batches: the merged spill stream arrives in many batches and + // groups span batch boundaries + .with_batch_size(64) + .set_bool("datafusion.execution.enable_migration_aggregate", !legacy); + let ctx = + SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); + ctx.register_table("t", Arc::new(nested_key_table())) + .unwrap(); + + let df = ctx + .sql( + "select st, count(v), count(distinct v), sum(v), avg(v), min(v), max(v) \ + from t group by st", + ) + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + + let task_ctx = ctx.task_ctx(); + let batches = collect_batches(Arc::clone(&plan), task_ctx) + .await + .expect("Query execution failed"); + + let spill_count = plan_spill_count(plan.as_ref()); + match memory_limit { + Some(_) => assert_ne!(spill_count, 0, "must have spilled"), + None => assert_eq!(spill_count, 0, "must not spill on unbounded memory"), + } + + batches_to_sort_string(&batches) + } + + let expected = run_nested_key_query(None, false).await; + assert_eq!( + run_nested_key_query(None, true).await, + expected, + "unbounded, legacy=true" + ); + + for legacy in [true, false] { + assert_eq!( + run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT), legacy).await, + expected, + "spilling, legacy={legacy}" + ); + } +} + /// A grouped `COUNT(DISTINCT <string>)` gets one accumulator per group, and /// each of those owns a hash set of the distinct values it has seen. Those /// sets used to be created pre-allocated, which costs far more than the diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index dd77abfeb6..d8f2543996 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -1379,15 +1379,16 @@ impl GroupedHashAggregateStream { // Recreate `group_values` for streaming merge so group ids are assigned // in first-seen order, as required by `GroupOrderingFull`. - // The pre-spill multi-column collector may use `vectorized_intern`, which - // can assign new group ids out of input order under hash collisions. + // The pre-spill collector may use `vectorized_intern`, which can assign + // new group ids out of input order under hash collisions. That is the + // multi-column collector, which also serves a single group column + // whose type has no specialized single-column collector (for example + // `Struct` or `Map`), so recreate unconditionally. let group_schema = self .spill_state .merging_group_by .group_schema(&self.spill_state.spill_schema)?; - if group_schema.fields().len() > 1 { - self.group_values = new_group_values(group_schema, &self.group_ordering)?; - } + self.group_values = new_group_values(group_schema, &self.group_ordering)?; // Use `OutOfMemoryMode::ReportError` from this point on // to ensure we don't spill the spilled data to disk again. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
