From e2cecbee350401adeacc7cd5f7835b79891bfd39 Mon Sep 17 00:00:00 2001 From: Shehab <11789402+shehabgamin@users.noreply.github.com> Date: Sun, 25 Oct 2026 04:28:22 -0700 Subject: [PATCH 1/2] fix: Compact scalars during merge --- .../functions-aggregate/src/first_last.rs | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index fba4aa697778e..bf19e112da223 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -977,6 +977,7 @@ impl Accumulator for TrivialFirstValueAccumulator { && !first.is_empty() { self.first = ScalarValue::try_from_array(first, 0)?; + self.first.compact(); self.is_set = true; } } @@ -2972,4 +2973,26 @@ mod tests { assert_eq!(int64_values(&out), vec![Some(500)]); Ok(()) } + + #[test] + fn test_first_value_merge_compacts_list_state() -> Result<()> { + let values: ArrayRef = + Arc::new(ListArray::from_iter_primitive::( + repeat_with(|| Some(vec![Some(1)])).take(10_000), + )); + + let mut updated = + TrivialFirstValueAccumulator::try_new(values.data_type(), false)?; + let mut merged = + TrivialFirstValueAccumulator::try_new(values.data_type(), false)?; + + updated.update_batch(&[Arc::clone(&values)])?; + + let flags: ArrayRef = Arc::new(BooleanArray::from(vec![true; values.len()])); + merged.merge_batch(&[values, flags])?; + + assert_eq!(merged.evaluate()?, updated.evaluate()?); + assert_eq!(merged.size(), updated.size()); + Ok(()) + } } From 7b84fff220f9a6aabe063cfa0d1eb2383909e304 Mon Sep 17 00:00:00 2001 From: Shehab <11789402+shehabgamin@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:32:16 -0700 Subject: [PATCH 2/2] add test case --- .../test_files/aggregate_memory_spill.slt | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 79d1213ce4316..0894a5f9099c1 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -239,6 +239,30 @@ FROM ( 08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[early_emit_count=] +# --- Case H: DISTINCT ON with nested FIRST_VALUE state after spilling --- +# Merged list states must not retain the entire replay batch for every group. +# Check the complete nested payload as well as row count. +statement ok +SET datafusion.execution.target_partitions = 1 + +query II +SELECT count(*), count(*) FILTER ( + WHERE contents = [named_struct('key', k, 'value', repeat('v', 224))] +) +FROM ( + SELECT DISTINCT ON (k) k, + [named_struct('key', k, 'value', repeat('v', 224))] AS contents + FROM ( + SELECT (v * 7) % 8192 AS k + FROM generate_series(0, 8191) AS t(v) + ) +) +---- +8192 8192 + +statement ok +SET datafusion.execution.target_partitions = 4 + # Restore settings to slt runner defaults statement ok RESET datafusion.runtime.memory_limit