Skip to content

Commit 64a6e55

Browse files
committed
refactor: separate aggregate group completion from input ordering
1 parent 0e9cafe commit 64a6e55

12 files changed

Lines changed: 199 additions & 100 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
@@ -523,7 +523,8 @@ impl GroupedHashAggregateStream {
523523
.collect::<Vec<_>>()
524524
.join(", ");
525525
let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})");
526-
let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?;
526+
let group_ordering =
527+
GroupOrdering::try_new_for_group_completion(&agg.group_completion_mode)?;
527528
let oom_mode = match (agg.mode, &group_ordering) {
528529
// In partial aggregation mode, always prefer to emit incomplete results early.
529530
(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
@@ -41,6 +41,7 @@ use super::aggregate_hash_table::{
4141
AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker,
4242
PartialSkipMarker,
4343
};
44+
use super::order::GroupCompletionMode;
4445
use super::ordered_final_stream::OrderedFinalAggregateStream;
4546
use super::skip_partial::SkipAggregationProbe;
4647
use crate::metrics::{
@@ -266,6 +267,7 @@ impl FinalSpillContext {
266267

267268
let mut final_agg = agg.clone();
268269
final_agg.input_order_mode = InputOrderMode::Sorted;
270+
final_agg.group_completion_mode = GroupCompletionMode::Full;
269271

270272
Ok(Self {
271273
final_agg,
@@ -354,7 +356,7 @@ impl FinalSpillContext {
354356
&context,
355357
partition,
356358
merged,
357-
&InputOrderMode::Sorted,
359+
&GroupCompletionMode::Full,
358360
baseline_metrics.clone(),
359361
metrics,
360362
None,

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

Lines changed: 50 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,7 @@ use datafusion_physical_expr_common::sort_expr::{
206206
use datafusion_expr::utils::AggregateOrderSensitivity;
207207
use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
208208
use itertools::Itertools;
209+
use order::GroupCompletionMode;
209210
use topk::hash_table::is_supported_hash_key_type;
210211
use 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)?;

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

Lines changed: 57 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,50 @@ 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+
///
38+
/// For example, when grouping by `key`, both inputs have fully contiguous
39+
/// groups within the input partition:
40+
///
41+
/// ```text
42+
/// sorted: A A B B C C
43+
/// not sorted: C C A A B B
44+
/// ```
45+
///
46+
/// In both cases, once the key changes, the previous key will not appear again,
47+
/// so its group is complete and can be emitted.
48+
#[derive(Clone, Debug, PartialEq, Eq)]
49+
pub(crate) enum GroupCompletionMode {
50+
/// No group can be known complete before the input ends.
51+
None,
52+
/// Rows with the same values at these grouping-expression indices form one
53+
/// contiguous range. When those values change, every group in the previous
54+
/// range is complete and can be emitted.
55+
///
56+
/// For example, with `GROUP BY (a, b)`, `Partial(vec![0])` means all rows
57+
/// for each value of `a` are contiguous, while an `(a, b)` tuple may recur
58+
/// within that range.
59+
Partial(Vec<usize>),
60+
/// Rows with the same complete grouping tuple form one contiguous range.
61+
/// When the tuple changes, the previous group can be emitted.
62+
Full,
63+
}
64+
65+
impl From<&InputOrderMode> for GroupCompletionMode {
66+
fn from(value: &InputOrderMode) -> Self {
67+
match value {
68+
InputOrderMode::Linear => Self::None,
69+
InputOrderMode::PartiallySorted(indices) => Self::Partial(indices.clone()),
70+
InputOrderMode::Sorted => Self::Full,
71+
}
72+
}
73+
}
74+
3175
/// Ordering information for each group in the hash table
3276
#[derive(Debug)]
3377
pub enum GroupOrdering {
@@ -40,15 +84,24 @@ pub enum GroupOrdering {
4084
}
4185

4286
impl GroupOrdering {
43-
/// Create a `GroupOrdering` for the specified ordering
87+
/// Create a `GroupOrdering` for the specified input order mode.
4488
pub fn try_new(mode: &InputOrderMode) -> Result<Self> {
89+
Self::try_new_for_group_completion(&GroupCompletionMode::from(mode))
90+
}
91+
92+
/// Create a `GroupOrdering` for the specified group-completion mode.
93+
pub(crate) fn try_new_for_group_completion(
94+
mode: &GroupCompletionMode,
95+
) -> Result<Self> {
4596
match mode {
46-
InputOrderMode::Linear => Ok(GroupOrdering::None),
47-
InputOrderMode::PartiallySorted(order_indices) => {
97+
GroupCompletionMode::None => Ok(GroupOrdering::None),
98+
GroupCompletionMode::Partial(order_indices) => {
4899
GroupOrderingPartial::try_new(order_indices.clone())
49100
.map(GroupOrdering::Partial)
50101
}
51-
InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())),
102+
GroupCompletionMode::Full => {
103+
Ok(GroupOrdering::Full(GroupOrderingFull::new()))
104+
}
52105
}
53106
}
54107

0 commit comments

Comments
 (0)