fix: reduce peak memory usage when round robin tiebreaker is disabled by ariel-miculas · Pull Request #23606 · apache/datafusion

ariel-miculas · 2026-07-15T12:38:21Z

Which issue does this PR close?

First part of Increased memory usage due to round robin tie-breaker feature in SortPreservingMergeStream #23604

Rationale for this change

Don't pay the memory price when the tie breaking feature is disabled.

What changes are included in this PR?

Mostly tests to show the issue

Are these changes tested?

Yes

Are there any user-facing changes?

No

codecov-commenter · 2026-07-16T13:00:23Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.90%. Comparing base (db0c31b) to head (468f0da).
⚠️ Report is 18 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23606      +/-   ##
==========================================
+ Coverage   80.89%   80.90%   +0.01%     
==========================================
  Files        1102     1101       -1     
  Lines      376111   377044     +933     
  Branches   376111   377044     +933     
==========================================
+ Hits       304251   305052     +801     
- Misses      53753    53795      +42     
- Partials    18107    18197      +90

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

yonipeleg33

LGTM, nice work

yonipeleg33 · 2026-07-16T13:16:15Z

+                if self.enable_round_robin_tie_breaker {
+                    self.prev_cursors.as_mut().expect("prev_cursor should be set when round robin tie breaker is enabled")[stream_idx] = taken;
+                }


The expect() check seems a bit redundant, just match on the option directly

Suggested change

if self.enable_round_robin_tie_breaker {

self.prev_cursors.as_mut().expect("prev_cursor should be set when round robin tie breaker is enabled")[stream_idx] = taken;

}

if let Some(prev_cursors) = self.prev_cursors.as_mut() {

prev_cursors[stream_idx] = taken;

}

this would silently break if for some reason a None would end up in prev_cursors, I prefer to fail loudly

ariel-miculas · 2026-07-20T14:14:24Z

coverage report useless because of #23647

ariel-miculas · 2026-08-04T10:34:36Z

This PR is a prerequisite for:
#23619

@kumarUjjawal @rluvaton @alamb could you help me get this PR merged?

kumarUjjawal · 2026-08-04T11:05:47Z

Thank you @ariel-miculas for working on this. I will review it today.

kumarUjjawal

Overall this looks good. I noticed a related existing issue the round-robin setting can be lost when we rebuilds the merge operator, it can remove the memory benefit added here.

Can you look into this?

kumarUjjawal · 2026-08-04T13:49:21Z

+                // 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;


Could we add a small test where both input streams contain the same sort values across multiple batches?

added a test where the sort column is a constant value across all batches

ariel-miculas · 2026-08-04T14:42:52Z

Overall this looks good. I noticed a related existing issue the round-robin setting can be lost when we rebuilds the merge operator, it can remove the memory benefit added here.

This PR doesn't really bring any memory benefit (since round robin tie breaking is enabled by default), its purpose is to show the memory overhead caused by the exising design. My goal with this PR is to make it clear how the next PR reduces the memory overhead with round robin tie breaking enabled

kumarUjjawal · 2026-08-04T15:26:03Z

My goal with this PR is to make it clear how the next PR reduces the memory overhead with round robin tie breaking enabled

Thanks for clarification.

…mers

…test

ariel-miculas · 2026-08-04T20:59:04Z

I had to switch to PeakRecordingPool instead of TrackConsumersPool because the latter was reporting 0 bytes as the peak, the reason being that the memory consumer was being unregistered before we could get the actual value. This used to work previously, before the refactor to SortPreservingMergeStream in #23407 and #23976

Longer explanation:
SortPreservingMergeStream's reservation is unregistered from the memory pool before the test reads .metrics(), because
of how the stream is now implemented.

Two commits on main rewrote SortPreservingMergeStream from a manual poll_next state machine into an async generator:

096012e — "Simplifying SortPreservingMergeStream to use generators instead of state machine" (chore: Simplifying SortPreservingMergeStream to use generators instead of state machine #23407)
77b172e — "simplify SortPreservingMergeStream to be as textbook-like as possible" (chore: simplify SortPreservingMergeStream to be as textbook-like as possible #23702)

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:

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.

That drop cascades: MemoryReservation::drop → SharedRegistration::drop (once the last Arc clone goes
away) → MemoryPool::unregister(). TrackConsumersPool::unregister removes the consumer from its
tracked_consumers map entirely so by the time the test calls .metrics() after the loop, the consumer
is already gone and its recorded peak is lost.

ariel-miculas · 2026-08-05T10:13:04Z

I've fixed the existings issues and the failing tests, @kumarUjjawal could you please take another look?

kumarUjjawal · 2026-08-05T12:08:58Z

+    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


what about alternating ranges, stream 0 uses the even slots and stream 1 uses the odd slots?

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?

My bad, I should have been clear what I meant. I wasn’t asking for another test case. I was saying that the comment above says stream 0 covers [0,1000) and stream 1 covers [1000,2000). But the code does
stream 0 → [0, 100), [200, 300), ...
stream 1 → [100, 200), [300, 400), ...

I'll fix the comment in the next PR, since it'll touch these tests anyway.

kumarUjjawal

Thanks for all the follow ups.

I will leave this open for 1-2 days if anyone else has any opinions.

…apache#23606) ## Which issue does this PR close? - First part of apache#23604 ## Rationale for this change Don't pay the memory price when the tie breaking feature is disabled. ## What changes are included in this PR? Mostly tests to show the issue ## Are these changes tested? Yes ## Are there any user-facing changes? No

…tie-breaking purposes (apache#23619) ## Which issue does this PR close? - Closes apache#23606. ## Rationale for this change See the linked issue. Note that this PR also contains the changes in apache#23606, so this will have to be rebased ## What changes are included in this PR? Store only the last row in prev_cursors instead of keeping the entire cursor ## Are these changes tested? Added a test to show the improvement ## Are there any user-facing changes? No

github-actions Bot added core Core DataFusion crate physical-plan Changes to the physical-plan crate labels Jul 15, 2026

mzabaluev reviewed Jul 15, 2026

View reviewed changes

Comment thread datafusion/physical-plan/src/sorts/merge.rs Outdated

Udaytaneja mentioned this pull request Jul 15, 2026

Increased memory usage due to round robin tie-breaker feature in SortPreservingMergeStream #23604

Open

ariel-miculas mentioned this pull request Jul 15, 2026

feat: Store only the last row in the previous cursor for round robin tie-breaking purposes #23619

Merged

yonipeleg33 reviewed Jul 16, 2026

View reviewed changes

Comment thread datafusion/physical-plan/src/sorts/merge.rs Outdated

yonipeleg33 approved these changes Jul 16, 2026

View reviewed changes

ariel-miculas force-pushed the reduce-memory-usage-when-rr-disabled branch from 95a0830 to d018aa3 Compare July 20, 2026 13:34

kumarUjjawal self-requested a review August 4, 2026 11:06

kumarUjjawal reviewed Aug 4, 2026

View reviewed changes

ariel-miculas added 8 commits August 4, 2026 19:36

fix: reduce peak memory usage when round robin tiebreaker is disabled

3109bdb

chore: remove useless code

d0518a7

refactor: simplify the expression

4ec6618

feat: add test for multi-column sort

c05609d

refactor: make prev_cursors optional

ed55e65

refactor: use prev_cursors as source of truth for tie-breaker config

72eb210

feat: add test case where all sort values are equal across all batches

80bec04

fix: use PeakRecordingPool to avoid missing unregistered memory consu…

4bf99ec

…mers

ariel-miculas force-pushed the reduce-memory-usage-when-rr-disabled branch from 82e7472 to 4bf99ec Compare August 4, 2026 19:42

ariel-miculas added 2 commits August 4, 2026 23:09

refactor: move PeakRecordingPool to datafusion-execution

fa82f5a

refactor: avoid circular dependency

4aaf8dd

github-actions Bot added the execution Related to the execution crate label Aug 4, 2026

ariel-miculas added 2 commits August 4, 2026 23:34

refactor: remove unused dependencies

a5f0136

refactor: remove unused dependency and regenerate cargo.lock

c1d3bf5

refactor: enable arrow_buffer_pool feature for the arrow memory pool …

79bda67

…test

ariel-miculas added 2 commits August 5, 2026 00:02

refactor: cannot have circular dependencies so gate via cfg instead

4e5ba44

fix: rust doctests

468f0da

kumarUjjawal reviewed Aug 5, 2026

View reviewed changes

kumarUjjawal approved these changes Aug 5, 2026

View reviewed changes

kumarUjjawal added this pull request to the merge queue Aug 7, 2026

Merged via the queue into apache:main with commit 0646a31 Aug 7, 2026
41 checks passed

Uh oh!

Conversation

ariel-miculas commented Jul 15, 2026 • edited by alamb Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

yonipeleg33 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ariel-miculas commented Jul 20, 2026

Uh oh!

ariel-miculas commented Aug 4, 2026

Uh oh!

kumarUjjawal commented Aug 4, 2026

Uh oh!

kumarUjjawal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ariel-miculas commented Aug 4, 2026

Uh oh!

kumarUjjawal commented Aug 4, 2026

Uh oh!

ariel-miculas commented Aug 4, 2026

Uh oh!

ariel-miculas commented Aug 5, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kumarUjjawal left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ariel-miculas commented Jul 15, 2026 •

edited by alamb

Loading

codecov-commenter commented Jul 16, 2026 •

edited

Loading