@@ -206,6 +206,7 @@ use datafusion_physical_expr_common::sort_expr::{
206206use datafusion_expr:: utils:: AggregateOrderSensitivity ;
207207use datafusion_physical_expr_common:: utils:: evaluate_expressions_to_arrays;
208208use itertools:: Itertools ;
209+ use order:: GroupCompletionMode ;
209210use topk:: hash_table:: is_supported_hash_key_type;
210211use topk:: heap:: is_supported_heap_type;
211212
@@ -894,6 +895,12 @@ pub struct AggregateExec {
894895 /// input order". When that is not possible, the constructor overwrites it
895896 /// with the unordered variant [`InputOrderMode::Linear`].
896897 input_order_mode : InputOrderMode ,
898+ /// Describes when the executor can determine that groups are complete.
899+ ///
900+ /// Input ordering describes a subset of the cases in which groups can be
901+ /// safely emitted before the input ends. Full group completion requires only
902+ /// that rows for each complete grouping tuple are contiguous.
903+ group_completion_mode : GroupCompletionMode ,
897904 cache : Arc < PlanProperties > ,
898905 /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]),
899906 /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node.
@@ -918,6 +925,7 @@ impl AggregateExec {
918925 required_input_ordering : self . required_input_ordering . clone ( ) ,
919926 metrics : ExecutionPlanMetricsSet :: new ( ) ,
920927 input_order_mode : self . input_order_mode . clone ( ) ,
928+ group_completion_mode : self . group_completion_mode . clone ( ) ,
921929 cache : Arc :: clone ( & self . cache ) ,
922930 mode : self . mode ,
923931 group_by : Arc :: clone ( & self . group_by ) ,
@@ -938,6 +946,7 @@ impl AggregateExec {
938946 required_input_ordering : self . required_input_ordering . clone ( ) ,
939947 metrics : ExecutionPlanMetricsSet :: new ( ) ,
940948 input_order_mode : self . input_order_mode . clone ( ) ,
949+ group_completion_mode : self . group_completion_mode . clone ( ) ,
941950 cache : Arc :: clone ( & self . cache ) ,
942951 mode : self . mode ,
943952 group_by : Arc :: clone ( & self . group_by ) ,
@@ -1060,6 +1069,8 @@ impl AggregateExec {
10601069 input_order_mode = InputOrderMode :: Linear ;
10611070 }
10621071
1072+ let group_completion_mode = GroupCompletionMode :: from ( & input_order_mode) ;
1073+
10631074 // construct a map from the input expression to the output expression of the Aggregation group by
10641075 let group_expr_mapping =
10651076 ProjectionMapping :: try_new ( group_by. expr . clone ( ) , & input. schema ( ) ) ?;
@@ -1090,6 +1101,7 @@ impl AggregateExec {
10901101 required_input_ordering,
10911102 limit_options : None ,
10921103 input_order_mode,
1104+ group_completion_mode,
10931105 cache : Arc :: new ( cache) ,
10941106 dynamic_filter : None ,
10951107 } ;
@@ -1272,47 +1284,51 @@ impl AggregateExec {
12721284 ) ;
12731285 }
12741286
1275- // Choose the execution path based on (aggregation mode, ordering).
1276- //
1277- // Note that `self.input_order_mode` represents both input ordering and output
1278- // order promise. See its comment for details.
1287+ // Choose the execution path based on aggregation mode and when groups
1288+ // are known to be complete.
12791289 use AggregateMode :: * ;
1280- use InputOrderMode :: * ;
1281- let stream = match ( self . mode , & self . input_order_mode ) {
1282- ( Partial , Linear ) => StreamType :: PartialHash (
1290+ let stream = match ( self . mode , & self . group_completion_mode ) {
1291+ ( Partial , GroupCompletionMode :: None ) => StreamType :: PartialHash (
12831292 PartialHashAggregateStream :: new ( self , context, partition) ?,
12841293 ) ,
1285- ( Partial , Sorted | PartiallySorted ( _) ) => {
1294+ ( Partial , GroupCompletionMode :: Partial ( _) | GroupCompletionMode :: Full ) => {
12861295 StreamType :: OrderedPartialAggregate ( OrderedPartialAggregateStream :: new (
12871296 self , context, partition,
12881297 ) ?)
12891298 }
1290- ( PartialReduce , Linear ) => StreamType :: PartialReduceHash (
1299+ ( PartialReduce , GroupCompletionMode :: None ) => StreamType :: PartialReduceHash (
12911300 PartialReduceHashAggregateStream :: new ( self , context, partition) ?,
12921301 ) ,
1293- ( PartialReduce , Sorted | PartiallySorted ( _) ) => {
1294- // See the comment above: the builder enforces `Linear` order for
1295- // `PartialReduce` mode.
1302+ (
1303+ PartialReduce ,
1304+ GroupCompletionMode :: Partial ( _) | GroupCompletionMode :: Full ,
1305+ ) => {
12961306 return internal_err ! (
1297- "PartialReduce aggregation must use InputOrderMode::Linear "
1307+ "PartialReduce aggregation must use GroupCompletionMode::None "
12981308 ) ;
12991309 }
1300- ( Final | FinalPartitioned , Linear ) => StreamType :: FinalHash (
1301- FinalHashAggregateStream :: new ( self , context, partition) ?,
1302- ) ,
1303- ( Final | FinalPartitioned , Sorted | PartiallySorted ( _) ) => {
1304- StreamType :: OrderedFinalAggregate ( OrderedFinalAggregateStream :: new (
1310+ ( Final | FinalPartitioned , GroupCompletionMode :: None ) => {
1311+ StreamType :: FinalHash ( FinalHashAggregateStream :: new (
13051312 self , context, partition,
13061313 ) ?)
13071314 }
1308- ( Single | SinglePartitioned , Linear ) => StreamType :: SingleHash (
1309- SingleHashAggregateStream :: new ( self , context, partition) ?,
1310- ) ,
1311- ( Single | SinglePartitioned , Sorted | PartiallySorted ( _) ) => {
1312- StreamType :: OrderedSingleAggregate ( OrderedSingleAggregateStream :: new (
1315+ (
1316+ Final | FinalPartitioned ,
1317+ GroupCompletionMode :: Partial ( _) | GroupCompletionMode :: Full ,
1318+ ) => StreamType :: OrderedFinalAggregate ( OrderedFinalAggregateStream :: new (
1319+ self , context, partition,
1320+ ) ?) ,
1321+ ( Single | SinglePartitioned , GroupCompletionMode :: None ) => {
1322+ StreamType :: SingleHash ( SingleHashAggregateStream :: new (
13131323 self , context, partition,
13141324 ) ?)
13151325 }
1326+ (
1327+ Single | SinglePartitioned ,
1328+ GroupCompletionMode :: Partial ( _) | GroupCompletionMode :: Full ,
1329+ ) => StreamType :: OrderedSingleAggregate ( OrderedSingleAggregateStream :: new (
1330+ self , context, partition,
1331+ ) ?) ,
13161332 } ;
13171333 Ok ( stream)
13181334 }
@@ -2340,6 +2356,8 @@ impl ExecutionPlan for AggregateExec {
23402356 required_input_ordering : _,
23412357 // Derived at construction from the input ordering and `group_by`.
23422358 input_order_mode : _,
2359+ // Derived at construction from `input_order_mode`.
2360+ group_completion_mode : _,
23432361 // Derived at construction by `Self::compute_properties`.
23442362 cache : _,
23452363 dynamic_filter,
@@ -4630,6 +4648,10 @@ mod tests {
46304648 aggregate. input_order_mode( ) ,
46314649 InputOrderMode :: PartiallySorted ( _)
46324650 ) ) ;
4651+ assert_eq ! (
4652+ aggregate. group_completion_mode,
4653+ GroupCompletionMode :: Partial ( vec![ 0 ] )
4654+ ) ;
46334655
46344656 let task_ctx = new_migrated_hash_ctx ( 2 ) ;
46354657 let stream = aggregate. execute_typed ( 0 , & task_ctx) ?;
@@ -4998,6 +5020,7 @@ mod tests {
49985020 ) ?;
49995021
50005022 assert_eq ! ( aggregate. input_order_mode( ) , & InputOrderMode :: Linear ) ;
5023+ assert_eq ! ( aggregate. group_completion_mode, GroupCompletionMode :: None ) ;
50015024 // This captures the behavior before #24438. When the source can declare
50025025 // `(key, time_bin)` group-contiguous, the corresponding case can use
50035026 // `EmissionType::Incremental`.
@@ -5158,6 +5181,10 @@ mod tests {
51585181 Arc :: clone ( & schema) ,
51595182 ) ?;
51605183 assert_eq ! ( final_aggregate. input_order_mode( ) , & InputOrderMode :: Sorted ) ;
5184+ assert_eq ! (
5185+ final_aggregate. group_completion_mode,
5186+ GroupCompletionMode :: Full
5187+ ) ;
51615188
51625189 let task_ctx = new_migrated_hash_ctx ( 2 ) ;
51635190 let stream = final_aggregate. execute_typed ( 0 , & task_ctx) ?;
0 commit comments