Fix concurrent partition draining in write execution to prevent memory accumulation

[Copilot]

Description

Partition streams were created sequentially before spawning async consumers, causing repartition producers to fill channels for idle partitions. This manifests as memory buildup in repartition queues with some partitions at 0 bytes while others explode.

Changes

• Moved execute(i, task_ctx) inside spawned tasks: Ensures all partition streams start concurrently rather than sequentially registering with repartition logic before consumption begins
• Applied to both write paths: Standard write and CDC write paths in write_execution_plan_v2

Before:

for i in 0..partition_count {
    let mut stream = inner_plan.execute(i, task_ctx)?;  // Sequential registration
    tokio::task::spawn(async move { /* consume stream */ });
}

After:

for i in 0..partition_count {
    tokio::task::spawn(async move {
        let mut stream = inner_plan.execute(i, task_ctx)?;  // Concurrent execution
        /* consume stream */
    });
}

Related Issue(s)

N/A

Documentation

N/A

Original prompt

Make sure you drain all output partitions concurrently for Upsert operation.
If your application (or the Delta upsert pipeline) consumes partitions sequentially, repartition producers can fill channels for “not currently drained” partitions, causing memory to accumulate in the repartition queues.

This matches the symptom “some partitions at 0 bytes” while others explode: it can happen when some consumers aren’t polling, so only the partitions being drained make progress and the others build up/backpressure patterns get weird.

Fix
Ensure you execute/collect all output partitions concurrently, not one-by-one.
If you’re using something like collect(partition) in a loop, replace with JoinSet / FuturesUnordered over all partitions and await them together (similar to what DataFusion’s own tests do in this file).
If the upsert path is single-threadedly pulling partitions, this is especially important.

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.