From d50b923f3978d85c6e547aaf952db64799af4e4c Mon Sep 17 00:00:00 2001 From: naman Date: Thu, 10 Sep 2026 14:49:19 +0530 Subject: [PATCH] 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. --- .../producer/expr/aggregate_function.rs | 2 +- .../producer/expr/window_function.rs | 3 +- datafusion/substrait/tests/cases/serialize.rs | 67 ++++++++++++++++++- 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs index 3713f8934f19f..d96f33f49f108 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/aggregate_function.rs @@ -65,7 +65,7 @@ pub fn from_aggregate_function( true => AggregationInvocation::Distinct as i32, false => AggregationInvocation::All as i32, }, - phase: AggregationPhase::Unspecified as i32, + phase: AggregationPhase::InitialToResult as i32, args: vec![], options: vec![], }), diff --git a/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs index d35771bf099d3..f449b8100c34d 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs @@ -20,6 +20,7 @@ use crate::logical_plan::producer::utils::substrait_sort_field; use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err}; use datafusion::logical_expr::expr::{WindowFunction, WindowFunctionParams}; use datafusion::logical_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits}; +use substrait::proto::AggregationPhase; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::expression::RexType; use substrait::proto::expression::WindowFunction as SubstraitWindowFunction; @@ -108,7 +109,7 @@ fn make_substrait_window_function( sorts, options: vec![], output_type: None, - phase: 0, // default to AGGREGATION_PHASE_UNSPECIFIED + phase: AggregationPhase::InitialToResult as i32, invocation: if distinct { AggregationInvocation::Distinct as i32 } else { diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 4a8413718edb9..809ed612a66d7 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -35,7 +35,7 @@ mod tests { use substrait::proto::plan_rel::RelType; use substrait::proto::rel_common::{Emit, EmitKind}; use substrait::proto::r#type::{I64, Kind as TypeKind, List, Nullability, Struct}; - use substrait::proto::{Expression, RelCommon, Type, rel}; + use substrait::proto::{AggregationPhase, Expression, RelCommon, Type, rel}; use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; @@ -321,6 +321,71 @@ mod tests { Ok(()) } + // Collects the `phase` of every aggregate and window function call in a plan. + fn collect_phases(rel_type: &rel::RelType, out: &mut Vec) { + let input = match rel_type { + rel::RelType::Aggregate(aggregate) => { + for measure in &aggregate.measures { + if let Some(function) = &measure.measure { + out.push(function.phase); + } + } + aggregate.input.as_ref() + } + rel::RelType::Project(project) => { + for expr in &project.expressions { + if let Some(RexType::WindowFunction(window)) = &expr.rex_type { + out.push(window.phase); + } + } + project.input.as_ref() + } + rel::RelType::Filter(filter) => filter.input.as_ref(), + rel::RelType::Sort(sort) => sort.input.as_ref(), + rel::RelType::Fetch(fetch) => fetch.input.as_ref(), + _ => None, + }; + if let Some(rel_type) = input.and_then(|input| input.rel_type.as_ref()) { + collect_phases(rel_type, out); + } + } + + /// Substrait requires `phase` on aggregate and window function calls, and + /// requires `INITIAL_TO_RESULT` for a complete invocation. A DataFusion + /// logical plan only ever describes complete aggregations, so that is the + /// phase every produced call should carry. + #[tokio::test] + async fn aggregate_and_window_functions_declare_initial_to_result() -> Result<()> { + let ctx = create_context().await?; + + for sql in [ + "SELECT sum(a) FROM data", + "SELECT a, count(*) FROM data GROUP BY a", + "SELECT RANK() OVER (PARTITION BY a) FROM data", + ] { + let plan = ctx.sql(sql).await?.into_optimized_plan()?; + let proto = to_substrait_plan(&plan, &ctx.state())?; + + let root = match proto.relations.first().unwrap().rel_type.as_ref() { + Some(RelType::Root(root)) => root.input.as_ref().unwrap(), + _ => panic!("expected Root"), + }; + let mut phases = vec![]; + collect_phases(root.rel_type.as_ref().unwrap(), &mut phases); + + assert!(!phases.is_empty(), "no function call found for `{sql}`"); + for phase in phases { + assert_eq!( + phase, + AggregationPhase::InitialToResult as i32, + "phase for `{sql}`" + ); + } + } + + Ok(()) + } + fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec) { assert_eq!( rel_common.unwrap().emit_kind.clone(),