Skip to content

Commit d50b923

Browse files
fix: Set Substrait aggregation phase to INITIAL_TO_RESULT
The producer left `phase` at its default on every aggregate and window function call it emits, so each one carried AGGREGATION_PHASE_UNSPECIFIED. That is not the same as omitting the field: the spec gives the value the meaning INTERMEDIATE_TO_RESULT, i.e. that the arguments are already intermediate aggregation state to be combined. A LogicalPlan::Aggregate is always a complete aggregation over its input rows. The partial/final split is a physical planning concern that the logical producer has no notion of. Both `AggregateFunction.phase` and `Expression.WindowFunction.phase` are documented as required, and as needing INITIAL_TO_RESULT for a complete invocation, so set that at both call sites. This stays invisible to a DataFusion-to-DataFusion round trip because the consumer never reads `phase`, but a consumer that honours the declaration reads a complete aggregation as one whose arguments are partial state.
1 parent 517ac6d commit d50b923

3 files changed

Lines changed: 69 additions & 3 deletions

File tree

‎datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ pub fn from_aggregate_function(
6565
true => AggregationInvocation::Distinct as i32,
6666
false => AggregationInvocation::All as i32,
6767
},
68-
phase: AggregationPhase::Unspecified as i32,
68+
phase: AggregationPhase::InitialToResult as i32,
6969
args: vec![],
7070
options: vec![],
7171
}),

‎datafusion/substrait/src/logical_plan/producer/expr/window_function.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use crate::logical_plan::producer::utils::substrait_sort_field;
2020
use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err};
2121
use datafusion::logical_expr::expr::{WindowFunction, WindowFunctionParams};
2222
use datafusion::logical_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits};
23+
use substrait::proto::AggregationPhase;
2324
use substrait::proto::aggregate_function::AggregationInvocation;
2425
use substrait::proto::expression::RexType;
2526
use substrait::proto::expression::WindowFunction as SubstraitWindowFunction;
@@ -108,7 +109,7 @@ fn make_substrait_window_function(
108109
sorts,
109110
options: vec![],
110111
output_type: None,
111-
phase: 0, // default to AGGREGATION_PHASE_UNSPECIFIED
112+
phase: AggregationPhase::InitialToResult as i32,
112113
invocation: if distinct {
113114
AggregationInvocation::Distinct as i32
114115
} else {

‎datafusion/substrait/tests/cases/serialize.rs‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ mod tests {
3535
use substrait::proto::plan_rel::RelType;
3636
use substrait::proto::rel_common::{Emit, EmitKind};
3737
use substrait::proto::r#type::{I64, Kind as TypeKind, List, Nullability, Struct};
38-
use substrait::proto::{Expression, RelCommon, Type, rel};
38+
use substrait::proto::{AggregationPhase, Expression, RelCommon, Type, rel};
3939

4040
use crate::cases::roundtrip_logical_plan::higher_order_function_ctx;
4141

@@ -321,6 +321,71 @@ mod tests {
321321
Ok(())
322322
}
323323

324+
// Collects the `phase` of every aggregate and window function call in a plan.
325+
fn collect_phases(rel_type: &rel::RelType, out: &mut Vec<i32>) {
326+
let input = match rel_type {
327+
rel::RelType::Aggregate(aggregate) => {
328+
for measure in &aggregate.measures {
329+
if let Some(function) = &measure.measure {
330+
out.push(function.phase);
331+
}
332+
}
333+
aggregate.input.as_ref()
334+
}
335+
rel::RelType::Project(project) => {
336+
for expr in &project.expressions {
337+
if let Some(RexType::WindowFunction(window)) = &expr.rex_type {
338+
out.push(window.phase);
339+
}
340+
}
341+
project.input.as_ref()
342+
}
343+
rel::RelType::Filter(filter) => filter.input.as_ref(),
344+
rel::RelType::Sort(sort) => sort.input.as_ref(),
345+
rel::RelType::Fetch(fetch) => fetch.input.as_ref(),
346+
_ => None,
347+
};
348+
if let Some(rel_type) = input.and_then(|input| input.rel_type.as_ref()) {
349+
collect_phases(rel_type, out);
350+
}
351+
}
352+
353+
/// Substrait requires `phase` on aggregate and window function calls, and
354+
/// requires `INITIAL_TO_RESULT` for a complete invocation. A DataFusion
355+
/// logical plan only ever describes complete aggregations, so that is the
356+
/// phase every produced call should carry.
357+
#[tokio::test]
358+
async fn aggregate_and_window_functions_declare_initial_to_result() -> Result<()> {
359+
let ctx = create_context().await?;
360+
361+
for sql in [
362+
"SELECT sum(a) FROM data",
363+
"SELECT a, count(*) FROM data GROUP BY a",
364+
"SELECT RANK() OVER (PARTITION BY a) FROM data",
365+
] {
366+
let plan = ctx.sql(sql).await?.into_optimized_plan()?;
367+
let proto = to_substrait_plan(&plan, &ctx.state())?;
368+
369+
let root = match proto.relations.first().unwrap().rel_type.as_ref() {
370+
Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
371+
_ => panic!("expected Root"),
372+
};
373+
let mut phases = vec![];
374+
collect_phases(root.rel_type.as_ref().unwrap(), &mut phases);
375+
376+
assert!(!phases.is_empty(), "no function call found for `{sql}`");
377+
for phase in phases {
378+
assert_eq!(
379+
phase,
380+
AggregationPhase::InitialToResult as i32,
381+
"phase for `{sql}`"
382+
);
383+
}
384+
}
385+
386+
Ok(())
387+
}
388+
324389
fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec<i32>) {
325390
assert_eq!(
326391
rel_common.unwrap().emit_kind.clone(),

0 commit comments

Comments
 (0)