Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,13 @@ use datafusion_common::assert_or_internal_err;
use datafusion_execution::memory_pool::proxy::VecAllocExt;
use datafusion_expr::EmitTo;

use crate::InputOrderMode;
use crate::PhysicalExpr;
use crate::aggregates::group_values::{
AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics,
GroupByMetrics, GroupValues, new_group_values,
};
use crate::aggregates::grouped_hash_stream::create_group_accumulator;
use crate::aggregates::order::GroupOrdering;
use crate::aggregates::order::{GroupCompletionMode, GroupOrdering};
use crate::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions,
evaluate_group_by,
Expand Down Expand Up @@ -113,7 +112,7 @@ impl OrderedAggregateTableMetrics {
/// `OrderedAggrMode` selects the aggregate semantics. For example,
/// `OrderedAggregateTable::<PartialMarker>::new(...)` consumes raw rows
/// and emits partial states, while
/// `OrderedAggregateTable::<FinalMarker>::new_with_input_order(...)`
/// `OrderedAggregateTable::<FinalMarker>::new_with_group_completion(...)`
/// consumes partial states and emits final values.
///
/// Shared methods live on `impl<T>`; single/partial/final behavior lives on
Expand Down Expand Up @@ -184,7 +183,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
output_schema: SchemaRef,
state_schema: SchemaRef,
batch_size: usize,
input_order_mode: &InputOrderMode,
group_completion_mode: &GroupCompletionMode,
aggregate_mode: &AggregateMode,
filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
metrics: OrderedAggregateTableMetrics,
Expand All @@ -194,7 +193,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
"OrderedAggregateTable requires config batch_size >= 1"
);

let group_ordering = GroupOrdering::try_new(input_order_mode)?;
let group_ordering = GroupOrdering::try_new_for_mode(group_completion_mode)?;
let group_schema = agg.group_by.group_schema(input_schema)?;
let group_values = new_group_values(group_schema, &group_ordering)?;
let aggregate_arguments = aggregate_expressions(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;

use crate::InputOrderMode;
use crate::aggregates::aggregate_hash_table::FinalMarker;
use crate::aggregates::order::GroupCompletionMode;
use crate::aggregates::{AggregateExec, AggregateMode, group_values::AccumulatorPhase};

use super::common::HashAggregateAccumulator;
Expand All @@ -42,12 +42,12 @@ use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}
///
/// See comments at [`OrderedAggregateTable`] for details.
impl OrderedAggregateTable<FinalMarker> {
pub(in crate::aggregates) fn new_with_input_order(
pub(in crate::aggregates) fn new_with_group_completion(
agg: &AggregateExec,
input_schema: &SchemaRef,
output_schema: SchemaRef,
batch_size: usize,
input_order_mode: &InputOrderMode,
group_completion_mode: &GroupCompletionMode,
metrics: OrderedAggregateTableMetrics,
) -> Result<Self> {
Self::new_for_mode(
Expand All @@ -56,7 +56,7 @@ impl OrderedAggregateTable<FinalMarker> {
output_schema,
Arc::clone(input_schema),
batch_size,
input_order_mode,
group_completion_mode,
&AggregateMode::Final,
vec![None; agg.aggr_expr.len()],
metrics,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ impl OrderedAggregateTable<PartialMarker> {
output_schema,
state_schema,
batch_size,
&agg.input_order_mode,
&agg.group_completion_mode,
&AggregateMode::Partial,
agg.filter_expr.iter().cloned().collect(),
metrics,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ impl OrderedAggregateTable<SingleMarker> {
output_schema,
state_schema,
batch_size,
&agg.input_order_mode,
&agg.group_completion_mode,
&agg.mode,
agg.filter_expr.iter().cloned().collect(),
metrics,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -504,7 +504,7 @@ impl GroupedHashAggregateStream {
.collect::<Vec<_>>()
.join(", ");
let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})");
let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?;
let group_ordering = GroupOrdering::try_new_for_mode(&agg.group_completion_mode)?;
let oom_mode = match (agg.mode, &group_ordering) {
// In partial aggregation mode, always prefer to emit incomplete results early.
(AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly,
Expand Down
4 changes: 3 additions & 1 deletion datafusion/physical-plan/src/aggregates/hash_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ use super::aggregate_hash_table::{
AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker,
PartialSkipMarker,
};
use super::order::GroupCompletionMode;
use super::ordered_final_stream::OrderedFinalAggregateStream;
use super::skip_partial::SkipAggregationProbe;
use crate::metrics::{
Expand Down Expand Up @@ -326,6 +327,7 @@ impl FinalSpillContext {

let mut final_agg = agg.clone();
final_agg.input_order_mode = InputOrderMode::Sorted;
final_agg.group_completion_mode = GroupCompletionMode::Full;

Ok(Self {
final_agg,
Expand Down Expand Up @@ -414,7 +416,7 @@ impl FinalSpillContext {
&context,
partition,
merged,
&InputOrderMode::Sorted,
&GroupCompletionMode::Full,
baseline_metrics.clone(),
metrics,
None,
Expand Down
Loading
Loading