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

[ariel-miculas]

Which issue does this PR close?

• Closes fix: reduce peak memory usage when round robin tiebreaker is disabled #23606.

Rationale for this change

See the linked issue.
Note that this PR also contains the changes in #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

[EmilyMatt]

I feel like this should return an Option, in case values() is empty

[ariel-miculas]

values() shouldn't be empty, there are some asserts in the code:

        assert!(rows.num_rows() > 0);
        assert!(array.len() > 0, "Empty array passed to FieldCursor");

I don't imagine a Cursor with empty values being useful.

impl RowValues {
    /// Create a new [`RowValues`] from `rows` and a `reservation`
    /// that tracks its memory. There must be at least one row
    ///
    /// Panics if the reservation is not for exactly `rows.size()`
    /// bytes or if `rows` is empty.
    pub fn new(rows: Arc<Rows>, reservation: MemoryReservation) -> Self {
        assert_eq!(
            rows.size(),
            reservation.size(),
            "memory reservation mismatch"
        );
        assert!(rows.num_rows() > 0);
        Self {
            rows,
            _reservation: reservation,
        }
    }
}

    pub fn new<A: CursorArray<Values = T>>(
        options: SortOptions,
        array: &A,
        reservation: MemoryReservation,
    ) -> Self {
        assert!(array.len() > 0, "Empty array passed to FieldCursor");
        let null_threshold = match options.nulls_first {
            true => array.null_count(),
            false => array.len() - array.null_count(),
        };

[ariel-miculas]

The previous code also assumed values is never empty:

    fn is_eq_to_prev_row_in_prev_batch(&self, other: &Self) -> bool {
        assert_eq!(self.offset, 0);
        T::eq(
            &self.values,
            self.offset,
            &other.values,
            other.values.len() - 1,
        )
    }

[ariel-miculas]

@Dandandan wondering whether we can change ReusableRows to only hold one slot after this change, since the previous cursor will no longer need to be kept alive for the round robin tie-breaking feature.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 81.48148% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.03%. Comparing base (0646a31) to head (06bde7b).
⚠️ Report is 75 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/sorts/merge.rs	11.11%	8 Missing ⚠️
datafusion/physical-plan/src/sorts/cursor.rs	95.55%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23619      +/-   ##
==========================================
+ Coverage   81.02%   81.03%   +0.01%     
==========================================
  Files        1105     1106       +1     
  Lines      380669   381026     +357     
  Branches   380669   381026     +357     
==========================================
+ Hits       308446   308778     +332     
- Misses      53994    53997       +3     
- Partials    18229    18251      +22     

☔ 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.

[ariel-miculas]

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 #23407 and #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.

[ariel-miculas]

I was wondering where this comment went, I meant to post it in #23606

[ariel-miculas]

@kumarUjjawal this is the follow-up to #23606, could you please take a look?

[kumarUjjawal]

Thanks you @ariel-miculas

I have left few comments. Let me know what you think.

[kumarUjjawal]

we should keep at least the single-column no-round-robin case added by #23606? Removing all three disabled-mode cases drops the only direct peak-memory regression coverage for that fix

[ariel-miculas]

We could keep all three, I removed them since there's no longer any difference in the reservation between round robin tie-breaker enabled / disabled

[kumarUjjawal]

Sure! Thanks

[ariel-miculas]

I've added back the no round robin test cases

[ariel-miculas]

Thanks for the review! As soon as this is merged I'll rebase #23802 on top of main and the tests should pass.

[ariel-miculas]

@kumarUjjawal could you please merge this PR?

[kumarUjjawal]

Thank you @ariel-miculas