Skip to content

Commit 0267ad1

Browse files
committed
refactor: separate aggregate group completion from input ordering
1 parent 5dfc10b commit 0267ad1

12 files changed

Lines changed: 109 additions & 54 deletions

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,14 +28,13 @@ use datafusion_common::assert_or_internal_err;
2828
use datafusion_execution::memory_pool::proxy::VecAllocExt;
2929
use datafusion_expr::EmitTo;
3030

31-
use crate::InputOrderMode;
3231
use crate::PhysicalExpr;
3332
use crate::aggregates::group_values::{
3433
AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics,
3534
GroupByMetrics, GroupValues, new_group_values,
3635
};
3736
use crate::aggregates::grouped_hash_stream::create_group_accumulator;
38-
use crate::aggregates::order::GroupOrdering;
37+
use crate::aggregates::order::{GroupCompletionMode, GroupOrdering};
3938
use crate::aggregates::{
4039
AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions,
4140
evaluate_group_by,
@@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics {
113112
/// `OrderedAggrMode` selects the aggregate semantics. For example,
114113
/// `OrderedAggregateTable::<PartialMarker>::new(...)` consumes raw rows
115114
/// and emits partial states, while
116-
/// `OrderedAggregateTable::<FinalMarker>::new_with_input_order(...)`
115+
/// `OrderedAggregateTable::<FinalMarker>::new_with_group_completion(...)`
117116
/// consumes partial states and emits final values.
118117
///
119118
/// Shared methods live on `impl<T>`; single/partial/final behavior lives on
@@ -184,7 +183,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
184183
output_schema: SchemaRef,
185184
state_schema: SchemaRef,
186185
batch_size: usize,
187-
input_order_mode: &InputOrderMode,
186+
group_completion_mode: &GroupCompletionMode,
188187
aggregate_mode: &AggregateMode,
189188
filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
190189
metrics: OrderedAggregateTableMetrics,
@@ -194,7 +193,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
194193
"OrderedAggregateTable requires config batch_size >= 1"
195194
);
196195

197-
let group_ordering = GroupOrdering::try_new(input_order_mode)?;
196+
let group_ordering =
197+
GroupOrdering::try_new_for_group_completion(group_completion_mode)?;
198198
let group_schema = agg.group_by.group_schema(input_schema)?;
199199
let group_values = new_group_values(group_schema, &group_ordering)?;
200200
let aggregate_arguments = aggregate_expressions(

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@ use arrow::datatypes::SchemaRef;
2525
use arrow::record_batch::RecordBatch;
2626
use datafusion_common::Result;
2727

28-
use crate::InputOrderMode;
2928
use crate::aggregates::aggregate_hash_table::FinalMarker;
29+
use crate::aggregates::order::GroupCompletionMode;
3030
use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase};
3131

3232
use super::common::HashAggregateAccumulator;
@@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}
4242
///
4343
/// See comments at [`OrderedAggregateTable`] for details.
4444
impl OrderedAggregateTable<FinalMarker> {
45-
pub(in crate::aggregates) fn new_with_input_order(
45+
pub(in crate::aggregates) fn new_with_group_completion(
4646
agg: &AggregateExec,
4747
input_schema: &SchemaRef,
4848
output_schema: SchemaRef,
4949
batch_size: usize,
50-
input_order_mode: &InputOrderMode,
50+
group_completion_mode: &GroupCompletionMode,
5151
metrics: OrderedAggregateTableMetrics,
5252
) -> Result<Self> {
5353
Self::new_for_mode(
@@ -56,7 +56,7 @@ impl OrderedAggregateTable<FinalMarker> {
5656
output_schema,
5757
Arc::clone(input_schema),
5858
batch_size,
59-
input_order_mode,
59+
group_completion_mode,
6060
&AggregateMode::Final,
6161
vec![None; agg.aggr_expr.len()],
6262
metrics,

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ impl OrderedAggregateTable<PartialMarker> {
6868
output_schema,
6969
state_schema,
7070
batch_size,
71-
&agg.input_order_mode,
71+
&agg.group_completion_mode,
7272
&AggregateMode::Partial,
7373
agg.filter_expr.iter().cloned().collect(),
7474
metrics,

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ impl OrderedAggregateTable<SingleMarker> {
5959
output_schema,
6060
state_schema,
6161
batch_size,
62-
&agg.input_order_mode,
62+
&agg.group_completion_mode,
6363
&agg.mode,
6464
agg.filter_expr.iter().cloned().collect(),
6565
metrics,

‎datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,8 @@ impl GroupedHashAggregateStream {
518518
.collect::<Vec<_>>()
519519
.join(", ");
520520
let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})");
521-
let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?;
521+
let group_ordering =
522+
GroupOrdering::try_new_for_group_completion(&agg.group_completion_mode)?;
522523
let oom_mode = match (agg.mode, &group_ordering) {
523524
// In partial aggregation mode, always prefer to emit incomplete results early.
524525
(AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly,

‎datafusion/physical-plan/src/aggregates/hash_stream.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ use super::aggregate_hash_table::{
4545
AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker,
4646
PartialSkipMarker,
4747
};
48+
use super::order::GroupCompletionMode;
4849
use super::ordered_final_stream::OrderedFinalAggregateStream;
4950
use super::skip_partial::SkipAggregationProbe;
5051
use crate::metrics::{
@@ -326,6 +327,7 @@ impl FinalSpillContext {
326327

327328
let mut final_agg = agg.clone();
328329
final_agg.input_order_mode = InputOrderMode::Sorted;
330+
final_agg.group_completion_mode = GroupCompletionMode::Full;
329331

330332
Ok(Self {
331333
final_agg,
@@ -414,7 +416,7 @@ impl FinalSpillContext {
414416
&context,
415417
partition,
416418
merged,
417-
&InputOrderMode::Sorted,
419+
&GroupCompletionMode::Full,
418420
baseline_metrics.clone(),
419421
metrics,
420422
None,

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,7 @@ use datafusion_physical_expr_common::sort_expr::{
203203
use datafusion_expr::utils::AggregateOrderSensitivity;
204204
use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
205205
use itertools::Itertools;
206+
use order::GroupCompletionMode;
206207
use topk::hash_table::is_supported_hash_key_type;
207208
use 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`.

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

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,33 @@ use crate::InputOrderMode;
2828
pub use full::GroupOrderingFull;
2929
pub use partial::GroupOrderingPartial;
3030

31+
/// Describes how an aggregate can determine that groups are complete.
32+
///
33+
/// This is distinct from [`InputOrderMode`], which describes the ordering of
34+
/// the input relative to the grouping expressions. Input ordering is one way
35+
/// to establish a group-completion mode, but the execution machinery only
36+
/// needs to know when it can safely emit completed groups.
37+
#[derive(Clone, Debug, PartialEq, Eq)]
38+
pub(crate) enum GroupCompletionMode {
39+
/// Groups cannot be completed before the input ends.
40+
None,
41+
/// Groups sharing the values at these grouping-expression indices form a
42+
/// contiguous range.
43+
Partial(Vec<usize>),
44+
/// Every complete grouping tuple forms a contiguous range.
45+
Full,
46+
}
47+
48+
impl From<&InputOrderMode> for GroupCompletionMode {
49+
fn from(value: &InputOrderMode) -> Self {
50+
match value {
51+
InputOrderMode::Linear => Self::None,
52+
InputOrderMode::PartiallySorted(indices) => Self::Partial(indices.clone()),
53+
InputOrderMode::Sorted => Self::Full,
54+
}
55+
}
56+
}
57+
3158
/// Ordering information for each group in the hash table
3259
#[derive(Debug)]
3360
pub enum GroupOrdering {
@@ -40,15 +67,24 @@ pub enum GroupOrdering {
4067
}
4168

4269
impl GroupOrdering {
43-
/// Create a `GroupOrdering` for the specified ordering
70+
/// Create a `GroupOrdering` for the specified input order mode.
4471
pub fn try_new(mode: &InputOrderMode) -> Result<Self> {
72+
Self::try_new_for_group_completion(&GroupCompletionMode::from(mode))
73+
}
74+
75+
/// Create a `GroupOrdering` for the specified group-completion mode.
76+
pub(crate) fn try_new_for_group_completion(
77+
mode: &GroupCompletionMode,
78+
) -> Result<Self> {
4579
match mode {
46-
InputOrderMode::Linear => Ok(GroupOrdering::None),
47-
InputOrderMode::PartiallySorted(order_indices) => {
80+
GroupCompletionMode::None => Ok(GroupOrdering::None),
81+
GroupCompletionMode::Partial(order_indices) => {
4882
GroupOrderingPartial::try_new(order_indices.clone())
4983
.map(GroupOrdering::Partial)
5084
}
51-
InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())),
85+
GroupCompletionMode::Full => {
86+
Ok(GroupOrdering::Full(GroupOrderingFull::new()))
87+
}
5288
}
5389
}
5490

0 commit comments

Comments
 (0)