diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 9172e34a87c7..850e07ee1d81 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/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) -> 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); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 3f6f3f8ce815..25dc1086841b 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/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.