@@ -29,6 +29,7 @@ use crate::physical_plan::joins::{
2929use crate :: physical_plan:: projection:: ProjectionExec ;
3030use crate :: physical_plan:: repartition:: RepartitionExec ;
3131use crate :: physical_plan:: sorts:: sort:: SortOptions ;
32+ use crate :: physical_plan:: union:: { can_interleave, InterleaveExec , UnionExec } ;
3233use crate :: physical_plan:: windows:: WindowAggExec ;
3334use crate :: physical_plan:: Partitioning ;
3435use crate :: physical_plan:: { with_new_children_if_necessary, Distribution , ExecutionPlan } ;
@@ -825,6 +826,35 @@ fn ensure_distribution(
825826 return Ok ( Transformed :: No ( plan) ) ;
826827 }
827828
829+ // special case for UnionExec: We want to "bubble up" hash-partitioned data. So instead of:
830+ //
831+ // Agg:
832+ // Repartition (hash):
833+ // Union:
834+ // - Agg:
835+ // Repartition (hash):
836+ // Data
837+ // - Agg:
838+ // Repartition (hash):
839+ // Data
840+ //
841+ // We can use:
842+ //
843+ // Agg:
844+ // Interleave:
845+ // - Agg:
846+ // Repartition (hash):
847+ // Data
848+ // - Agg:
849+ // Repartition (hash):
850+ // Data
851+ if let Some ( union_exec) = plan. as_any ( ) . downcast_ref :: < UnionExec > ( ) {
852+ if can_interleave ( union_exec. inputs ( ) ) {
853+ let plan = InterleaveExec :: try_new ( union_exec. inputs ( ) . clone ( ) ) ?;
854+ return Ok ( Transformed :: Yes ( Arc :: new ( plan) ) ) ;
855+ }
856+ }
857+
828858 let required_input_distributions = plan. required_input_distribution ( ) ;
829859 let children: Vec < Arc < dyn ExecutionPlan > > = plan. children ( ) ;
830860 assert_eq ! ( children. len( ) , required_input_distributions. len( ) ) ;
@@ -2134,4 +2164,42 @@ mod tests {
21342164 assert_optimized ! ( expected, exec) ;
21352165 Ok ( ( ) )
21362166 }
2167+
2168+ #[ test]
2169+ fn union_to_interleave ( ) -> Result < ( ) > {
2170+ // group by (a as a1)
2171+ let left = aggregate_exec_with_alias (
2172+ parquet_exec ( ) ,
2173+ vec ! [ ( "a" . to_string( ) , "a1" . to_string( ) ) ] ,
2174+ ) ;
2175+ // group by (a as a2)
2176+ let right = aggregate_exec_with_alias (
2177+ parquet_exec ( ) ,
2178+ vec ! [ ( "a" . to_string( ) , "a1" . to_string( ) ) ] ,
2179+ ) ;
2180+
2181+ // Union
2182+ let plan = Arc :: new ( UnionExec :: new ( vec ! [ left, right] ) ) ;
2183+
2184+ // final agg
2185+ let plan =
2186+ aggregate_exec_with_alias ( plan, vec ! [ ( "a1" . to_string( ) , "a2" . to_string( ) ) ] ) ;
2187+
2188+ // Only two RepartitionExecs added, no final RepartionExec required
2189+ let expected = & [
2190+ "AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]" ,
2191+ "AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]" ,
2192+ "InterleaveExec" ,
2193+ "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]" ,
2194+ "RepartitionExec: partitioning=Hash([Column { name: \" a1\" , index: 0 }], 10), input_partitions=1" ,
2195+ "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]" ,
2196+ "ParquetExec: limit=None, partitions={1 group: [[x]]}, projection=[a, b, c, d, e]" ,
2197+ "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]" ,
2198+ "RepartitionExec: partitioning=Hash([Column { name: \" a1\" , index: 0 }], 10), input_partitions=1" ,
2199+ "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]" ,
2200+ "ParquetExec: limit=None, partitions={1 group: [[x]]}, projection=[a, b, c, d, e]" ,
2201+ ] ;
2202+ assert_optimized ! ( expected, plan) ;
2203+ Ok ( ( ) )
2204+ }
21372205}
0 commit comments