Fix channel-based multi-output routing and output partitioning

- Fix ProducingOutput to carry partition index alongside batch, ensuring
  correct routing when emit_next_partition eagerly advances the index
- Add use_channels() method to centralize the decision of when to use
  the channel-based multi-output path
- Add update_cache_partitioning() to keep output partitioning in sync
  when limit_options or num_agg_partitions change
- Fix with_new_limit_options to recalculate output partitioning (prevents
  TopK aggregation from claiming Hash partitioning when channels won't
  be used)
- Guard CombinePartialFinalAggregate from combining when Partial has
  num_agg_partitions > 1
- Set num_agg_partitions in physical planner when repartitioning is enabled
- Keep spawned task references alive via Arc in output streams to prevent
  abort-on-drop

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Dandandan

datafusion/core/src/physical_planner.rs

@@ -1043,19 +1043,39 @@ impl DefaultPhysicalPlanner {
                     input_exec
                 };
 
-                let initial_aggr = Arc::new(AggregateExec::try_new(
-                    AggregateMode::Partial,
-                    groups.clone(),
-                    aggregates,
-                    filters.clone(),
-                    input_exec,
-                    Arc::clone(&physical_input_schema),
-                )?);
-
                 let can_repartition = !groups.is_empty()
                     && session_state.config().target_partitions() > 1
                     && session_state.config().repartition_aggregations();
 
+                let initial_aggr = if can_repartition {
+                    // When repartitioning is enabled, set num_agg_partitions
+                    // so the partial aggregate produces target_partitions
+                    // output partitions directly (eliminating the need for
+                    // a downstream RepartitionExec).
+                    Arc::new(
+                        AggregateExec::try_new(
+                            AggregateMode::Partial,
+                            groups.clone(),
+                            aggregates,
+                            filters.clone(),
+                            input_exec,
+                            Arc::clone(&physical_input_schema),
+                        )?
+                        .with_num_agg_partitions(
+                            session_state.config().target_partitions(),
+                        ),
+                    )
+                } else {
+                    Arc::new(AggregateExec::try_new(
+                        AggregateMode::Partial,
+                        groups.clone(),
+                        aggregates,
+                        filters.clone(),
+                        input_exec,
+                        Arc::clone(&physical_input_schema),
+                    )?)
+                };
+
                 // Some aggregators may be modified during initialization for
                 // optimization purposes. For example, a FIRST_VALUE may turn
                 // into a LAST_VALUE with the reverse ordering requirement.

datafusion/physical-optimizer/src/combine_partial_final_agg.rs

@@ -73,6 +73,7 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
             };
 
             let transformed = if *input_agg_exec.mode() == AggregateMode::Partial
+                && input_agg_exec.num_agg_partitions() <= 1
                 && can_combine(
                     (
                         agg_exec.group_expr(),

datafusion/physical-plan/src/aggregates/mod.rs

@@ -634,8 +634,7 @@ enum AggRepartitionState {
     /// Input tasks have been spawned and are sending to output channels.
     Consuming {
         /// One receiver per output partition. `None` if already taken by `execute()`.
-        receivers:
-            Vec<Option<tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>>>,
+        receivers: Vec<Option<tokio::sync::mpsc::UnboundedReceiver<Result<RecordBatch>>>>,
         /// Background tasks; dropped to abort if the exec is dropped.
         _abort_helper: Vec<SpawnedTask<()>>,
     },
@@ -745,7 +744,7 @@ impl AggregateExec {
 
     /// Clone this exec, overriding only the limit hint.
     pub fn with_new_limit_options(&self, limit_options: Option<LimitOptions>) -> Self {
-        Self {
+        let mut new = Self {
             limit_options,
             // clone the rest of the fields
             required_input_ordering: self.required_input_ordering.clone(),
@@ -762,7 +761,9 @@ impl AggregateExec {
             dynamic_filter: self.dynamic_filter.clone(),
             num_agg_partitions: self.num_agg_partitions,
             repartition_state: Arc::new(Mutex::new(AggRepartitionState::default())),
-        }
+        };
+        new.update_cache_partitioning();
+        new
     }
 
     pub fn cache(&self) -> &PlanProperties {
@@ -910,27 +911,55 @@ impl AggregateExec {
         &self.mode
     }
 
+    /// Number of internal hash table partitions.
+    pub fn num_agg_partitions(&self) -> usize {
+        self.num_agg_partitions
+    }
+
     /// Set the number of internal hash table partitions for partial aggregation.
     /// When n > 1, the aggregate also acts as a repartitioning operator,
     /// producing n output partitions.
     pub fn with_num_agg_partitions(mut self, n: usize) -> Self {
         self.num_agg_partitions = n;
-        if n > 1 {
-            // Update output partitioning to reflect the repartitioned output.
-            let group_exprs = self.output_group_expr();
-            let mut cache = (*self.cache).clone();
-            cache.partitioning = Partitioning::Hash(group_exprs, n);
-            self.cache = Arc::new(cache);
-        }
+        self.update_cache_partitioning();
         self
     }
 
     /// Set the limit options for this AggExec
     pub fn with_limit_options(mut self, limit_options: Option<LimitOptions>) -> Self {
         self.limit_options = limit_options;
+        self.update_cache_partitioning();
         self
     }
 
+    /// Returns true if the channel-based multi-output path will be used
+    /// in execute().
+    fn use_channels(&self) -> bool {
+        self.num_agg_partitions > 1
+            && !self.group_by.is_true_no_grouping()
+            && (self.limit_options.is_none()
+                || self.is_unordered_unfiltered_group_by_distinct())
+    }
+
+    /// Update the cached output partitioning based on whether channels
+    /// will be used.
+    fn update_cache_partitioning(&mut self) {
+        if self.use_channels() {
+            let group_exprs = self.output_group_expr();
+            let mut cache = (*self.cache).clone();
+            cache.partitioning = Partitioning::Hash(group_exprs, self.num_agg_partitions);
+            self.cache = Arc::new(cache);
+        } else if self.num_agg_partitions > 1 {
+            // Channels won't be used (e.g. limit/TopK was set), revert
+            // to single-partition output matching the non-channel path.
+            let mut cache = (*self.cache).clone();
+            cache.partitioning = Partitioning::UnknownPartitioning(
+                self.input.output_partitioning().partition_count(),
+            );
+            self.cache = Arc::new(cache);
+        }
+    }
+
     /// Get the limit options (if set)
     pub fn limit_options(&self) -> Option<LimitOptions> {
         self.limit_options
@@ -1028,8 +1057,7 @@ impl AggregateExec {
                 Err(e) => {
                     let e = Arc::new(e);
                     for sender in &senders {
-                        let _ =
-                            sender.send(Err(DataFusionError::from(&e)));
+                        let _ = sender.send(Err(DataFusionError::from(&e)));
                     }
                     return;
                 }
@@ -1553,8 +1581,7 @@ impl ExecutionPlan for AggregateExec {
         if self.num_agg_partitions > 1 {
             let group_exprs = me.output_group_expr();
             let mut cache = (*me.cache).clone();
-            cache.partitioning =
-                Partitioning::Hash(group_exprs, self.num_agg_partitions);
+            cache.partitioning = Partitioning::Hash(group_exprs, self.num_agg_partitions);
             me.cache = Arc::new(cache);
         }
 
@@ -1566,7 +1593,7 @@ impl ExecutionPlan for AggregateExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        if self.num_agg_partitions <= 1 {
+        if !self.use_channels() {
             return self
                 .execute_typed(partition, &context)
                 .map(|stream| stream.into());
@@ -1593,18 +1620,12 @@ impl ExecutionPlan for AggregateExec {
             // Spawn one task per input partition
             let mut tasks = Vec::with_capacity(num_input);
             for input_idx in 0..num_input {
-                let stream = GroupedHashAggregateStream::new(
-                    self, &context, input_idx,
-                )?;
-                let senders_clone: Vec<_> =
-                    senders.iter().map(|s| s.clone()).collect();
+                let stream = GroupedHashAggregateStream::new(self, &context, input_idx)?;
+                let senders_clone: Vec<_> = senders.iter().map(|s| s.clone()).collect();
                 let n = num_output;
 
-                let task = SpawnedTask::spawn(Self::pull_and_route(
-                    stream,
-                    n,
-                    senders_clone,
-                ));
+                let task =
+                    SpawnedTask::spawn(Self::pull_and_route(stream, n, senders_clone));
                 tasks.push(task);
             }
 
@@ -1615,8 +1636,7 @@ impl ExecutionPlan for AggregateExec {
         }
 
         // Take this partition's receiver
-        let AggRepartitionState::Consuming { receivers, .. } = &mut *state
-        else {
+        let AggRepartitionState::Consuming { receivers, .. } = &mut *state else {
             unreachable!()
         };
 
@@ -1625,12 +1645,17 @@ impl ExecutionPlan for AggregateExec {
             .and_then(|r| r.take())
             .expect("partition receiver already consumed");
         let schema = Arc::clone(&self.schema);
+        // Keep a reference to the shared state so spawned tasks aren't
+        // aborted when the AggregateExec is dropped but streams are
+        // still being polled.
+        let _state_ref = Arc::clone(&self.repartition_state);
         drop(state);
 
         // Wrap the receiver as a RecordBatchStream
-        let stream = futures::stream::unfold(rx, |mut rx| async move {
-            rx.recv().await.map(|batch| (batch, rx))
-        });
+        let stream =
+            futures::stream::unfold((rx, _state_ref), |(mut rx, state_ref)| async move {
+                rx.recv().await.map(|batch| (batch, (rx, state_ref)))
+            });
         Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
     }
 
@@ -4348,9 +4373,7 @@ mod tests {
             Arc::clone(&schema),
             vec![
                 Arc::new(Int32Array::from(vec![1, 2, 3, 1, 2, 3])),
-                Arc::new(Float64Array::from(vec![
-                    10.0, 20.0, 30.0, 40.0, 50.0, 60.0,
-                ])),
+                Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0])),
             ],
         )?;
         let batch2 = RecordBatch::try_new(
@@ -4364,25 +4387,17 @@ mod tests {
 
         let task_ctx = Arc::new(TaskContext::default());
 
-        let group_expr =
-            vec![(col("group_col", &schema)?, "group_col".to_string())];
+        let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())];
         let aggr_expr = vec![Arc::new(
-            AggregateExprBuilder::new(
-                sum_udaf(),
-                vec![col("value_col", &schema)?],
-            )
-            .schema(Arc::clone(&schema))
-            .alias("sum_value")
-            .build()?,
+            AggregateExprBuilder::new(sum_udaf(), vec![col("value_col", &schema)?])
+                .schema(Arc::clone(&schema))
+                .alias("sum_value")
+                .build()?,
         )];
 
         let num_output_partitions = 3;
 
-        let exec = TestMemoryExec::try_new(
-            &input_partitions,
-            Arc::clone(&schema),
-            None,
-        )?;
+        let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?;
         let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec)));
 
         let aggregate_exec = Arc::new(
@@ -4399,7 +4414,10 @@ mod tests {
 
         // Verify output partitioning is Hash with num_output_partitions
         assert_eq!(
-            aggregate_exec.properties().output_partitioning().partition_count(),
+            aggregate_exec
+                .properties()
+                .output_partitioning()
+                .partition_count(),
             num_output_partitions
         );
 
@@ -4408,8 +4426,7 @@ mod tests {
             std::collections::HashMap::new();
 
         for partition in 0..num_output_partitions {
-            let stream =
-                aggregate_exec.execute(partition, Arc::clone(&task_ctx))?;
+            let stream = aggregate_exec.execute(partition, Arc::clone(&task_ctx))?;
             let batches: Vec<_> = stream
                 .collect::<Vec<_>>()
                 .await
@@ -4428,8 +4445,7 @@ mod tests {
                     .downcast_ref::<Float64Array>()
                     .unwrap();
                 for i in 0..batch.num_rows() {
-                    *all_results.entry(groups.value(i)).or_insert(0.0) +=
-                        sums.value(i);
+                    *all_results.entry(groups.value(i)).or_insert(0.0) += sums.value(i);
                 }
             }
         }