You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/15 18:05:50 UTC
(datafusion) branch main updated: test: bound merge fan-in in ordered aggregate spill tests (#25252)
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 22651d24cc test: bound merge fan-in in ordered aggregate spill tests (#25252)
22651d24cc is described below
commit 22651d24cc8196f3206e09a36a437eda9e766b87
Author: Kumar Ujjawal <uj...@gmail.com>
AuthorDate: Tue Sep 15 17:56:37 2026 +0000
test: bound merge fan-in in ordered aggregate spill tests (#25252)
## Which issue does this PR close?
Closes #25047.
## Rationale for this change
`ordered_aggregate_spill.slt` can fail when spill-merge buffers leave
too little memory for aggregate replay while another partition retains
aggregate state in the shared greedy pool. Limit merge fan-in to two to
reduce this contention while preserving the 600 KiB memory limit, two
partitions, and existing query results.
## What changes are included in this PR?
- Set and reset the SQL test's spill merge fan-in.
- Add controlled two-partition tests for grouped results, spilling, and
reservation cleanup on completion, input error, and cancellation.
- Check spilled rows without depending on the formatted byte unit, and
remove exact spill-count claims from comments.
## What is the testing strategy for this PR?
Local validation:
- 320 SQL test executions (16 concurrent copies per run) and 50
repetitions of the focused tests.
- Full extended workspace tests, including all 518 SQL test files.
- Workspace Clippy with all targets and features, formatting, and
submission checks.
## Are there any user-facing changes?
No. This changes tests only; production execution and defaults are
unchanged.
---
.../src/aggregates/ordered_final_stream.rs | 283 +++++++++++++++++++++
.../test_files/ordered_aggregate_spill.slt | 16 +-
2 files changed, 295 insertions(+), 4 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
index 8823dcb32f..ac6f2491b4 100644
--- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
@@ -900,3 +900,286 @@ impl RecordBatchStream for OrderedFinalAggregateStream {
Arc::clone(&self.schema)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::ExecutionPlan;
+ use crate::aggregates::PhysicalGroupBy;
+ use crate::common::collect;
+ use crate::stream::RecordBatchStreamAdapter;
+ use crate::test::TestMemoryExec;
+ use arrow::array::{Int64Array, StringViewArray};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_execution::config::SessionConfig;
+ use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool};
+ use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+ 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::FutureExt;
+ use futures::channel::mpsc;
+ use std::collections::BTreeMap;
+
+ #[derive(Clone, Copy)]
+ enum Finish {
+ Collect,
+ DropDuringMerge,
+ DropDuringReplay,
+ InputError,
+ }
+
+ #[tokio::test]
+ async fn spill_replay_with_another_ordered_partition() -> Result<()> {
+ for input_batches in [28, 36, 55, 63] {
+ run_shared_pool_case(input_batches, 600 * 1024, Finish::Collect).await?;
+ }
+ // The same input also produces the reference results without spilling.
+ run_shared_pool_case(63, 10 * 1024 * 1024, Finish::Collect).await
+ }
+
+ #[tokio::test]
+ async fn spill_replay_releases_memory_on_drop() -> Result<()> {
+ run_shared_pool_case(36, 600 * 1024, Finish::DropDuringMerge).await?;
+ run_shared_pool_case(36, 600 * 1024, Finish::DropDuringReplay).await
+ }
+
+ #[tokio::test]
+ async fn ordered_spill_releases_memory_on_input_error() -> Result<()> {
+ run_shared_pool_case(36, 600 * 1024, Finish::InputError).await
+ }
+
+ async fn run_shared_pool_case(
+ input_batches: i64,
+ limit: usize,
+ finish: Finish,
+ ) -> Result<()> {
+ 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 ordering =
+ LexOrdering::new([PhysicalSortExpr::new_default(col("a", &partial_schema)?)])
+ .unwrap();
+ let input = TestMemoryExec::try_new(
+ &[vec![], vec![]],
+ Arc::clone(&partial_schema),
+ None,
+ )?
+ .try_with_sort_information(vec![ordering])?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::FinalPartitioned,
+ groups.as_final(),
+ expressions,
+ vec![None; 2],
+ Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
+ schema,
+ )?;
+ assert_eq!(
+ aggregate.input_order_mode(),
+ &InputOrderMode::PartiallySorted(vec![0])
+ );
+
+ 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).set_bool(
+ "datafusion.execution.enable_migration_aggregate",
+ true,
+ ),
+ )
+ .with_runtime(
+ RuntimeEnvBuilder::new()
+ .with_max_spill_merge_fan_in(2)
+ .with_memory_pool(Arc::clone(&pool))
+ .build_arc()?,
+ ),
+ );
+ let mut streams = vec![];
+ let mut senders = vec![];
+ for partition in 0..2 {
+ let (sender, receiver) = mpsc::unbounded();
+ let input = Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&partial_schema),
+ receiver,
+ ));
+ let stream = OrderedFinalAggregateStream::new_with_input(
+ &aggregate,
+ &context,
+ partition,
+ input,
+ aggregate.input_order_mode(),
+ )?;
+ 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()
+ };
+ // Keep partition 1's incomplete ordered run live while partition 0 spills
+ // and replays. Channel inputs return Pending after each supplied batch,
+ // making this interleaving independent of task scheduling.
+ for batch in 0..55 {
+ senders[1]
+ .unbounded_send(Ok(make_batch(1, batch * 128)))
+ .unwrap();
+ assert!(streams[1].next().now_or_never().is_none());
+ }
+ let held = streams[1].reservation.size();
+ assert!(held > 500 * 1024);
+ for batch in 0..input_batches {
+ // Repeated keys cross spill runs, so replay must merge their sums.
+ senders[0]
+ .unbounded_send(Ok(make_batch(0, batch * 128)))
+ .unwrap();
+ assert!(streams[0].next().now_or_never().is_none());
+ }
+ if limit == 600 * 1024 {
+ assert!(aggregate.metrics().unwrap().spill_count().unwrap() > 0);
+ }
+ let mut first = streams.remove(0);
+ match finish {
+ Finish::Collect => {
+ senders[0].close_channel();
+ let mut output = collect(Box::pin(first)).await?;
+ assert_eq!(pool.reserved(), held);
+ senders[1].close_channel();
+ output.extend(collect(Box::pin(streams.remove(0))).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!(
+ aggregate.metrics().unwrap().spill_count().unwrap() > 0,
+ limit == 600 * 1024
+ );
+ }
+ Finish::DropDuringMerge => {
+ senders[0].close_channel();
+ let _ = first.next().now_or_never();
+ assert!(matches!(
+ first.state.as_ref(),
+ Some(OrderedFinalAggregateState::MergingSpills { .. })
+ ));
+ drop(first);
+ }
+ 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/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt
index 67fa20101d..548dfc69b1 100644
--- a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt
+++ b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt
@@ -35,6 +35,11 @@ SET datafusion.optimizer.prefer_existing_sort = true
statement ok
SET datafusion.execution.enable_migration_aggregate = true
+# Limit merge buffers so replay can allocate while another partition retains
+# its aggregate state in the shared greedy pool.
+statement ok
+SET datafusion.runtime.max_spill_merge_fan_in = 2
+
statement ok
SET datafusion.runtime.memory_limit = '1M'
@@ -108,7 +113,7 @@ GROUP BY round(v1, -4), v1 % 5000
----
45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12
-# Round 2: The same query spills five times with a 600 KB limit.
+# Round 2: The same query spills with a 600 KB limit.
statement ok
SET datafusion.runtime.memory_limit = '600K'
@@ -132,7 +137,7 @@ GROUP BY round(v1, -4), v1 % 5000
----
45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12
-# Round 3: The same query spills six times with a 500 KB limit.
+# Round 3: The same query spills with a 500 KB limit.
statement ok
SET datafusion.runtime.memory_limit = '500K'
@@ -162,7 +167,7 @@ GROUP BY round(v1, -4), v1 % 5000
statement ok
SET datafusion.runtime.memory_limit = '600K'
-# Ensures final aggregate has spill_count > 0
+# Check spilled rows because extra merge passes can change the byte unit.
query TT
EXPLAIN ANALYZE
SELECT round(v1, -4), v1 % 5000,
@@ -171,7 +176,7 @@ FROM generate_series(20000) AS t1(v1)
GROUP BY round(v1, -4), v1 % 5000
----
Plan with Metrics
-01)AggregateExec: mode=FinalPartitioned,<slt:ignore>aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[<slt:ignore>spilled_bytes=<slt:ignore>KB,<slt:ignore>]
+01)AggregateExec: mode=FinalPartitioned,<slt:ignore>aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[<slt:ignore>spilled_rows=<slt:ignore> K,<slt:ignore>]
02)--RepartitionExec:<slt:ignore>input_partitions=1, maintains_sort_order=true<slt:ignore>
03)----AggregateExec: mode=Partial,<slt:ignore>aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
<slt:ignore>
@@ -249,5 +254,8 @@ RESET datafusion.execution.batch_size
statement ok
SET datafusion.execution.target_partitions = 4
+statement ok
+RESET datafusion.runtime.max_spill_merge_fan_in
+
statement ok
RESET datafusion.catalog.create_default_catalog_and_schema
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org