From 4a16706ab4bb4ac89b55de5c11b9c378fc7e1275 Mon Sep 17 00:00:00 2001 From: Alex Huang Date: Thu, 17 Sep 2026 17:23:44 +0800 Subject: [PATCH 1/3] Release consumed sort batches after partial emit --- datafusion/physical-plan/src/sorts/builder.rs | 141 ++++++++++++++---- 1 file changed, 113 insertions(+), 28 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..71ed2667c5f43 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -157,34 +157,7 @@ impl BatchBuilder { // Remove consumed indices, keeping any remaining for the next call. self.indices.drain(..rows_to_emit); - // Only clean up fully-consumed batches when all indices are drained, - // because remaining indices may still reference earlier batches. - // In the overflow/partial-emit case this may retain some extra memory - // across a few drain polls, but avoids costly index scanning on the - // hot path. The retention is bounded and short-lived since leftover - // rows are drained over subsequent polls. - if self.indices.is_empty() { - // New cursors are only created once the previous cursor for the stream - // is finished. This means all remaining rows from all but the last batch - // for each stream have been yielded to the newly created record batch - // - // We can therefore drop all but the last batch for each stream - let mut batch_idx = 0; - let mut retained = 0; - self.batches.retain(|(stream_idx, batch)| { - let stream_cursor = &mut self.cursors[*stream_idx]; - let retain = stream_cursor.batch_idx == batch_idx; - batch_idx += 1; - - if retain { - stream_cursor.batch_idx = retained; - retained += 1; - } else { - self.batches_mem_used -= get_record_batch_memory_size(batch); - } - retain - }); - } + self.retain_live_batches(); // Release excess memory back to the pool, but never shrink below // initial_reservation to maintain the anti-starvation guarantee @@ -197,6 +170,49 @@ impl BatchBuilder { RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) } + fn retain_live_batches(&mut self) { + let mut retain_batch = vec![false; self.batches.len()]; + for (batch_idx, _) in &self.indices { + retain_batch[*batch_idx] = true; + } + + 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; + } + } + + let mut batch_idx = 0; + let mut retained = 0; + let mut remap = vec![usize::MAX; self.batches.len()]; + self.batches.retain(|(_, batch)| { + let retain = retain_batch[batch_idx]; + if retain { + remap[batch_idx] = retained; + retained += 1; + } else { + self.batches_mem_used -= get_record_batch_memory_size(batch); + } + batch_idx += 1; + retain + }); + + for (batch_idx, _) in &mut self.indices { + *batch_idx = remap[*batch_idx]; + } + for (stream_idx, cursor) in self.cursors.iter_mut().enumerate() { + if retain_cursor[stream_idx] { + cursor.batch_idx = remap[cursor.batch_idx]; + } + } + } + /// Drains the in_progress row indexes, and builds a new RecordBatch from them /// /// Will then drop any batches for which all rows have been yielded to the output. @@ -280,6 +296,7 @@ mod tests { use arrow::array::{Array, ArrayDataBuilder, Int32Array, ListArray}; use arrow::buffer::Buffer; use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; @@ -303,6 +320,74 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap() } + fn reservation() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("test").register(&pool) + } + + fn int_batch(values: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + fn push_n_rows(builder: &mut BatchBuilder, stream_idx: usize, n: usize) { + for _ in 0..n { + builder.push_row(stream_idx); + } + } + + fn emit_n_rows(builder: &mut BatchBuilder, n: usize) -> RecordBatch { + let columns = builder + .try_interleave_columns(&builder.indices[..n]) + .unwrap(); + builder.finish_record_batch(n, columns).unwrap() + } + + fn assert_int_output(batch: &RecordBatch, expected: &[i32]) { + let actual = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(); + assert_eq!(actual, expected); + } + + #[test] + fn test_partial_emit_releases_unreferenced_and_retains_live_batches() { + let batch0 = int_batch(vec![10, 11]); + let batch1 = int_batch(vec![20, 21]); + let batch2 = int_batch(vec![30, 31]); + let batch1_size = get_record_batch_memory_size(&batch1); + let batch2_size = get_record_batch_memory_size(&batch2); + let schema = batch0.schema(); + let mut builder = BatchBuilder::new(Arc::clone(&schema), 3, 6, reservation()); + + builder.push_batch(0, batch0).unwrap(); + push_n_rows(&mut builder, 0, 2); + builder.push_batch(1, batch1).unwrap(); + push_n_rows(&mut builder, 1, 2); + // Keep one stream empty so stale default cursors cannot retain consumed batches. + builder.push_batch(0, batch2).unwrap(); + + let output = emit_n_rows(&mut builder, 2); + assert_int_output(&output, &[10, 11]); + + assert_eq!(builder.len(), 2); + assert_eq!(builder.batches.len(), 2); + assert_eq!(builder.batches_mem_used, batch1_size + batch2_size); + assert_eq!(builder.reservation.size(), batch1_size + batch2_size); + + push_n_rows(&mut builder, 0, 2); + let output = emit_n_rows(&mut builder, 4); + assert_int_output(&output, &[20, 21, 30, 31]); + + assert!(builder.is_empty()); + assert!(builder.batches.is_empty()); + assert_eq!(builder.batches_mem_used, 0); + assert_eq!(builder.reservation.size(), 0); + } + #[test] fn test_retry_interleave_halves_rows_until_success() { let mut attempts = Vec::new(); From b562850abd1776b79a4b13f2afbc3386b4b02e88 Mon Sep 17 00:00:00 2001 From: Alex Huang Date: Thu, 17 Sep 2026 18:54:23 +0800 Subject: [PATCH 2/3] Preserve fast path for fully drained sort output --- datafusion/physical-plan/src/sorts/builder.rs | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 71ed2667c5f43..13ca81067c59a 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -157,7 +157,11 @@ impl BatchBuilder { // Remove consumed indices, keeping any remaining for the next call. self.indices.drain(..rows_to_emit); - self.retain_live_batches(); + if self.indices.is_empty() { + self.retain_cursor_batches(); + } else { + self.retain_live_batches(); + } // Release excess memory back to the pool, but never shrink below // initial_reservation to maintain the anti-starvation guarantee @@ -170,6 +174,30 @@ impl BatchBuilder { RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) } + fn retain_cursor_batches(&mut self) { + // New cursors are only created once the previous cursor for the stream + // is finished. This means all remaining rows from all but the last batch + // for each stream have been yielded to the newly created record batch + // + // We can therefore drop all but the last live cursor batch for each stream + let mut batch_idx = 0; + let mut retained = 0; + self.batches.retain(|(stream_idx, batch)| { + let stream_cursor = &mut self.cursors[*stream_idx]; + let retain = stream_cursor.batch_idx == batch_idx + && stream_cursor.row_idx < batch.num_rows(); + batch_idx += 1; + + if retain { + stream_cursor.batch_idx = retained; + retained += 1; + } else { + self.batches_mem_used -= get_record_batch_memory_size(batch); + } + retain + }); + } + fn retain_live_batches(&mut self) { let mut retain_batch = vec![false; self.batches.len()]; for (batch_idx, _) in &self.indices { From 6576ed741a70986c736e507a624888728021ca70 Mon Sep 17 00:00:00 2001 From: Alex Huang Date: Thu, 17 Sep 2026 22:06:29 +0800 Subject: [PATCH 3/3] Address released sort cursors --- datafusion/physical-plan/src/sorts/builder.rs | 73 +++++++++++++++---- 1 file changed, 57 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 13ca81067c59a..89763efc4d75c 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -26,7 +26,7 @@ use datafusion_execution::memory_pool::MemoryReservation; use log::warn; use std::sync::Arc; -#[derive(Debug, Copy, Clone, Default)] +#[derive(Debug, Copy, Clone)] struct BatchCursor { /// The index into BatchBuilder::batches batch_idx: usize, @@ -34,6 +34,15 @@ struct BatchCursor { 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, + }; +} + /// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`] #[derive(Debug)] pub struct BatchBuilder { @@ -81,7 +90,7 @@ impl BatchBuilder { Self { schema, batches: Vec::with_capacity(stream_count * 2), - cursors: vec![BatchCursor::default(); stream_count], + cursors: vec![BatchCursor::RELEASED; stream_count], indices: Vec::with_capacity(batch_size), reservation, batches_mem_used: 0, @@ -108,6 +117,14 @@ impl BatchBuilder { /// Append the next row from `stream_idx` pub fn push_row(&mut self, stream_idx: usize) { + debug_assert!( + self.batches + .get(self.cursors[stream_idx].batch_idx) + .is_some_and(|(_, batch)| { + self.cursors[stream_idx].row_idx < batch.num_rows() + }), + "push_row on stream {stream_idx} with no live batch" + ); let cursor = &mut self.cursors[stream_idx]; let row_idx = cursor.row_idx; cursor.row_idx += 1; @@ -184,14 +201,17 @@ impl BatchBuilder { let mut retained = 0; self.batches.retain(|(stream_idx, batch)| { 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); } retain @@ -204,15 +224,13 @@ impl BatchBuilder { retain_batch[*batch_idx] = true; } - 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() - }, - ) { + for cursor in &self.cursors { + if self + .batches + .get(cursor.batch_idx) + .is_some_and(|(_, batch)| cursor.row_idx < batch.num_rows()) + { retain_batch[cursor.batch_idx] = true; - retain_cursor[stream_idx] = true; } } @@ -234,9 +252,10 @@ impl BatchBuilder { for (batch_idx, _) in &mut self.indices { *batch_idx = remap[*batch_idx]; } - 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` means the cursor's batch was released. + cursor.batch_idx = *new_idx; } } } @@ -395,7 +414,7 @@ mod tests { push_n_rows(&mut builder, 0, 2); builder.push_batch(1, batch1).unwrap(); push_n_rows(&mut builder, 1, 2); - // Keep one stream empty so stale default cursors cannot retain consumed batches. + // Keep one stream empty so an unloaded cursor cannot retain consumed batches. builder.push_batch(0, batch2).unwrap(); let output = emit_n_rows(&mut builder, 2); @@ -416,6 +435,28 @@ mod tests { assert_eq!(builder.reservation.size(), 0); } + #[test] + fn test_released_cursor_accepts_new_batch_for_stream() { + let batch0 = int_batch(vec![10]); + let batch1 = int_batch(vec![20]); + let schema = batch0.schema(); + let mut builder = BatchBuilder::new(Arc::clone(&schema), 1, 1, reservation()); + + builder.push_batch(0, batch0).unwrap(); + builder.push_row(0); + let output = emit_n_rows(&mut builder, 1); + assert_int_output(&output, &[10]); + assert!(builder.batches.is_empty()); + + builder.push_batch(0, batch1).unwrap(); + builder.push_row(0); + let output = emit_n_rows(&mut builder, 1); + assert_int_output(&output, &[20]); + assert!(builder.batches.is_empty()); + assert_eq!(builder.batches_mem_used, 0); + assert_eq!(builder.reservation.size(), 0); + } + #[test] fn test_retry_interleave_halves_rows_until_success() { let mut attempts = Vec::new();