ariel-miculas commented on code in PR #23606:
URL: https://github.com/apache/datafusion/pull/23606#discussion_r3720506933
##########
datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs:
##########
@@ -290,6 +301,250 @@ 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,
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,
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,
false).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,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, false,
true).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false, 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,
+ tied_values: 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
Review Comment:
what's the rationale for adding all these edge cases? is it trying to figure
out whether SortPreservingMerge over-reserves memory with certain data inputs?
--
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]
