@@ -203,6 +203,7 @@ use datafusion_physical_expr_common::sort_expr::{
203203use datafusion_expr:: utils:: AggregateOrderSensitivity ;
204204use datafusion_physical_expr_common:: utils:: evaluate_expressions_to_arrays;
205205use itertools:: Itertools ;
206+ use order:: GroupCompletionMode ;
206207use topk:: hash_table:: is_supported_hash_key_type;
207208use topk:: heap:: is_supported_heap_type;
208209
@@ -877,6 +878,8 @@ pub struct AggregateExec {
877878 required_input_ordering : Option < OrderingRequirements > ,
878879 /// Describes how the input is ordered relative to the group by columns
879880 input_order_mode : InputOrderMode ,
881+ /// Describes how the executor can determine that groups are complete.
882+ group_completion_mode : GroupCompletionMode ,
880883 cache : Arc < PlanProperties > ,
881884 /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]),
882885 /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node.
@@ -901,6 +904,7 @@ impl AggregateExec {
901904 required_input_ordering : self . required_input_ordering . clone ( ) ,
902905 metrics : ExecutionPlanMetricsSet :: new ( ) ,
903906 input_order_mode : self . input_order_mode . clone ( ) ,
907+ group_completion_mode : self . group_completion_mode . clone ( ) ,
904908 cache : Arc :: clone ( & self . cache ) ,
905909 mode : self . mode ,
906910 group_by : Arc :: clone ( & self . group_by ) ,
@@ -921,6 +925,7 @@ impl AggregateExec {
921925 required_input_ordering : self . required_input_ordering . clone ( ) ,
922926 metrics : ExecutionPlanMetricsSet :: new ( ) ,
923927 input_order_mode : self . input_order_mode . clone ( ) ,
928+ group_completion_mode : self . group_completion_mode . clone ( ) ,
924929 cache : Arc :: clone ( & self . cache ) ,
925930 mode : self . mode ,
926931 group_by : Arc :: clone ( & self . group_by ) ,
@@ -1043,6 +1048,8 @@ impl AggregateExec {
10431048 input_order_mode = InputOrderMode :: Linear ;
10441049 }
10451050
1051+ let group_completion_mode = GroupCompletionMode :: from ( & input_order_mode) ;
1052+
10461053 // construct a map from the input expression to the output expression of the Aggregation group by
10471054 let group_expr_mapping =
10481055 ProjectionMapping :: try_new ( group_by. expr . clone ( ) , & input. schema ( ) ) ?;
@@ -1073,6 +1080,7 @@ impl AggregateExec {
10731080 required_input_ordering,
10741081 limit_options : None ,
10751082 input_order_mode,
1083+ group_completion_mode,
10761084 cache : Arc :: new ( cache) ,
10771085 dynamic_filter : None ,
10781086 } ;
@@ -1281,7 +1289,7 @@ impl AggregateExec {
12811289
12821290 fn should_use_partial_hash_stream ( & self , _context : & TaskContext ) -> bool {
12831291 self . mode == AggregateMode :: Partial
1284- && self . input_order_mode == InputOrderMode :: Linear
1292+ && self . group_completion_mode == GroupCompletionMode :: None
12851293 && !self . group_by . is_true_no_grouping ( )
12861294 && self . group_by . is_single ( )
12871295 && self . limit_options_supported_by_hash_stream ( )
@@ -1292,7 +1300,7 @@ impl AggregateExec {
12921300 _context : & TaskContext ,
12931301 ) -> bool {
12941302 self . mode == AggregateMode :: Partial
1295- && self . input_order_mode != InputOrderMode :: Linear
1303+ && self . group_completion_mode != GroupCompletionMode :: None
12961304 && !self . group_by . is_true_no_grouping ( )
12971305 && self . group_by . is_single ( )
12981306 && self . limit_options_supported_by_hash_stream ( )
@@ -1303,7 +1311,7 @@ impl AggregateExec {
13031311 self . mode,
13041312 AggregateMode :: Final | AggregateMode :: FinalPartitioned
13051313 ) && self . limit_options_supported_by_hash_stream ( )
1306- && self . input_order_mode == InputOrderMode :: Linear
1314+ && self . group_completion_mode == GroupCompletionMode :: None
13071315 && !self . group_by . is_true_no_grouping ( )
13081316 && self . group_by . is_single ( )
13091317 }
@@ -1316,7 +1324,7 @@ impl AggregateExec {
13161324
13171325 self . mode == AggregateMode :: PartialReduce
13181326 && self . limit_options . is_none ( )
1319- && self . input_order_mode == InputOrderMode :: Linear
1327+ && self . group_completion_mode == GroupCompletionMode :: None
13201328 && !self . group_by . is_true_no_grouping ( )
13211329 && self . group_by . is_single ( )
13221330 }
@@ -1326,7 +1334,7 @@ impl AggregateExec {
13261334 self . mode,
13271335 AggregateMode :: Single | AggregateMode :: SinglePartitioned
13281336 ) && self . limit_options . is_none ( )
1329- && self . input_order_mode == InputOrderMode :: Linear
1337+ && self . group_completion_mode == GroupCompletionMode :: None
13301338 && !self . group_by . is_true_no_grouping ( )
13311339 && self . group_by . is_single ( )
13321340 }
@@ -1336,7 +1344,7 @@ impl AggregateExec {
13361344 self . mode,
13371345 AggregateMode :: Single | AggregateMode :: SinglePartitioned
13381346 ) && self . limit_options . is_none ( )
1339- && self . input_order_mode != InputOrderMode :: Linear
1347+ && self . group_completion_mode != GroupCompletionMode :: None
13401348 && !self . group_by . is_true_no_grouping ( )
13411349 && self . group_by . is_single ( )
13421350 }
@@ -1346,7 +1354,7 @@ impl AggregateExec {
13461354 self . mode,
13471355 AggregateMode :: Final | AggregateMode :: FinalPartitioned
13481356 ) && self . limit_options_supported_by_hash_stream ( )
1349- && self . input_order_mode != InputOrderMode :: Linear
1357+ && self . group_completion_mode != GroupCompletionMode :: None
13501358 && !self . group_by . is_true_no_grouping ( )
13511359 && self . group_by . is_single ( )
13521360 }
@@ -2354,6 +2362,8 @@ impl ExecutionPlan for AggregateExec {
23542362 required_input_ordering : _,
23552363 // Derived at construction from the input ordering and `group_by`.
23562364 input_order_mode : _,
2365+ // Derived at construction from `input_order_mode`.
2366+ group_completion_mode : _,
23572367 // Derived at construction by `Self::compute_properties`.
23582368 cache : _,
23592369 dynamic_filter,
@@ -4592,6 +4602,7 @@ mod tests {
45924602 ) ?;
45934603
45944604 assert_eq ! ( aggregate. input_order_mode( ) , & InputOrderMode :: Linear ) ;
4605+ assert_eq ! ( aggregate. group_completion_mode, GroupCompletionMode :: None ) ;
45954606 // This captures the behavior before #24438. When the source can declare
45964607 // `(key, time_bin)` group-contiguous, the corresponding case can use
45974608 // `EmissionType::Incremental`.
0 commit comments