(datafusion) branch main updated: fix: report correct number of items left in the iterator for `IncrementalSortIterator` (#25107)

[gi...@apache.org]

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 c0b2ebe178 fix: report correct number of items left in the iterator for `IncrementalSortIterator` (#25107)
c0b2ebe178 is described below

commit c0b2ebe17853d681b002bbbe5066629da25b642f
Author: Raz Luvaton <16...@users.noreply.github.com>
AuthorDate: Thu Sep 10 11:05:59 2026 +0000

    fix: report correct number of items left in the iterator for `IncrementalSortIterator` (#25107)
    
    ## Which issue does this PR close?
    
    N/A
    
    ## Rationale for this change
    `size_hint` for `IncrementalSortIterator` returns the initial number of
    items rather than how many items left in the iterator
    
    ## What changes are included in this PR?
    
    update size hint to consider the number of elements already produced,
    impl `ExactSizeIterator` and added test
    
    ## What is the testing strategy for this PR?
    unit test
    
    ## Are there any user-facing changes?
    no
---
 datafusion/physical-plan/src/sorts/stream.rs | 82 +++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 13 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs
index bb9c009493..e2a840fded 100644
--- a/datafusion/physical-plan/src/sorts/stream.rs
+++ b/datafusion/physical-plan/src/sorts/stream.rs
@@ -376,8 +376,10 @@ 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<usize>) {
-        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))
@@ -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,44 @@ mod tests {
             assert_eq!(Arc::strong_count(&hold_ref), 1);
         }
     }
+
+    fn assert_iterator_size_hint(iter: &IncrementalSortIterator, expected_len: usize) {
+        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_size_hint(&iterator, 4);
+
+        let batch = iterator.next().unwrap()?;
+        assert_eq!(batch.num_rows(), batch_size);
+
+        assert_iterator_size_hint(&iterator, 3);
+
+        let batch = iterator.next().unwrap()?;
+        assert_eq!(batch.num_rows(), batch_size);
+
+        assert_iterator_size_hint(&iterator, 2);
+
+        let batch = iterator.next().unwrap()?;
+        assert_eq!(batch.num_rows(), batch_size);
+
+        assert_iterator_size_hint(&iterator, 1);
+
+        let batch = iterator.next().unwrap()?;
+        // left over
+        assert_eq!(batch.num_rows(), 1);
+
+        assert_iterator_size_hint(&iterator, 0);
+
+        assert!(iterator.next().is_none());
+
+        Ok(())
+    }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org