fix: duplicate groups after spilling in legacy hash aggregation with a nested group key

When GroupedHashAggregateStream spills and switches to merging the sorted
spill files it relies on GroupOrderingFull, which requires group ids to be
assigned in first-seen order. It recreated its group values collector to
guarantee that only when there was more than one group column, assuming a
single column always uses a sequential single-column collector.

A single nested column (Struct, Map) has no specialized single-column
collector and is handled by GroupValuesColumn through a row-backed column,
whose vectorized interning assigns ids out of input order. The ordering then
emitted groups that were still in progress and the next batch reopened them
as new groups, so the same key came out more than once with its aggregates
split between the rows.

Always recreate the collector for the merge phase.

Author: rluvaton

datafusion/core/tests/memory_limit/mod.rs

@@ -26,9 +26,13 @@ mod nlj_spill_unmatched;
 mod repartition_mem_limit;
 mod union_nullable_spill;
 mod view_spill_compaction;
-use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray};
+use arrow::array::{
+    ArrayRef, DictionaryArray, Int32Array, Int64Array, Int64Builder, ListBuilder,
+    RecordBatch, 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;
@@ -43,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};
@@ -133,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.
@@ -220,6 +224,117 @@ mod count_distinct_spill {
     }
 }
 
+/// `GROUP BY` on a single nested key in the legacy `GroupedHashAggregateStream`
+/// under a memory limit.
+///
+/// After spilling, the legacy stream re-aggregates the merged spill files with
+/// `GroupOrderingFull`, which requires group ids in first-seen order. A single
+/// nested key has no specialized single-column group values implementation and
+/// is handled by the multi-column one, whose vectorized interning does not
+/// guarantee that order. If the stream keeps that implementation for the merge
+/// phase, groups are emitted while still in progress and come out duplicated.
+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 for the merge of the spilled runs. A `FairSpillPool`
+/// caps every stage at its share, so the partial stages emit early and the
+/// final stages spill repeatedly, which produces the many small merged batches
+/// this bug needs.
+const NESTED_KEY_MEMORY_LIMIT: usize = 4 * 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),
+    ])
+}
+
+/// `st` is mostly `{list: [g, g + 1], num: g}` for group `g`, with a sprinkle
+/// of null lists, empty lists, null nums and null structs so that keys of
+/// different shapes meet in the same batches. `v` is unique.
+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()
+}
+
+const NESTED_KEY_QUERY: &str = "select st, count(v), count(distinct v), sum(v), avg(v), min(v), max(v) \
+                     from t group by st";
+
+/// Runs the query on the legacy stream, with or without a memory limit.
+async fn run_nested_key_query(memory_limit: Option<usize>) -> 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", false);
+    let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap());
+    ctx.register_table("t", Arc::new(nested_key_table()))
+        .unwrap();
+    let batches = ctx
+        .sql(NESTED_KEY_QUERY)
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    batches_to_sort_string(&batches)
+}
+
+#[tokio::test]
+async fn legacy_stream_nested_key_spill_keeps_groups_unique() {
+    let expected = run_nested_key_query(None).await;
+    let actual = run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT)).await;
+    // A duplicated group shows up as extra rows with the counts split
+    assert_eq!(actual, expected);
+}
+
 #[tokio::test]
 async fn join_by_key_multiple_partitions() {
     let config = SessionConfig::new().with_target_partitions(2);

datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs

@@ -1372,15 +1372,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.