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:42:43 UTC
(datafusion) branch gh-readonly-queue/main/pr-23283-fb4aa45c82e81f89d1876b74cbb8a3259c204e15 created (now 8155cb7771)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23283-fb4aa45c82e81f89d1876b74cbb8a3259c204e15
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 8155cb7771 feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
This branch includes the following new commits:
new 8155cb7771 feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: feat: Eagerly drop last finished stream in `FusedStreams` (#23283)
Posted by gi...@apache.org.
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-23283-fb4aa45c82e81f89d1876b74cbb8a3259c204e15
in repository https://gitbox.apache.org/repos/asf/datafusion.git
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