This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6700d02160 test: bound merge fan-in in the final hash aggregate spill
test (#25424)
6700d02160 is described below
commit 6700d02160b1d2aee228455a715813e4a586cfc3
Author: Kumar Ujjawal <[email protected]>
AuthorDate: Sun Sep 20 14:08:09 2026 +0000
test: bound merge fan-in in the final hash aggregate spill test (#25424)
## Which issue does this PR close?
- Closes #25423.
## Rationale for this change
Case G of `aggregate_memory_spill.slt` runs four final partitions in a 1
MiB greedy pool. A partition that replays its spills cannot spill again.
Its replay table fails to grow when its own merge buffers and the state
of the other partitions fill the pool. CI hit this on two merge-queue
runs (see the issue).
## What changes are included in this PR?
- `aggregate_memory_spill.slt`: set
`datafusion.runtime.max_spill_merge_fan_in = 2` for Case G and reset it
at the end, as #25252 did for `ordered_aggregate_spill.slt`. The memory
limit stays at 1 MiB. Cases A to F are unchanged.
- `FinalHashAggregateStream::new_with_input` builds the stream over a
given input, like `OrderedFinalAggregateStream::new_with_input`. `new`
calls it.
- Tests in `hash_stream.rs` spill and replay partition 0 while
partitions 1 to 3 keep their state in a shared 1 MiB greedy pool.
Channel inputs return `Pending` after each batch, so the interleaving
does not depend on task scheduling. Input sizes follow the observed
reservations: the held partitions read batches until they reserve 512
KiB, and partition 0 reads batches until it has spilled the requested
number of times. So the cases do not depend on the memory accounting of
a platform or feature set, such as `force_hash_collisions`. The tests
check the grouped results against a reference, the spill metrics, and
that the pool returns to the held size after completion, after a drop
during replay, and after an input error, and to zero after all streams
drop. These tests cover the spill and replay lifecycle; the SLT case
covers the fan-in regression.
## What is the testing strategy for this PR?
Local runs on this branch:
- The SLT file 480 times as 16 concurrent copies, each in its own
working directory. Before the change: 12 failures, all in Case G with
the `ResourcesExhausted` error from the issue. After the change: 0
failures.
- The new tests 50 times with default features and 20 times with
`force_hash_collisions` enabled in `datafusion-common`, once with
`--profile ci` and that feature as the extended CI job runs them, and
all `datafusion-physical-plan` unit tests.
- `cargo test -p datafusion-sqllogictest --test sqllogictests` for all
files.
- `cargo fmt --all -- --check` and `cargo clippy --all-targets
--all-features -- -D warnings`.
## Are there any user-facing changes?
No.
---
.../physical-plan/src/aggregates/hash_stream.rs | 299 ++++++++++++++++++++-
.../test_files/aggregate_memory_spill.slt | 10 +
2 files changed, 306 insertions(+), 3 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs
b/datafusion/physical-plan/src/aggregates/hash_stream.rs
index 8e74e7bec7..d340bf5d43 100644
--- a/datafusion/physical-plan/src/aggregates/hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs
@@ -736,8 +736,18 @@ impl FinalHashAggregateStream {
));
debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear);
- let schema = Arc::clone(&agg.schema);
let input = agg.input.execute(partition, Arc::clone(context))?;
+ Self::new_with_input(agg, context, partition, input)
+ }
+
+ /// Builds the stream over `input` instead of executing the plan's input.
+ pub(in crate::aggregates) fn new_with_input(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ input: SendableRecordBatchStream,
+ ) -> Result<Self> {
+ let schema = Arc::clone(&agg.schema);
let input_schema = input.schema();
let batch_size = context.session_config().batch_size();
let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
@@ -1022,18 +1032,24 @@ mod tests {
use super::*;
use crate::aggregates::{AggregateMode, PhysicalGroupBy};
+ use crate::common::collect;
use crate::execution_plan::ExecutionPlan;
use crate::test::TestMemoryExec;
use crate::test::exec::BarrierExec;
- use arrow::array::{AsArray, Int32Array, Int64Array};
+ use arrow::array::{AsArray, Int32Array, Int64Array, StringViewArray};
use arrow::datatypes::{DataType, Field, Int32Type, Schema};
use datafusion_common::Result;
+ use datafusion_execution::config::SessionConfig;
+ use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool};
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_functions_aggregate::count::count_udaf;
+ use datafusion_functions_aggregate::{min_max::min_udaf, sum::sum_udaf};
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_physical_expr::expressions::col;
- use futures::StreamExt;
+ use futures::channel::mpsc;
+ use futures::{FutureExt, StreamExt};
+ use std::collections::BTreeMap;
#[tokio::test]
async fn test_partial_hash_stream_double_emission_race_condition_bug() ->
Result<()> {
@@ -1607,4 +1623,281 @@ mod tests {
Ok(())
}
+
+ #[derive(Clone, Copy)]
+ enum Finish {
+ Collect,
+ DropDuringReplay,
+ InputError,
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_replay_with_other_partitions_holding_state() ->
Result<()> {
+ for spills in [1, 2, 3] {
+ run_shared_pool_case(spills, 1024 * 1024, Finish::Collect).await?;
+ }
+ // An unlimited pool produces the reference results without spilling.
+ run_shared_pool_case(0, 10 * 1024 * 1024, Finish::Collect).await
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_replay_releases_memory_on_drop() -> Result<()> {
+ run_shared_pool_case(1, 1024 * 1024, Finish::DropDuringReplay).await
+ }
+
+ #[tokio::test]
+ async fn final_hash_spill_releases_memory_on_input_error() -> Result<()> {
+ run_shared_pool_case(1, 1024 * 1024, Finish::InputError).await
+ }
+
+ /// Partition 0 spills `spills` times and replays while partitions 1..3
+ /// keep their aggregate state in the same greedy pool. With `spills` at 0,
+ /// partition 0 reads a fixed input that fits in memory.
+ ///
+ /// Input sizes follow the observed reservations, so the cases do not
+ /// depend on the memory accounting of a platform or feature set.
+ ///
+ /// These cases cover the spill and replay lifecycle in a shared pool:
+ /// results, spill metrics, and memory release. Case G in
+ /// `aggregate_memory_spill.slt` covers the merge fan-in regression for
+ /// issue #25423.
+ async fn run_shared_pool_case(
+ spills: usize,
+ limit: usize,
+ finish: Finish,
+ ) -> Result<()> {
+ const PARTITIONS: usize = 4;
+ /// Partitions 1..3 hold at least this much state.
+ const HELD_BYTES: usize = 512 * 1024;
+ /// A case that never reaches its target fails instead of looping.
+ const MAX_BATCHES: i64 = 1000;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ Field::new("v", DataType::Int64, false),
+ Field::new("s", DataType::Utf8View, false),
+ ]));
+ let groups = PhysicalGroupBy::new_single(vec![
+ (col("a", &schema)?, "a".into()),
+ (col("b", &schema)?, "b".into()),
+ ]);
+ let expressions = vec![
+ Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum")
+ .build()?,
+ ),
+ Arc::new(
+ AggregateExprBuilder::new(min_udaf(), vec![col("s", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("min")
+ .build()?,
+ ),
+ ];
+ let empty = TestMemoryExec::try_new_exec(&[vec![]],
Arc::clone(&schema), None)?;
+ let partial = AggregateExec::try_new(
+ AggregateMode::Partial,
+ groups.clone(),
+ expressions.clone(),
+ vec![None; 2],
+ empty,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let input = TestMemoryExec::try_new_exec(
+ &vec![vec![]; PARTITIONS],
+ Arc::clone(&partial_schema),
+ None,
+ )?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::FinalPartitioned,
+ groups.as_final(),
+ expressions,
+ vec![None; 2],
+ input,
+ schema,
+ )?;
+ assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear);
+
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(limit));
+ let context = Arc::new(
+ TaskContext::default()
+ .with_session_config(SessionConfig::new().with_batch_size(128))
+ .with_runtime(
+ RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&pool))
+ .build_arc()?,
+ ),
+ );
+ let mut streams = vec![];
+ let mut senders = vec![];
+ for partition in 0..PARTITIONS {
+ let (sender, receiver) = mpsc::unbounded();
+ let input = Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&partial_schema),
+ receiver,
+ ));
+ let stream = FinalHashAggregateStream::new_with_input(
+ &aggregate, &context, partition, input,
+ )?
+ .into_stream();
+ senders.push(sender);
+ streams.push(stream);
+ }
+ let mut expected = BTreeMap::new();
+ let mut make_batch = |partition: i64, start: i64| {
+ for value in start..start + 128 {
+ let entry = expected
+ .entry((
+ partition,
+ if partition == 0 && value % 128 == 0 {
+ 0
+ } else {
+ value
+ },
+ ))
+ .or_insert_with(|| (0, (value % 2).to_string()));
+ entry.0 += value * 2;
+ }
+ RecordBatch::try_new(
+ Arc::clone(&partial_schema),
+ vec![
+ Arc::new(Int64Array::from(vec![partition; 128])),
+ Arc::new(Int64Array::from_iter_values((start..start +
128).map(
+ |value| {
+ if partition == 0 && value % 128 == 0 {
+ 0
+ } else {
+ value
+ }
+ },
+ ))),
+ Arc::new(Int64Array::from_iter_values(
+ (start..start + 128).map(|v| v * 2),
+ )),
+ Arc::new(StringViewArray::from_iter_values(
+ (start..start + 128).map(|v| if v % 2 == 0 { "0" }
else { "1" }),
+ )),
+ ],
+ )
+ .unwrap()
+ };
+ // Channel inputs return Pending after each supplied batch, so the
+ // interleaving below does not depend on task scheduling.
+ let mut feed = |partition: usize, batch: i64| {
+ senders[partition]
+ .unbounded_send(Ok(make_batch(partition as i64, batch * 128)))
+ .unwrap();
+ assert!(streams[partition].next().now_or_never().is_none());
+ };
+ // Keep the state of partitions 1..3 live while partition 0 spills and
+ // replays.
+ let mut held_batches = 0;
+ while pool.reserved() < HELD_BYTES {
+ assert!(held_batches < MAX_BATCHES, "held state stays small");
+ for partition in 1..PARTITIONS {
+ feed(partition, held_batches);
+ }
+ held_batches += 1;
+ }
+ let held = pool.reserved();
+ let spill_count = ||
aggregate.metrics().unwrap().spill_count().unwrap();
+ assert_eq!(spill_count(), 0);
+ // Feed partition 0 until it has spilled `spills` times. Key (0, 0)
+ // repeats in every batch, so replay must merge its sum across runs.
+ let mut batches = 0;
+ loop {
+ let done = if spills == 0 {
+ batches == 70
+ } else {
+ spill_count() >= spills
+ };
+ if done {
+ break;
+ }
+ assert!(batches < MAX_BATCHES, "partition 0 did not spill");
+ feed(0, batches);
+ batches += 1;
+ }
+ // Add groups after the last spill so replay also merges the final
+ // in-memory run.
+ for _ in 0..8 {
+ feed(0, batches);
+ batches += 1;
+ }
+ assert!(spill_count() >= spills);
+ assert_eq!(spill_count() == 0, spills == 0);
+ let mut first = streams.remove(0);
+ match finish {
+ Finish::Collect => {
+ senders[0].close_channel();
+ let mut output = collect(first).await?;
+ assert_eq!(pool.reserved(), held);
+ for sender in &senders[1..] {
+ sender.close_channel();
+ }
+ for stream in streams.drain(..) {
+ output.extend(collect(stream).await?);
+ }
+ let mut actual = BTreeMap::new();
+ for batch in output {
+ let a = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let b = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let sum = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let min = batch
+ .column(3)
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .unwrap();
+ for row in 0..batch.num_rows() {
+ assert!(
+ actual
+ .insert(
+ (a.value(row), b.value(row)),
+ (sum.value(row),
min.value(row).to_string())
+ )
+ .is_none()
+ );
+ }
+ }
+ assert_eq!(actual, expected);
+ assert_eq!(spill_count() == 0, spills == 0);
+ }
+ Finish::DropDuringReplay => {
+ senders[0].close_channel();
+ first.next().await.unwrap()?;
+ assert!(pool.reserved() > held);
+ drop(first);
+ assert_eq!(pool.reserved(), held);
+ }
+ Finish::InputError => {
+ senders[0]
+ .unbounded_send(datafusion_common::exec_err!(
+ "injected input failure"
+ ))
+ .unwrap();
+ let error = first.next().await.unwrap().unwrap_err();
+ assert!(error.to_string().contains("injected input failure"));
+ assert_eq!(pool.reserved(), held);
+ drop(first);
+ }
+ }
+ drop(streams);
+ assert_eq!(pool.reserved(), 0);
+ Ok(())
+ }
}
diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
index 53d515bf7f..48622127bc 100644
--- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
+++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt
@@ -198,6 +198,13 @@ FROM (
statement ok
SET datafusion.execution.target_partitions = 4
+# Four final partitions share the greedy pool. Limit merge buffers so a
+# partition that replays its spills can allocate while the other partitions
+# retain their aggregate state.
+# See https://github.com/apache/datafusion/issues/25423
+statement ok
+SET datafusion.runtime.max_spill_merge_fan_in = 2
+
query II
SELECT count(*), sum(total)
FROM (
@@ -230,6 +237,9 @@ FROM (
statement ok
RESET datafusion.runtime.memory_limit
+statement ok
+RESET datafusion.runtime.max_spill_merge_fan_in
+
statement ok
RESET datafusion.execution.batch_size
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
