This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25156-1ec9ede5da5c67cbc2d3e0110a7ddbc2a47796a6 in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
