rluvaton commented on code in PR #23286:
URL: https://github.com/apache/datafusion/pull/23286#discussion_r3570668105
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -891,6 +969,63 @@ mod tests {
Ok(())
}
+
+ /// Same as [`respill_halves_the_merge_output_batch_size`], but under a
budget tight
+ /// enough that *both* runs must be re-spilled before the merge fits - the
scenario
+ /// where the batch-size reduction could compound. Because the reduction
is tracked
+ /// per-run (each run capped at half) rather than by halving the global
batch size on
+ /// every split, the merged output is emitted in 4096-row batches - half,
not a
+ /// quarter. A global-halving implementation would have halved once per
re-spill and
+ /// emitted 2048-row batches.
+ #[tokio::test]
+ async fn respilling_two_skewed_runs_halves_the_output_without_compounding()
+ -> Result<()> {
+ let env = Arc::new(RuntimeEnv::default());
+ let schema = test_schema();
+ let spill_manager = build_spill_manager(&env, &schema);
+
+ let n: i64 = 16384;
+ let f0 = make_sorted_spill_file(&spill_manager, &schema,
(0..n).collect());
+ let f1 = make_sorted_spill_file(&spill_manager, &schema,
(0..n).collect());
+ let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory);
+
+ // 2.5*m is tight enough that even after halving one run the two still
don't
+ // fit, so *both* runs are re-spilled once before the merge succeeds.
(3.5*m,
+ // as in the single-split test, would let the pair fit after one
split.) This
+ // is exactly the scenario where a compounding, global-halving
implementation
+ // would drive the output batch size down to a quarter.
+ let initial_batch_size = 8192;
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 5 /
2));
+
+ let builder = build_merge_builder(
+ spill_manager,
+ Arc::clone(&schema),
+ vec![f0, f1],
+ &pool,
+ initial_batch_size,
+ );
+ let stream = builder.create_spillable_merge_stream();
+ let batches: Vec<RecordBatch> = stream.try_collect().await?;
+
+ // All rows are still present.
+ let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(total_rows, (2 * n) as usize);
+
+ // Each run was re-spilled once, so each is capped at half the
original batch
+ // size and the merge caps its output at that half - NOT a quarter. A
global
+ // halving-per-split implementation would have emitted 2048-row
batches here.
+ let expected_batch_size = initial_batch_size / 2;
+ let max_batch_rows = batches.iter().map(|b|
b.num_rows()).max().unwrap_or(0);
+ assert_eq!(
+ max_batch_rows, expected_batch_size,
+ "two re-spills must halve (not quarter) the output: expected \
+ {expected_batch_size}-row batches, got a largest batch of \
+ {max_batch_rows} rows"
+ );
+
+ Ok(())
+ }
Review Comment:
I would prefer more integration test like and not access private functions
but I see that all the other tests here are like that so we can do that in a
separate pr
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
