You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/03 13:53:29 UTC
(datafusion) branch main updated: feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8155cb7771 feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
8155cb7771 is described below
commit 8155cb7771a4cc0b9ea854385d59f8373964f6dd
Author: Raz Luvaton <16...@users.noreply.github.com>
AuthorDate: Fri Jul 3 16:42:20 2026 +0300
feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
## Which issue does this PR close?
N/A
## Rationale for this change
Drop finished stream that might hold on memory, this will reduce merge
sort memory that one of the stream hold on memory reservation
## What changes are included in this PR?
replaced finished stream in `FusedStreams` with `EmptyRecordBatchStream`
once finished
## Are these changes tested?
yes
## Are there any user-facing changes?
no
---
datafusion/physical-plan/src/sorts/stream.rs | 113 ++++++++++++++++++++++++++-
1 file changed, 109 insertions(+), 4 deletions(-)
diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs
index ff7f259dd1..107631074e 100644
--- a/datafusion/physical-plan/src/sorts/stream.rs
+++ b/datafusion/physical-plan/src/sorts/stream.rs
@@ -15,8 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use crate::SendableRecordBatchStream;
use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues};
+use crate::{EmptyRecordBatchStream, SendableRecordBatchStream};
use crate::{PhysicalExpr, PhysicalSortExpr};
use arrow::array::{Array, UInt32Array};
use arrow::compute::take_record_batch;
@@ -73,9 +73,21 @@ impl FusedStreams {
stream_idx: usize,
) -> Poll<Option<Result<RecordBatch>>> {
loop {
- match ready!(self.0[stream_idx].poll_next_unpin(cx)) {
- Some(Ok(b)) if b.num_rows() == 0 => continue,
- r => return Poll::Ready(r),
+ let poll_result = self.0[stream_idx].poll_next_unpin(cx);
+ match &poll_result {
+ Poll::Pending => return Poll::Pending,
+ Poll::Ready(Some(Ok(b))) if b.num_rows() == 0 => continue,
+ Poll::Ready(Some(Ok(_))) => return poll_result,
+ Poll::Ready(None) | Poll::Ready(Some(Err(_))) => {
+ let stream_schema = self.0[stream_idx].get_ref().schema();
+
+ // Replace the stream with an empty stream, so we can drop memory usage
+ let empty_stream: SendableRecordBatchStream =
+ Box::pin(EmptyRecordBatchStream::new(stream_schema));
+ self.0[stream_idx] = empty_stream.fuse();
+
+ return poll_result;
+ }
}
}
}
@@ -388,8 +400,12 @@ mod tests {
use super::*;
use arrow::array::{AsArray, Int32Array};
use arrow::datatypes::{DataType, Field, Int32Type};
+ use arrow_schema::SchemaRef;
use datafusion_common::DataFusionError;
+ use datafusion_execution::RecordBatchStream;
use datafusion_physical_expr::expressions::col;
+ use futures::Stream;
+ use std::pin::Pin;
/// Verifies that `take_record_batch` in `IncrementalSortIterator` actually
/// copies the data into a new allocation rather than returning a zero-copy
@@ -435,4 +451,93 @@ mod tests {
assert_eq!(total_rows, original_len);
Ok(())
}
+
+ #[test]
+ fn test_fused_stream_drop_finished_streams() {
+ #[derive(Clone)]
+ struct SingleItemManualStream {
+ // Held only so its `Arc` strong count reveals when the stream is dropped.
+ #[expect(dead_code)]
+ hold_ref: Arc<()>,
+ record_batch: RecordBatch,
+ should_finish: bool,
+ }
+
+ impl Stream for SingleItemManualStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ _cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ if !self.should_finish {
+ self.should_finish = true;
+ return Poll::Ready(Some(Ok(self.record_batch.clone())));
+ }
+
+ Poll::Ready(None)
+ }
+ }
+
+ impl RecordBatchStream for SingleItemManualStream {
+ fn schema(&self) -> SchemaRef {
+ self.record_batch.schema()
+ }
+ }
+
+ let hold_ref = Arc::new(());
+ let record_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
+ vec![Arc::new(Int32Array::from(vec![1]))],
+ )
+ .unwrap();
+
+ let stream_1 = SingleItemManualStream {
+ hold_ref: Arc::clone(&hold_ref),
+ should_finish: false,
+ record_batch: record_batch.clone(),
+ };
+ let stream_2 = stream_1.clone();
+
+ let stream_1: SendableRecordBatchStream = Box::pin(stream_1);
+ let stream_2: SendableRecordBatchStream = Box::pin(stream_2);
+
+ let mut fused_stream = FusedStreams(vec![stream_1.fuse(), stream_2.fuse()]);
+
+ let waker = futures::task::noop_waker();
+ let mut cx = Context::from_waker(&waker);
+
+ // The original plus one clone held by each of the two streams.
+ assert_eq!(Arc::strong_count(&hold_ref), 3);
+
+ // First fetch from stream 0 yields its single batch.
+ // the stream is not finished yet, so nothing is dropped.
+ let poll = fused_stream.poll_next(&mut cx, 0);
+ assert!(matches!(poll, Poll::Ready(Some(Ok(_)))));
+ assert_eq!(Arc::strong_count(&hold_ref), 3);
+
+ // Second fetch from stream 0 returns `None`, so it is replaced with an
+ // empty stream and dropped, releasing its `hold_ref` clone.
+ // running 3 times to make sure the stream is fused correctly
+ for _ in 0..3 {
+ let poll = fused_stream.poll_next(&mut cx, 0);
+ assert!(matches!(poll, Poll::Ready(None)));
+ assert_eq!(Arc::strong_count(&hold_ref), 2);
+ }
+
+ // First fetch from stream 1 yields its single batch
+ // the stream is not finished yet, so nothing is dropped.
+ let poll = fused_stream.poll_next(&mut cx, 1);
+ assert!(matches!(poll, Poll::Ready(Some(Ok(_)))));
+ assert_eq!(Arc::strong_count(&hold_ref), 2);
+
+ // Second fetch from stream 1 returns `None`, so it is replaced with an
+ // empty stream and dropped, releasing its `hold_ref` clone.
+ // running 3 times to make sure the stream is fused correctly
+ for _ in 0..3 {
+ let poll = fused_stream.poll_next(&mut cx, 1);
+ assert!(matches!(poll, Poll::Ready(None)));
+ assert_eq!(Arc::strong_count(&hold_ref), 1);
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org