ariel-miculas commented on code in PR #23619:
URL: https://github.com/apache/datafusion/pull/23619#discussion_r3714795493
##########
datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs:
##########
@@ -290,6 +302,220 @@ 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;
+ let sort_col: Int32Array =
+ (base..base + num_rows_per_batch as i32).collect();
+ let payload_col: StringArray =
+ std::iter::repeat_n(large_string.as_str(),
num_rows_per_batch)
+ .map(Some)
+ .collect();
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(sort_col), Arc::new(payload_col)],
+ )
+ .unwrap()
+ })
+ .collect();
+
+ // Spill to disk then read back: each RecordBatch is now IPC-backed,
+ // meaning all columns share a single parent buffer. As a result,
+ // get_buffer_memory_size() on the sort_key column returns the full
+ // parent-buffer capacity (≈ batch size of both columns combined)
rather
+ // than just the key data (num_rows * 4 bytes).
+ let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+ let manager =
+ SpillManager::new(Arc::clone(&spill_env), metrics,
Arc::clone(&schema));
+ let spill_file = manager
+ .spill_record_batch_and_finish(&batches, "stream")?
+ .expect("non-empty input should produce a spill file");
+
+ let mut stream = manager.read_spill_as_stream(spill_file, None)?;
+ let mut ipc_batches: Vec<RecordBatch> = Vec::new();
+ while let Some(batch) = stream.next().await {
+ ipc_batches.push(batch?);
+ }
+ partition_batches.push(ipc_batches);
+ }
+
+ let ipc_batch_size =
get_record_batch_memory_size(&partition_batches[0][0]);
+
+ // Build a 2-partition plan from the IPC-recovered batches.
+ let input =
+ MemorySourceConfig::try_new_exec(&partition_batches,
Arc::clone(&schema), None)?;
+
+ let sort_key_expr = PhysicalSortExpr {
+ expr: col("sort_key", &schema)?,
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ };
+ // `payload` has the same value in every row, so adding it as a secondary
+ // sort key doesn't change the resulting order — it only forces the merge
+ // onto the row-oriented (`RowValues`/`RowCursorStream`) comparison path
+ // used whenever more than one sort expression is present.
+ let mut sort_exprs = vec![sort_key_expr];
+ if multi_column_sort {
+ sort_exprs.push(PhysicalSortExpr {
+ expr: col("payload", &schema)?,
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ });
+ }
+
+ // When sorting by more than one column, the merge switches to the
+ // row-oriented `RowValues`/`RowCursorStream` path
+ //
+ // `RowCursorStream` also tracks one *shared* (not per-partition)
+ // reservation sized to `converter.size()` (`stream.rs`:
+ // `self.reservation.try_resize(self.converter.size())`) — the
+ // `RowConverter`'s own fixed internal state, separate from the `Rows`
+ // it produces per batch.
+ let (row_batch_size, converter_size) = if multi_column_sort {
+ let sort_fields = sort_exprs
+ .iter()
+ .map(|s| {
+ let data_type = s.expr.data_type(&schema)?;
+ Ok(SortField::new_with_options(data_type, s.options))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let converter = RowConverter::new(sort_fields)?;
+ let cols = sort_exprs
+ .iter()
+ .map(|s| {
+ s.expr
+ .evaluate(&partition_batches[0][0])?
+ .into_array(partition_batches[0][0].num_rows())
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let rows = converter.convert_columns(&cols)?;
+ (rows.size(), converter.size())
+ } else {
+ (0, 0)
+ };
+
+ let merge = Arc::new(
+ SortPreservingMergeExec::new(LexOrdering::new(sort_exprs).unwrap(),
input)
+ .with_round_robin_repartition(round_robin),
+ );
+
+ // TrackConsumersPool wraps an unbounded pool so the merge never OOMs;
+ // we keep the typed Arc to call .metrics() after the run.
+ let tracking_pool = Arc::new(TrackConsumersPool::new(
+ UnboundedMemoryPool::default(),
+ NonZeroUsize::new(10).unwrap(),
+ ));
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&tracking_pool) as Arc<dyn MemoryPool>)
+ .build()?;
+ let task_ctx = Arc::new(
+ TaskContext::default()
+
.with_session_config(SessionConfig::new().with_batch_size(num_rows_per_batch))
+ .with_runtime(Arc::new(runtime)),
+ );
+
+ let mut output = merge.execute(0, task_ctx)?;
+ let mut total_rows = 0usize;
+ while let Some(batch) = output.next().await {
+ total_rows += batch?.num_rows();
+ }
+ assert_eq!(total_rows, 2 * num_batches * num_rows_per_batch);
+
+ let mut metrics = tracking_pool.metrics();
+ metrics.sort_by_key(|m| std::cmp::Reverse(m.peak));
+ let peak_bytes: usize = metrics.iter().map(|m| m.peak).sum();
+
+ // in the single column case, the cursor takes up an ipc_batch_size worth
of memory due to the
+ // IPC roundtrip issue
+ // for the multi-column case, we've calculated row_batch_size above
+ let cursor_unit = if multi_column_sort {
+ row_batch_size
+ } else {
+ ipc_batch_size
+ };
+
+ // BatchBuilder needs to hold 3 Record batches simultaneously to merge two
+ // streams (because a stream can cross a record batch boundary)
+ // there is also one cursor needed per stream
+ let max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size;
+
+ assert!(
Review Comment:
tldr: This no longer works after the refactor of SortPreservingMergeStream
because peak_bytes reports 0.
SortPreservingMergeStream's reservation is unregistered from the memory pool
before the test reads .metrics(), because of how the stream is now (meaning
after https://github.com/apache/datafusion/pull/23407 and
https://github.com/apache/datafusion/pull/23702) implemented.
Before: SortPreservingMergeStream implemented Stream directly (`fn
poll_next(self: Pin<&mut Self>, ...)`). The struct — including its
MemoryReservation — was boxed and returned as the `SendableRecordBatchStream`
itself. It stayed alive for exactly as long as the caller held the stream, and
was only dropped (triggering unregister()) when the caller dropped it — in the
test, after tracking_pool.metrics() was read. Peak was correctly captured.
After: the implementation moved to (merge.rs:225-226):
```
fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
async_try_stream(|mut emitter| async move { /* ...self... */ })
}
```
self (and everything it owns, including the reservation) is now captured
inside the generator's async block. Once that block runs to completion — which
happens on the final poll_next() call that returns None, i.e. during the test's
draining loop, not when the caller later drops the stream object — self is
dropped right there.
--
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]
