`RepartitionExec` withholds all output from an unbounded input until `batch_size` rows accumulate

[calvinchengx]

Describe the bug

RepartitionExec coalesces on the producer side. OutputChannel::coalesce
pushes each batch into a LimitedBatchCoalescer and forwards only batches that
have reached target_batch_size (SessionConfig::batch_size(), default 8192).
Whatever is left over is flushed by SharedCoalescer::finalize, which runs when
the last input sender finishes.

An unbounded input never finishes, so the residual is never flushed. A stream
that produces rows slowly, or in small batches, delivers nothing downstream
until 8192 rows have accumulated — however long that takes. The plan is not
deadlocked and no data is lost: it is buffered, invisibly, for an unbounded
amount of time.

The coalescer is created for every non-preserve-order repartition, without
consulting the input's boundedness:

https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/repartition/mod.rs#L528

let shared_coalescer = (!preserve_order).then(|| {
    SharedCoalescer::new(
        input.schema(),
        context.session_config().batch_size(),
        num_input_partitions,
    )
});

ExecutionPlan::boundedness() reports Boundedness::Unbounded for such an
input, but is not consulted here.

To Reproduce

Full program, depending only on datafusion = "54.1.0", futures, and tokio.
A PartitionStream emits one row every 100 ms and never ends; it is fed through
RepartitionExec and the output is drained for six seconds.

src/main.rs

//! `RepartitionExec` withholds all output from an unbounded input until
//! `batch_size` rows accumulate.
//!
//! A source emits one row every 100 ms and never ends. It is fed through
//! `RepartitionExec` with round-robin partitioning, and the output is drained
//! for six seconds. About 60 rows are produced in that window.
//!
//! Run with `cargo run --release`.

use std::sync::Arc;
use std::time::{Duration, Instant};

use datafusion::arrow::array::{Int64Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::common::Result;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_expr::Partitioning;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec};
use datafusion::physical_plan::{ExecutionPlan, execute_stream};
use datafusion::physical_expr::LexOrdering;
use datafusion::prelude::{SessionConfig, SessionContext};
use futures::StreamExt;

/// One row every 100 ms, forever.
#[derive(Debug)]
struct TickStream {
    schema: SchemaRef,
}

impl PartitionStream for TickStream {
    fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
        let schema = Arc::clone(&self.schema);
        let stream = futures::stream::unfold(0i64, move |i| {
            let schema = Arc::clone(&schema);
            async move {
                tokio::time::sleep(Duration::from_millis(100)).await;
                let batch = RecordBatch::try_new(
                    Arc::clone(&schema),
                    vec![Arc::new(Int64Array::from(vec![i]))],
                );
                Some((batch.map_err(Into::into), i + 1))
            }
        });
        Box::pin(RecordBatchStreamAdapter::new(
            Arc::clone(&self.schema),
            stream,
        ))
    }
}

/// Drains the plan for `secs` and returns how many rows arrived.
async fn rows_delivered(batch_size: usize, secs: u64, repartition: bool) -> Result<usize> {
    let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));

    let source = Arc::new(StreamingTableExec::try_new(
        Arc::clone(&schema),
        vec![Arc::new(TickStream {
            schema: Arc::clone(&schema),
        }) as Arc<dyn PartitionStream>],
        None,
        std::iter::empty::<LexOrdering>(),
        true, // infinite
        None,
    )?) as Arc<dyn ExecutionPlan>;

    let plan: Arc<dyn ExecutionPlan> = if repartition {
        Arc::new(RepartitionExec::try_new(
            source,
            Partitioning::RoundRobinBatch(4),
        )?)
    } else {
        source
    };

    let ctx = SessionContext::new_with_config(SessionConfig::new().with_batch_size(batch_size));
    let mut stream = execute_stream(plan, ctx.task_ctx())?;

    let deadline = Instant::now() + Duration::from_secs(secs);
    let mut rows = 0usize;
    while Instant::now() < deadline {
        match tokio::time::timeout(Duration::from_millis(250), stream.next()).await {
            Ok(Some(batch)) => rows += batch?.num_rows(),
            Ok(None) => break,
            Err(_) => {} // no batch this interval; keep waiting
        }
    }
    Ok(rows)
}

#[tokio::main]
async fn main() -> Result<()> {
    let secs = 6;
    println!("source emits 1 row / 100 ms, so ~{} rows per run\n", secs * 10);

    let baseline = rows_delivered(8192, secs, false).await?;
    println!("no repartition,   batch_size = 8192 -> {baseline} rows");

    let repart_default = rows_delivered(8192, secs, true).await?;
    println!("RoundRobinBatch(4), batch_size = 8192 -> {repart_default} rows");

    let repart_small = rows_delivered(1, secs, true).await?;
    println!("RoundRobinBatch(4), batch_size =    1 -> {repart_small} rows");

    Ok(())
}

Output of cargo run --release:

source emits 1 row / 100 ms, so ~60 rows per run

no repartition,   batch_size = 8192 -> 59 rows
RoundRobinBatch(4), batch_size = 8192 -> 0 rows
RoundRobinBatch(4), batch_size =    1 -> 59 rows

The middle line is the bug: the same source, the same six seconds, one
RepartitionExec in between, and nothing arrives. The first line shows the
source itself is fine, and the third isolates the coalescer — the repartition is
still there, only batch_size changed.

Reproduced on 54.1.0; the construction on main is unchanged.

Expected behavior

Rows reach the consumer at a bounded latency, as they do when the same plan runs
without a repartition, or with a small batch_size.

For an unbounded input, waiting for a full batch is not a bounded wait. A few
directions, in the order I would guess you prefer them — though this is your
design call, which is why this is an issue and not a pull request:

1. Skip the producer-side coalescer when the input is Boundedness::Unbounded.
2. Keep it, but flush on a time bound as well as a size bound.
3. Expose it as a config so an unbounded pipeline can turn it off.

Additional context

Some notes to save duplicated work:

• This is not a regression from Integrate BatchCoalescer into RepartitionExec  #18782. The producer-side coalescer arrived
  with that issue, which also removed the coalesce_batches optimizer rule that
  previously inserted a CoalesceBatchesExec after RepartitionExec. I checked
  the older rule (52.0.0) and it did not consider boundedness either, so the
  blind spot predates the integration — Integrate BatchCoalescer into RepartitionExec  #18782 moved it inside the operator
  rather than introducing it.
• FilterExec looks like it has the same shape — it holds a
  LimitedBatchCoalescer and I did not see a boundedness check near it. I have
  not verified that one behaves the same way, so treat it as a lead rather
  than a second report.
• Found while debugging a streaming engine built on DataFusion, where a
  round-robin repartition sat between an unbounded source and a sink and the
  query produced no output at all. Worked around downstream by not repartitioning
  streaming plans; the underlying behaviour is unchanged, which is why I am
  reporting it here.

[jayzhan211]

I think we need emission rules beyond a size-based trigger for unbounded streams. Possible strategies include:

1. Time-based: Emit if no new input arrives within a configurable time interval.
2. Count-based: Emit after N consecutive polls that return no input, indicating the stream is idle.
3. Watermark/checkpoint barrier-based: Emit when a watermark or checkpoint barrier indicates that all relevant events have been observed, following event-time semantics.