This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23802-e16d453516453bb01b4626df62ed36756ab376c9 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 006dffce955346bb031daa0cf23075221b7d3baf Author: Ariel Miculas-Trif <[email protected]> AuthorDate: Tue Aug 18 14:58:08 2026 +0000 feat: remove extra Rows from ReusableRows (#23802) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #23801 ## Rationale for this change This change reduces the peak allocated memory by removing the unnecessary Rows from being cached. It will only work after https://github.com/apache/datafusion/pull/23619 is merged, until then test_round_robin_tie_breaker_success will fail: Error: Internal("Rows from RowCursorStream is still in use by consumer") test sorts::sort_preserving_merge::tests::test_round_robin_tie_breaker_success ... FAILED The failure is triggered by prev_cursors from SortPreservingMergeStream keeping the previous Cursor alive for round robin tie breaking purposes. The optimization from #23619 only keeps the last Row, so there's no longer a need to keep two Rows cached in ReusableRows. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? There's an existing test that currenly fails with this change and it will pass once #23619 is merged ## Are there any user-facing changes? No --- datafusion/physical-plan/src/sorts/stream.rs | 28 +++++++++------------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index 107631074e..652276c26d 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -93,22 +93,17 @@ impl FusedStreams { } } -/// A pair of `Arc<Rows>` that can be reused +/// An `Arc<Rows>` that can be reused #[derive(Debug)] struct ReusableRows { - // inner[stream_idx] holds a two Arcs: - // at start of a new poll - // .0 is the rows from the previous poll (at start), - // .1 is the one that is being written to - // at end of a poll, .0 will be swapped with .1, - inner: Vec<[Option<Arc<Rows>>; 2]>, + inner: Vec<Option<Arc<Rows>>>, } 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][1].take().unwrap()).map_err(|_| { + Arc::try_unwrap(self.inner[stream_idx].take().unwrap()).map_err(|_| { internal_datafusion_err!( "Rows from RowCursorStream is still in use by consumer" ) @@ -116,16 +111,13 @@ impl ReusableRows { } // save the Rows fn save(&mut self, stream_idx: usize, rows: &Arc<Rows>) { - self.inner[stream_idx][1] = Some(Arc::clone(rows)); - // swap the current with the previous one, so that the next poll can reuse the Rows from the previous poll - let [a, b] = &mut self.inner[stream_idx]; - mem::swap(a, b); + self.inner[stream_idx] = Some(Arc::clone(rows)); } } /// A [`PartitionedStream`] that wraps a set of [`SendableRecordBatchStream`] /// and computes [`RowValues`] based on the provided [`PhysicalSortExpr`] -/// Note: the stream returns an error if the consumer buffers more than one RowValues (i.e. holds on to two RowValues +/// Note: the stream returns an error if the consumer buffers even one RowValues (i.e. holds on to one RowValues /// from the same partition at the same time). #[derive(Debug)] pub struct RowCursorStream { @@ -137,8 +129,9 @@ pub struct RowCursorStream { streams: FusedStreams, /// Tracks the memory used by `converter` reservation: MemoryReservation, - /// Allocated rows for each partition, we keep two to allow for buffering one - /// in the consumer of the stream + /// Reused `Rows` allocation for each partition. The consumer must not + /// buffer the `RowValues` returned for a partition, since the old + /// `Arc<Rows>` must be dropped before that partition can be polled again. rows: ReusableRows, } @@ -162,10 +155,7 @@ impl RowCursorStream { 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))), - Some(Arc::new(converter.empty_rows(0, 0))), - ]); + rows.push(Some(Arc::new(converter.empty_rows(0, 0)))); } Ok(Self { converter, --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
