From 5e9642695e6d455767c90049c031136cd6ac3f72 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:50:06 +0300 Subject: [PATCH 1/2] fix: report correct number of items left in the iterator for IncrementalSortIterator --- datafusion/physical-plan/src/sorts/stream.rs | 86 +++++++++++++++++--- 1 file changed, 73 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index bb9c00949369e..24ded8a534ff0 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -377,13 +377,15 @@ impl Iterator for IncrementalSortIterator { } fn size_hint(&self) -> (usize, Option) { - let num_rows = self.batch.num_rows(); + let num_rows = self.batch.num_rows().saturating_sub(self.cursor); let batch_size = self.batch_size; let num_batches = num_rows.div_ceil(batch_size); (num_batches, Some(num_batches)) } } +impl ExactSizeIterator for IncrementalSortIterator {} + impl FusedIterator for IncrementalSortIterator {} #[cfg(test)] @@ -398,19 +400,14 @@ mod tests { 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 - /// slice of the original batch. If the output arrays were slices, their - /// underlying buffer length would match the original array's length; a true - /// copy will have a buffer sized to fit only the chunk. - #[test] - fn incremental_sort_iterator_copies_data() -> Result<()> { - let original_len = 10; - let batch_size = 3; - + fn create_incremental_sort_iter_on( + input_batch_len: usize, + output_batch_size: usize, + ) -> Result<(IncrementalSortIterator, RecordBatch)> { // Build a batch with a single Int32 column of descending values let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let col_a: Int32Array = Int32Array::from(vec![0; original_len]); + let col_a: Int32Array = + Int32Array::from_iter_values((0..input_batch_len as i32).rev()); let batch = RecordBatch::try_new(schema, vec![Arc::new(col_a)])?; // Sort ascending on column "a" @@ -420,8 +417,27 @@ mod tests { )?)]) .unwrap(); + let iter = + IncrementalSortIterator::new(batch.clone(), expressions, output_batch_size); + + Ok((iter, batch)) + } + + /// Verifies that `take_record_batch` in `IncrementalSortIterator` actually + /// copies the data into a new allocation rather than returning a zero-copy + /// slice of the original batch. If the output arrays were slices, their + /// underlying buffer length would match the original array's length; a true + /// copy will have a buffer sized to fit only the chunk. + #[test] + fn incremental_sort_iterator_copies_data() -> Result<()> { + let original_len = 10; + let batch_size = 3; + + let (mut iter, batch) = + create_incremental_sort_iter_on(original_len, batch_size)?; + let mut total_rows = 0; - IncrementalSortIterator::new(batch.clone(), expressions, batch_size).try_for_each( + iter.try_for_each( |result| { let chunk = result?; total_rows += chunk.num_rows(); @@ -531,4 +547,48 @@ mod tests { assert_eq!(Arc::strong_count(&hold_ref), 1); } } + + fn assert_iterator_len_and_batch_len( + iter: &IncrementalSortIterator, + expected_len: usize, + ) { + assert_eq!(iter.len(), expected_len); + assert_eq!(iter.size_hint(), (expected_len, Some(expected_len))); + } + + #[test] + fn incremental_sort_iterator_report_correct_len() -> Result<()> { + let original_len = 10; + let batch_size = 3; + + let (mut iterator, _) = + create_incremental_sort_iter_on(original_len, batch_size)?; + + assert_iterator_len_and_batch_len(&iterator, 4); + + let batch = iterator.next().unwrap()?; + assert_eq!(batch.num_rows(), batch_size); + + assert_iterator_len_and_batch_len(&iterator, 3); + + let batch = iterator.next().unwrap()?; + assert_eq!(batch.num_rows(), batch_size); + + assert_iterator_len_and_batch_len(&iterator, 2); + + let batch = iterator.next().unwrap()?; + assert_eq!(batch.num_rows(), batch_size); + + assert_iterator_len_and_batch_len(&iterator, 1); + + let batch = iterator.next().unwrap()?; + // left over + assert_eq!(batch.num_rows(), 1); + + assert_iterator_len_and_batch_len(&iterator, 0); + + assert!(iterator.next().is_none()); + + Ok(()) + } } From 1d2c0942f1c12a75b3e41e62125a612d0082f68d Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:34:46 +0300 Subject: [PATCH 2/2] remove exact size iterator --- datafusion/physical-plan/src/sorts/stream.rs | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index 24ded8a534ff0..e2a840fdedfd5 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -376,6 +376,8 @@ impl Iterator for IncrementalSortIterator { } } + // Not implementing ExactSizeIterator since in case of an error we stop and don't emit any more + // so the length would be wrong fn size_hint(&self) -> (usize, Option) { let num_rows = self.batch.num_rows().saturating_sub(self.cursor); let batch_size = self.batch_size; @@ -384,8 +386,6 @@ impl Iterator for IncrementalSortIterator { } } -impl ExactSizeIterator for IncrementalSortIterator {} - impl FusedIterator for IncrementalSortIterator {} #[cfg(test)] @@ -548,11 +548,7 @@ mod tests { } } - fn assert_iterator_len_and_batch_len( - iter: &IncrementalSortIterator, - expected_len: usize, - ) { - assert_eq!(iter.len(), expected_len); + fn assert_iterator_size_hint(iter: &IncrementalSortIterator, expected_len: usize) { assert_eq!(iter.size_hint(), (expected_len, Some(expected_len))); } @@ -564,28 +560,28 @@ mod tests { let (mut iterator, _) = create_incremental_sort_iter_on(original_len, batch_size)?; - assert_iterator_len_and_batch_len(&iterator, 4); + assert_iterator_size_hint(&iterator, 4); let batch = iterator.next().unwrap()?; assert_eq!(batch.num_rows(), batch_size); - assert_iterator_len_and_batch_len(&iterator, 3); + assert_iterator_size_hint(&iterator, 3); let batch = iterator.next().unwrap()?; assert_eq!(batch.num_rows(), batch_size); - assert_iterator_len_and_batch_len(&iterator, 2); + assert_iterator_size_hint(&iterator, 2); let batch = iterator.next().unwrap()?; assert_eq!(batch.num_rows(), batch_size); - assert_iterator_len_and_batch_len(&iterator, 1); + assert_iterator_size_hint(&iterator, 1); let batch = iterator.next().unwrap()?; // left over assert_eq!(batch.num_rows(), 1); - assert_iterator_len_and_batch_len(&iterator, 0); + assert_iterator_size_hint(&iterator, 0); assert!(iterator.next().is_none());