fix: Not all rows are accounted in RowCursorStream

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<Rows>, 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();

datafusion/physical-plan/src/sorts/stream.rs

@@ -94,25 +94,62 @@ impl FusedStreams {
     }
 }
 
-/// An `Arc<Rows>` that can be reused
+/// An `Arc<Rows>` 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<Option<Arc<Rows>>>,
+    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<Rows> {
-        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<Rows> {
+        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<Rows>) {
+
+    /// 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<Rows>) -> 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::<Result<Vec<_>>>()?;
 
         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<Option<Self::Output>> {
-        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<dyn MemoryPool>,
+    }
+
+    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<Self> {
+            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<SendableRecordBatchStream> = (0..partitions)
+                .map(|_| {
+                    let batches: Vec<RecordBatch> =
+                        (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<dyn MemoryPool> =
+                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<Option<RowValues>> {
+            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<bool> {
+            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<usize> {
+            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(())
+    }
 }