Release consumed sort input batches after partial output

[Weijun-H]

Which issue does this PR close?

• Closes Release consumed sort input batches after partial output #25405.

Rationale for this change

External sort can keep already-consumed input batches alive after producing a partial output batch. Releasing those batches earlier reduces how long memory stays reserved for sorts that emit output in multiple chunks.

What changes are included in this PR?

This PR updates BatchBuilder cleanup after output emission to retain only batches that are still referenced by pending row indices or live cursors. It remaps remaining indices and cursors after dropping consumed batches.

The change is limited to earlier release of consumed input batches. It does not add sort byte-targeting, spill admission changes, or new configuration.

The tradeoff is that partial-output cleanup now marks live batches and remaps remaining row indices/cursors so consumed batches can be released earlier. I do not have an independent performance measurement for this first change yet.

What is the testing strategy for this PR?

Added a BatchBuilder test that emits a partial output batch and checks both directions:

• the consumed input batch is released and its memory reservation is returned;
• batches still needed by pending indices or a live cursor remain available and produce the expected later output.

Also ran:

• cargo fmt --check
• cargo test -p datafusion-physical-plan sorts::builder --lib
• git diff --check HEAD^ HEAD

Are there any user-facing changes?

No API or configuration changes. This only releases internal sort input batch memory earlier.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.28%. Comparing base (e5469e1) to head (b562850).

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #25406   +/-   ##
=======================================
  Coverage   82.28%   82.28%           
=======================================
  Files        1137     1137           
  Lines      430211   430309   +98     
  Branches   430211   430309   +98     
=======================================
+ Hits       354003   354091   +88     
- Misses      54784    54789    +5     
- Partials    21424    21429    +5     

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

[Copilot]

🟡 Changes recommended

The new cleanup adds avoidable allocations and cursor scanning to every normal output batch.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Balanced

[jayzhan211]

Thanks @Weijun-H , here is a suggestion:

Both retain paths drop a cursor's batch but leave cursor.batch_idx pointing at the old slot, so after compaction a stale cursor indexes some other stream's batch. Before this PR the cursor batch was always retained, so batch_idx was always valid. Now a push_row on an exhausted stream would silently emit a row from the wrong stream, and retain_live_batches needs the batch_stream_idx == stream_idx check and a second retain_cursor vec only to detect that case.

A usize::MAX sentinel removes both, lets the remap loop produce the sentinel for free, and allows a debug_assert in push_row.

-#[derive(Debug, Copy, Clone, Default)]
+#[derive(Debug, Copy, Clone)]
 struct BatchCursor {
     batch_idx: usize,
     row_idx: usize,
 }
+
+impl BatchCursor {
+    /// A cursor whose batch has been released. `push_row` must not be called
+    /// for the stream until `push_batch` installs a new cursor.
+    const RELEASED: Self = Self {
+        batch_idx: usize::MAX,
+        row_idx: 0,
+    };
+}

-            cursors: vec![BatchCursor::default(); stream_count],
+            cursors: vec![BatchCursor::RELEASED; stream_count],

     pub fn push_row(&mut self, stream_idx: usize) {
+        debug_assert!(
+            self.batches
+                .get(self.cursors[stream_idx].batch_idx)
+                .is_some_and(|(_, b)| self.cursors[stream_idx].row_idx < b.num_rows()),
+            "push_row on stream {stream_idx} with no live batch"
+        );
         let cursor = &mut self.cursors[stream_idx];

             let stream_cursor = &mut self.cursors[*stream_idx];
-            let retain = stream_cursor.batch_idx == batch_idx
-                && stream_cursor.row_idx < batch.num_rows();
+            let is_cursor_batch = stream_cursor.batch_idx == batch_idx;
+            let retain = is_cursor_batch && stream_cursor.row_idx < batch.num_rows();
             batch_idx += 1;

             if retain {
                 stream_cursor.batch_idx = retained;
                 retained += 1;
             } else {
+                if is_cursor_batch {
+                    *stream_cursor = BatchCursor::RELEASED;
+                }
                 self.batches_mem_used -= get_record_batch_memory_size(batch);
             }

-        let mut retain_cursor = vec![false; self.cursors.len()];
-        for (stream_idx, cursor) in self.cursors.iter().enumerate() {
-            if self.batches.get(cursor.batch_idx).is_some_and(
-                |(batch_stream_idx, batch)| {
-                    *batch_stream_idx == stream_idx && cursor.row_idx < batch.num_rows()
-                },
-            ) {
-                retain_batch[cursor.batch_idx] = true;
-                retain_cursor[stream_idx] = true;
-            }
-        }
+        for cursor in &self.cursors {
+            if let Some((_, batch)) = self.batches.get(cursor.batch_idx)
+                && cursor.row_idx < batch.num_rows()
+            {
+                retain_batch[cursor.batch_idx] = true;
+            }
+        }

-        for (stream_idx, cursor) in self.cursors.iter_mut().enumerate() {
-            if retain_cursor[stream_idx] {
-                cursor.batch_idx = remap[cursor.batch_idx];
-            }
-        }
+        for cursor in &mut self.cursors {
+            if let Some(new_idx) = remap.get(cursor.batch_idx) {
+                // `usize::MAX` for a released batch, i.e. `BatchCursor::RELEASED`.
+                cursor.batch_idx = *new_idx;
+            }
+        }