diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index ff7f259dd1347..107631074ed3d 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>> { 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; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + 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); + } + } }