fix

Author: xudong963

datafusion/physical-plan/src/sorts/merge.rs

@@ -208,6 +208,19 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<RecordBatch>>> {
         if self.done {
+            // When `build_record_batch()` hits an i32 offset overflow (e.g.
+            // combined string offsets exceed 2 GB), it emits a partial batch
+            // and keeps the remaining rows in `self.in_progress.indices`.
+            // Drain those leftover rows before terminating the stream,
+            // otherwise they would be silently dropped.
+            // Repeated overflows are fine — each poll emits another partial
+            // batch until `in_progress` is fully drained.
+            if !self.in_progress.is_empty() {
+                let before = self.in_progress.len();
+                let result = self.in_progress.build_record_batch();
+                self.produced += before - self.in_progress.len();
+                return Poll::Ready(result.transpose());
+            }
             return Poll::Ready(None);
         }
         // Once all partitions have set their corresponding cursors for the loser tree,

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -1539,4 +1539,47 @@ mod tests {
             Err(_) => exec_err!("SortPreservingMerge caused a deadlock"),
         }
     }
+
+    /// Test that SortPreservingMerge with FETCH does not silently drop rows
+    /// when an interleave offset overflow forces a partial batch.
+    #[tokio::test]
+    async fn test_sort_merge_fetch_interleave_overflow() -> Result<()> {
+        // Each string is ~768 MB. Three rows total → ~2.3 GB > i32::MAX,
+        // which forces `build_record_batch` to emit a partial batch.
+        let big_str: String = "x".repeat(768 * 1024 * 1024);
+
+        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
+
+        // Create 3 single-row partitions, each with one large string.
+        let mut partitions = Vec::new();
+        for _ in 0..3 {
+            let array = StringArray::from(vec![big_str.as_str()]);
+            let batch = RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(array) as ArrayRef],
+            )?;
+            partitions.push(vec![batch]);
+        }
+
+        let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?;
+
+        let sort_exprs: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(
+            Column::new("s", 0),
+        )
+            as Arc<dyn PhysicalExpr>)]
+        .into();
+
+        let spm = SortPreservingMergeExec::new(sort_exprs, input).with_fetch(Some(3));
+
+        let task_ctx = Arc::new(TaskContext::default());
+        let batches = collect(Arc::new(spm), task_ctx).await?;
+
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert_eq!(
+            total_rows, 3,
+            "Expected all 3 rows to be emitted despite interleave overflow, got {total_rows}"
+        );
+
+        Ok(())
+    }
 }