ariel-miculas commented on code in PR #23606:
URL: https://github.com/apache/datafusion/pull/23606#discussion_r3715948732
##########
datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs:
##########
@@ -290,6 +302,226 @@ async fn
test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
Ok(())
}
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, true).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false, true).await
+}
+
+/// Intended to measure the maximum number of record batches held in memory by
+/// the SortPreservingMergeStream in a convoluted way by measuring the peak
+/// memory reservation. Relevant for merging spilled streams, where the
produced
+/// record batches suffer from the following issue:
+/// https://github.com/apache/arrow-rs/issues/6363
+///
+/// After an IPC roundtrip, all columns in a [`RecordBatch`] share a single
+/// parent buffer. It causes the memory reservation to be inflated, but the
+/// bigger issue is the increase in the peak allocated memory caused by
+/// prev_cursors in SortPreservingMergeExec. The increase is caused by the fact
+/// that the FieldCursor inside prev_cursors holds a reference for the entire
+/// Buffer allocated for the input record batch, preventing it from being
+/// dropped and thus increasing the number of concomitent input record batches
+/// living during the merging phase
+async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
+ round_robin: bool,
+ multi_column_sort: bool,
+) -> Result<()> {
+ let num_batches = 10usize;
+ let num_rows_per_batch = 100usize;
+ // payload is ~100x larger than the sort key (i32 = 4 bytes, string ≈ 400
bytes)
+ let large_string = "x".repeat(400);
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("sort_key", DataType::Int32, false),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+
+ // Unbounded env used only for spilling the input; the merge runs under its
+ // own pool below.
+ let spill_env = Arc::new(RuntimeEnvBuilder::new().build()?);
+
+ let mut partition_batches: Vec<Vec<RecordBatch>> = Vec::new();
+
+ for stream_idx in 0..2usize {
+ // Each stream covers a non-overlapping key range so both are
individually
+ // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000).
+ let batches: Vec<RecordBatch> = (0..num_batches)
+ .map(|b| {
+ // Interleave streams: stream 0 → even slots [0,200,400,...],
+ // stream 1 → odd slots [100,300,500,...] so the merge
+ // alternates between them on every batch.
+ let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32;
Review Comment:
added a test where the sort column is a constant value across all batches
--
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]
