chenkovsky commented on code in PR #16983:
URL: https://github.com/apache/datafusion/pull/16983#discussion_r2244222665


##########
datafusion/core/src/physical_planner.rs:
##########
@@ -749,16 +749,26 @@ impl DefaultPhysicalPlanner {
                     Arc::clone(&physical_input_schema),
                 )?);
 
-                let can_repartition = !groups.is_empty()
-                    && session_state.config().target_partitions() > 1
-                    && session_state.config().repartition_aggregations();
-
                 // 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.
                 // To reflect such changes to subsequent stages, use the 
updated
                 // `AggregateFunctionExpr`/`PhysicalSortExpr` objects.
                 let updated_aggregates = initial_aggr.aggr_expr().to_vec();
+                let final_grouping_set = initial_aggr.group_expr().as_final();
+
+                // For aggregations with grouping sets,
+                // the group by expressions differ between the partial and 
final stages,
+                // requiring a shuffle.
+                let has_grouping_id = final_grouping_set
+                    .expr()
+                    .iter()
+                    .any(|(_, name)| name == Aggregate::INTERNAL_GROUPING_ID);
+
+                let can_repartition = !groups.is_empty()
+                    && (session_state.config().target_partitions() > 1
+                        || has_grouping_id)
+                    && session_state.config().repartition_aggregations();
 
                 let next_partition_mode = if can_repartition {

Review Comment:
   maybe we should put code here `if can_repartition || has_grouping_id {`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org

Reply via email to