From f267045251a0d04747822d688e5336fca99ce4de Mon Sep 17 00:00:00 2001
From: goutamadwant <workwithgoutam@gmail.com>
Date: Sat, 8 Aug 2026 14:15:23 -0700
Subject: [PATCH] fix: avoid buffering unbounded repartition output

Skip producer-side batch coalescing for unbounded inputs so partial batches are not held until stream completion. Preserve bounded-input coalescing and cover the behavior with a pending unbounded source.

Signed-off-by: goutamadwant <workwithgoutam@gmail.com>
---
 .../physical-plan/src/repartition/mod.rs      | 64 +++++++++++++++++--
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index 873f35fd6aed9..ae28ab9838358 100644
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
@@ -455,6 +455,7 @@ impl RepartitionExecState {
 
         let num_input_partitions = streams_and_metrics.len();
         let num_output_partitions = partitioning.partition_count();
+        let coalesce_batches = !preserve_order && !input.boundedness().is_unbounded();
 
         let spill_manager = Arc::new(spill_manager);
 
@@ -523,9 +524,10 @@ impl RepartitionExecState {
 
             // Coalesce on the producer side, before the channel's gate, so
             // the consumer never sees the per-input-task small batches.
-            // Skip in preserve-order mode: each input has its own dedicated
-            // channel and `StreamingMergeBuilder` handles batching.
-            let shared_coalescer = (!preserve_order).then(|| {
+            // Skip in preserve-order mode, where `StreamingMergeBuilder`
+            // handles batching, and for unbounded inputs, where a residual
+            // batch could otherwise be withheld indefinitely.
+            let shared_coalescer = coalesce_batches.then(|| {
                 SharedCoalescer::new(
                     input.schema(),
                     context.session_config().batch_size(),
@@ -1126,7 +1128,8 @@ impl BatchPartitioner {
 /// Repartitioning one [`RecordBatch`] implies creating multiple smaller batches, potentially
 /// as many as the number of output partitions. [`RepartitionExec`] makes sure that the returned
 /// batches adhere to the configured `datafusion.execution.batch_size` for efficient operations,
-/// and for that, it will automatically coalesce batches right after repartitioning.
+/// and for that, it will automatically coalesce batches right after repartitioning for bounded
+/// inputs. Coalescing is skipped for unbounded inputs so partial batches are emitted promptly.
 ///
 /// For this, one shared [`LimitedBatchCoalescer`] per output partition is used:
 ///
@@ -2224,6 +2227,7 @@ mod tests {
     use super::*;
     use crate::empty::EmptyExec;
     use crate::projection::ProjectionExpr;
+    use crate::streaming::{PartitionStream, StreamingTableExec};
     use crate::test::TestMemoryExec;
     use crate::{
         test::{
@@ -2248,6 +2252,27 @@ mod tests {
     use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint};
     use insta::assert_snapshot;
 
+    #[derive(Debug)]
+    struct UnboundedTestPartition {
+        schema: SchemaRef,
+        batch: RecordBatch,
+    }
+
+    impl PartitionStream for UnboundedTestPartition {
+        fn schema(&self) -> &SchemaRef {
+            &self.schema
+        }
+
+        fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+            let stream = futures::stream::iter([Ok(self.batch.clone())])
+                .chain(futures::stream::pending());
+            Box::pin(RecordBatchStreamAdapter::new(
+                Arc::clone(&self.schema),
+                stream,
+            ))
+        }
+    }
+
     #[test]
     fn strength_reduced_u64_remainder_matches_modulo() {
         let divisors = [
@@ -2985,6 +3010,37 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn unbounded_input_emits_before_batch_size() -> Result<()> {
+        let schema = test_schema(false);
+        let batch = create_batch();
+        let source = Arc::new(StreamingTableExec::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(UnboundedTestPartition {
+                schema: Arc::clone(&schema),
+                batch: batch.clone(),
+            })],
+            None,
+            vec![],
+            true,
+            None,
+        )?);
+        let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(1))?;
+        let session_config = SessionConfig::new().with_batch_size(batch.num_rows() * 2);
+        let task_ctx =
+            Arc::new(TaskContext::default().with_session_config(session_config));
+
+        let mut stream = exec.execute(0, task_ctx)?;
+        let output =
+            tokio::time::timeout(std::time::Duration::from_secs(5), stream.next())
+                .await
+                .expect("unbounded repartition withheld a partial batch")
+                .expect("unbounded input ended unexpectedly")?;
+
+        assert_eq!(batch, output);
+        Ok(())
+    }
+
     fn test_schema(nullable: bool) -> Arc<Schema> {
         Arc::new(Schema::new(vec![Field::new(
             "c0",