From 9b992dd5c4a1649d3c7af6eedcf9599172480dcf Mon Sep 17 00:00:00 2001 From: Emily Matheys Date: Wed, 16 Sep 2026 12:49:32 +0300 Subject: [PATCH] fix: Not all rows are accounted in RowCursorStream --- datafusion/physical-plan/src/sorts/cursor.rs | 11 +- datafusion/physical-plan/src/sorts/stream.rs | 287 +++++++++++++++++-- 2 files changed, 270 insertions(+), 28 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index f12e0daab5909..7ee63bfc63d53 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -206,12 +206,13 @@ 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. + /// The reservation must cover exactly `rows.size()` bytes, or be empty when + /// the caller accounts for `rows` elsewhere for its whole lifetime. + /// + /// Panics if the reservation is neither of those, or if `rows` is empty. pub fn new(rows: Arc, reservation: MemoryReservation) -> Self { - assert_eq!( - rows.size(), - reservation.size(), + assert!( + reservation.size() == 0 || reservation.size() == rows.size(), "memory reservation mismatch" ); let len = rows.num_rows(); diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index e2a840fdedfd5..978acaaf90d81 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -94,25 +94,62 @@ impl FusedStreams { } } -/// An `Arc` that can be reused +/// An `Arc` that can be reused. +/// +/// Owns the reservation covering every retained buffer for as long as it is +/// retained, so the cache is visible to the pool rather than held off-book. +/// A retained buffer keeps its capacity, so the cost is the high-water mark of +/// each stream, not the size of the batch currently in flight. #[derive(Debug)] struct ReusableRows { inner: Vec>>, + reservation: MemoryReservation, } impl ReusableRows { // return a Rows for writing, // does not clone if the existing rows can be reused - fn take_next(&mut self, stream_idx: usize) -> Result { - Arc::try_unwrap(self.inner[stream_idx].take().unwrap()).map_err(|_| { - internal_datafusion_err!( - "Rows from RowCursorStream is still in use by consumer" - ) - }) + fn take_next(&mut self, stream_idx: usize, converter: &RowConverter) -> Result { + match self.inner[stream_idx].take() { + Some(rows) => Arc::try_unwrap(rows).map_err(|_| { + internal_datafusion_err!( + "Rows from RowCursorStream is still in use by consumer" + ) + }), + // Nothing retained yet, or already released, so start over. + None => Ok(converter.empty_rows(0, 0)), + } } - // save the Rows - fn save(&mut self, stream_idx: usize, rows: &Arc) { + + /// Account for a freshly built buffer, and retain it for reuse. + /// + /// The reservation is mandatory rather than best-effort. The buffer is live in the + /// cursor whether or not this slot keeps a handle to it, so declining to reserve + /// would hide it from the pool instead of avoiding it, and it frees nothing at this + /// point either, since the cursor holds the same `Arc`. Retention on top of the + /// reservation is free for the same reason. A pool that cannot cover the buffer + /// fails the query here, as it did before the buffer was cached at all. + fn save(&mut self, stream_idx: usize, rows: &Arc) -> Result<()> { self.inner[stream_idx] = Some(Arc::clone(rows)); + let retained = self.retained_size(); + debug_assert!(retained >= self.reservation.size()); + if let Err(e) = self.reservation.try_resize(retained) { + self.inner[stream_idx] = None; + return Err(e); + } + Ok(()) + } + + // drop whatever a finished stream was holding + fn release(&mut self, stream_idx: usize) { + debug_assert!(self.reservation.size() >= self.retained_size()); + if let Some(rows) = self.inner[stream_idx].take() { + self.reservation.shrink(rows.size()); + } + } + + fn retained_size(&self) -> usize { + self.inner.iter().flatten().map(|rows| rows.size()).sum() } } @@ -152,18 +189,18 @@ impl RowCursorStream { .collect::>>()?; let streams: Vec<_> = streams.into_iter().map(|s| s.fuse()).collect(); + let stream_count = streams.len(); let converter = RowConverter::new(sort_fields)?; - let mut rows = Vec::with_capacity(streams.len()); - for _ in &streams { - // Initialize each stream with an empty Rows - rows.push(Some(Arc::new(converter.empty_rows(0, 0)))); - } + let rows_reservation = reservation.new_empty(); Ok(Self { converter, reservation, column_expressions: expressions.iter().map(|x| Arc::clone(&x.expr)).collect(), streams: FusedStreams(streams), - rows: ReusableRows { inner: rows }, + rows: ReusableRows { + inner: vec![None; stream_count], + reservation: rows_reservation, + }, }) } @@ -175,7 +212,7 @@ impl RowCursorStream { let cols = evaluate_expressions_to_arrays(&self.column_expressions, batch)?; // At this point, ownership should of this Rows should be unique - let mut rows = self.rows.take_next(stream_idx)?; + let mut rows = self.rows.take_next(stream_idx, &self.converter)?; rows.clear(); @@ -184,12 +221,13 @@ impl RowCursorStream { let rows = Arc::new(rows); - self.rows.save(stream_idx, &rows); + self.rows.save(stream_idx, &rows)?; - // track the memory in the newly created Rows. - let rows_reservation = self.reservation.new_empty(); - rows_reservation.try_grow(rows.size())?; - Ok(RowValues::new(rows, rows_reservation)) + // `self.rows` now holds the reservation for this buffer unconditionally, and + // holds it for at least as long as the cursor does. `take_next` cannot reclaim + // the slot while the consumer still owns the `Arc`, so the cursor is handed an + // empty reservation rather than accounting for the same bytes a second time. + Ok(RowValues::new(rows, self.reservation.new_empty())) } } @@ -205,7 +243,12 @@ impl PartitionedStream for RowCursorStream { cx: &mut Context<'_>, stream_idx: usize, ) -> Poll> { - Poll::Ready(ready!(self.streams.poll_next(cx, stream_idx)).map(|r| { + let polled = ready!(self.streams.poll_next(cx, stream_idx)); + if polled.is_none() { + // The stream is finished, so its retained buffer will never be reused. + self.rows.release(stream_idx); + } + Poll::Ready(polled.map(|r| { r.and_then(|batch| { let cursor = self.convert_batch(&batch, stream_idx)?; Ok((cursor, batch)) @@ -391,11 +434,15 @@ impl FusedIterator for IncrementalSortIterator {} #[cfg(test)] mod tests { use super::*; - use arrow::array::{AsArray, Int32Array}; + use crate::memory::MemoryStream; + use arrow::array::{AsArray, Int32Array, StringArray}; use arrow::datatypes::{DataType, Field, Int32Type}; use arrow_schema::SchemaRef; use datafusion_common::DataFusionError; use datafusion_execution::RecordBatchStream; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; use datafusion_physical_expr::expressions::col; use futures::Stream; use std::pin::Pin; @@ -587,4 +634,198 @@ mod tests { Ok(()) } + + /// Drives a `RowCursorStream` over several partitions and reports the pool + /// reservation, so the tests below can watch it move. + struct RowCursorHarness { + stream: RowCursorStream, + pool: Arc, + } + + impl RowCursorHarness { + /// `batches_per_partition` batches of `rows` rows each, per partition. + /// The sort key spans two columns so this takes the `Rows` path rather + /// than the specialized single-column `FieldCursorStream`. + fn new( + partitions: usize, + batches_per_partition: usize, + rows: usize, + str_width: usize, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let make_batch = |seq: usize| { + let base = (seq * rows) as i32; + let a = Int32Array::from_iter_values((0..rows as i32).map(|i| base + i)); + let b = StringArray::from_iter_values( + (0..rows).map(|_| "x".repeat(str_width)), + ); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(a), Arc::new(b)]) + .unwrap() + }; + + let streams: Vec = (0..partitions) + .map(|_| { + let batches: Vec = + (0..batches_per_partition).map(make_batch).collect(); + Box::pin( + MemoryStream::try_new(batches, Arc::clone(&schema), None) + .unwrap(), + ) as SendableRecordBatchStream + }) + .collect(); + + let expressions = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(col("a", &schema)?), + PhysicalSortExpr::new_default(col("b", &schema)?), + ]) + .unwrap(); + + // Generous limit: this measures what is reserved, not what is refused. + let pool: Arc = + Arc::new(GreedyMemoryPool::new(256 * 1024 * 1024)); + let reservation = MemoryConsumer::new("test").register(&pool); + let stream = + RowCursorStream::try_new(&schema, &expressions, streams, reservation)?; + + Ok(Self { stream, pool }) + } + + fn reserved(&self) -> usize { + self.pool.reserved() + } + + /// Polls `stream_idx`, handing back the cursor so the caller controls when + /// it is dropped. `None` once the partition is finished. + fn poll_cursor(&mut self, stream_idx: usize) -> Result> { + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + match self.stream.poll_next(&mut cx, stream_idx) { + Poll::Ready(Some(Ok((values, _batch)))) => Ok(Some(values)), + Poll::Ready(Some(Err(e))) => Err(e), + Poll::Ready(None) => Ok(None), + Poll::Pending => unreachable!("MemoryStream is never pending"), + } + } + + /// Polls `stream_idx` and immediately drops the returned cursor, which is + /// what the merge does once a cursor is exhausted. Returns `false` once + /// the partition is finished. + fn poll_and_drop_cursor(&mut self, stream_idx: usize) -> Result { + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + match self.stream.poll_next(&mut cx, stream_idx) { + Poll::Ready(Some(Ok((values, _batch)))) => { + drop(values); + Ok(true) + } + Poll::Ready(Some(Err(e))) => Err(e), + Poll::Ready(None) => Ok(false), + Poll::Pending => unreachable!("MemoryStream is never pending"), + } + } + } + + // `RowCursorStream` keeps one `Rows` buffer per partition alive between + // polls so the allocation can be reused. That buffer is real memory, and it + // outlives the cursor handed to the consumer: the merge drops a cursor as + // soon as it is exhausted, but the buffer behind it stays cached. + // + // So dropping the cursor must not change what the pool reports. If it does, + // the bytes still held by the cache have gone off the books and the pool is + // under-counting a live allocation. + #[test] + fn dropping_a_cursor_does_not_unaccount_its_retained_buffer() -> Result<()> { + let mut harness = RowCursorHarness::new(1, 4, 512, 64)?; + + let cursor = harness.poll_cursor(0)?.expect("first batch"); + let with_cursor_alive = harness.reserved(); + + drop(cursor); + let after_drop = harness.reserved(); + + assert_eq!( + with_cursor_alive, after_drop, + "the cached `Rows` buffer outlives the cursor, so dropping the cursor \ + must not release its bytes (before {with_cursor_alive}, after {after_drop})" + ); + + // And the buffer is genuinely on the books, not merely unchanged at zero. + let baseline = RowCursorHarness::new(1, 4, 512, 64)?.reserved(); + assert!( + after_drop > baseline, + "retained buffer should be reserved (baseline {baseline}, now {after_drop})" + ); + + Ok(()) + } + + // The retained bytes are per partition, so what the pool reports has to + // scale with partition count rather than staying flat at the size of + // whichever batch happens to be in flight. + #[test] + fn retained_row_buffer_reservation_scales_with_partitions() -> Result<()> { + let measure = |partitions: usize| -> Result { + let mut harness = RowCursorHarness::new(partitions, 4, 512, 64)?; + for stream_idx in 0..partitions { + assert!(harness.poll_and_drop_cursor(stream_idx)?); + } + Ok(harness.reserved()) + }; + + let few = measure(2)?; + let many = measure(16)?; + + // 8x the partitions. The converter reservation is shared and does not + // scale, so this is deliberately loose - the point is that it grows with + // the number of cached buffers, not that it grows by an exact factor. + assert!( + many > few * 4, + "reservation should scale with the number of retained buffers, \ + but 2 partitions reserved {few} and 16 reserved {many}" + ); + + Ok(()) + } + + // The flip side: a partition that will never be polled again must hand its + // buffer back. Otherwise the reservation only ever grows and a long merge + // holds every partition's high-water mark until the whole stream drops. + #[test] + fn exhausted_partitions_release_their_retained_buffers() -> Result<()> { + let partitions = 8; + let batches_per_partition = 3; + let mut harness = + RowCursorHarness::new(partitions, batches_per_partition, 512, 64)?; + + let baseline = harness.reserved(); + + for stream_idx in 0..partitions { + for _ in 0..batches_per_partition { + assert!(harness.poll_and_drop_cursor(stream_idx)?); + } + } + let peak = harness.reserved(); + assert!( + peak > baseline, + "expected the cached buffers to be reserved" + ); + + // Poll each partition once more so it reports exhaustion. + for stream_idx in 0..partitions { + assert!(!harness.poll_and_drop_cursor(stream_idx)?); + } + + let after_release = harness.reserved(); + assert!( + after_release < peak, + "exhausted partitions should release their buffers \ + (peak {peak}, after {after_release})" + ); + + Ok(()) + } }