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 57f861942d test: make `nested_key_spill_keeps_groups_unique`
independent of scheduling (#25156)
57f861942d is described below
commit 57f861942da849d338b660a029a3d58bcbcf20a0
Author: Raz Luvaton <[email protected]>
AuthorDate: Thu Sep 10 19:33:57 2026 +0000
test: make `nested_key_spill_keeps_groups_unique` independent of scheduling
(#25156)
## Which issue does this PR close?
N/A
## Rationale for this change
The test is flaky for non legacy with limited memory
## What changes are included in this PR?
The migrated final stream replays the merged spill through an ordered
table that accounts for every group of a merged batch before emitting
the completed ones. With 16 groups of ~675 KiB `count(distinct)` state
each and 64 row batches over 4 hash partitioned finals, that replay
needed up to ~2.7 MiB, which under a `FairSpillPool` only fit once the
other partitions had finished and released their share. On CI 7
consumers were still registered and the run failed with
`ResourcesExhausted`.
The legacy run keeps the original shape, which is the only one that
reproduces the bug it guards against, and its replay stays well within
its worst case share. The migrated run now uses a single final stream
(no hash repartition) with 8 row batches, so it is the only spillable
consumer left at replay time and a merged batch holds one or two groups.
## What is the testing strategy for this PR?
N/A
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
No
Co-authored-by: Claude Fable 5.1 <[email protected]>
---
datafusion/core/tests/memory_limit/mod.rs | 62 ++++++++++++++++++++++++-------
1 file changed, 48 insertions(+), 14 deletions(-)
diff --git a/datafusion/core/tests/memory_limit/mod.rs
b/datafusion/core/tests/memory_limit/mod.rs
index 1e86c1fb7f..2cda54d678 100644
--- a/datafusion/core/tests/memory_limit/mod.rs
+++ b/datafusion/core/tests/memory_limit/mod.rs
@@ -233,10 +233,46 @@ async fn nested_key_spill_keeps_groups_unique() {
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 that the migrated final stream can hold one merged
- /// batch under a `FairSpillPool` shared by four partitions.
- const NESTED_KEY_MEMORY_LIMIT: usize = 8 * 1024 * 1024;
+ /// Small enough that the final stage must spill its `count(distinct)`
state
+ /// under the plan shape each stream implementation runs with, see
+ /// `nested_key_session_config` for why they differ.
+ const LEGACY_MEMORY_LIMIT: usize = 8 * 1024 * 1024;
+ const MIGRATED_MEMORY_LIMIT: usize = 2 * 1024 * 1024;
+
+ /// A `FairSpillPool` splits the limit evenly between the spillable
consumers
+ /// registered at the time of each allocation, so what a stream may hold
+ /// depends on which other streams are still alive. Each implementation
gets
+ /// a plan shape under which that does not matter.
+ ///
+ /// The legacy bug needs several new groups per merged batch: 64 row
batches
+ /// over 4 hash partitioned final streams reproduce it, and the legacy
replay
+ /// of the merged spill holds at most ~340 KiB, within the ~680 KiB share
it
+ /// gets even while all 12 consumers (partials, repartitions and finals)
are
+ /// registered.
+ ///
+ /// The migrated final stream replays the merged spill through an ordered
+ /// table that accounts for every group of a merged batch before emitting
+ /// the completed ones, and one group's `count(distinct)` state here is
+ /// ~675 KiB. With 64 row batches over 4 partitions that reached ~2.7 MiB,
+ /// which only fit once the other partitions had finished and released
+ /// their share, so the run depended on scheduling. It therefore runs a
+ /// single final stream (no hash repartition), which is the only spillable
+ /// consumer left once the partials are done, with 8 row batches so a
+ /// merged batch holds one or two groups.
+ fn nested_key_session_config(legacy: bool) -> SessionConfig {
+ let config = SessionConfig::new()
+ .with_target_partitions(4)
+ .set_bool("datafusion.execution.enable_migration_aggregate",
!legacy);
+ if legacy {
+ // small batches: the merged spill stream arrives in many batches
and
+ // groups span batch boundaries
+ config.with_batch_size(64)
+ } else {
+ config
+ .with_batch_size(8)
+ .set_bool("datafusion.optimizer.repartition_aggregations",
false)
+ }
+ }
fn nested_key_struct_fields() -> Fields {
Fields::from(vec![
@@ -297,14 +333,10 @@ async fn nested_key_spill_keeps_groups_unique() {
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",
!legacy);
- let ctx =
- SessionContext::new_with_config_rt(config,
runtime.build_arc().unwrap());
+ let ctx = SessionContext::new_with_config_rt(
+ nested_key_session_config(legacy),
+ runtime.build_arc().unwrap(),
+ );
ctx.register_table("t", Arc::new(nested_key_table()))
.unwrap();
@@ -338,9 +370,11 @@ async fn nested_key_spill_keeps_groups_unique() {
"unbounded, legacy=true"
);
- for legacy in [true, false] {
+ for (legacy, memory_limit) in
+ [(true, LEGACY_MEMORY_LIMIT), (false, MIGRATED_MEMORY_LIMIT)]
+ {
assert_eq!(
- run_nested_key_query(Some(NESTED_KEY_MEMORY_LIMIT), legacy).await,
+ run_nested_key_query(Some(memory_limit), legacy).await,
expected,
"spilling, legacy={legacy}"
);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
