From 6e8b7b30ab07596cd8c4c87a0d8d79f02efa3c43 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Mon, 7 Sep 2026 14:38:59 -0700 Subject: [PATCH 1/2] fix: reject unsupported Substrait aggregation phases --- .../consumer/expr/aggregate_function.rs | 19 +- .../consumer/expr/window_function.rs | 2 + .../tests/cases/aggregation_tests.rs | 193 +++++++++++++++++- 3 files changed, 211 insertions(+), 3 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs index 096eef7ae3b0e..32944a4ae1930 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs @@ -18,11 +18,25 @@ use crate::logical_plan::consumer::{ SubstraitConsumer, from_substrait_func_args, substrait_fun_name, }; -use datafusion::common::{DFSchema, ScalarValue, not_impl_datafusion_err, plan_err}; +use datafusion::common::{ + DFSchema, ScalarValue, not_impl_datafusion_err, not_impl_err, plan_datafusion_err, + plan_err, +}; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::{Expr, SortExpr, expr}; use std::sync::Arc; -use substrait::proto::AggregateFunction; +use substrait::proto::{AggregateFunction, AggregationPhase}; + +pub(super) fn validate_aggregation_phase(phase: i32) -> datafusion::common::Result<()> { + match AggregationPhase::try_from(phase) + .map_err(|e| plan_datafusion_err!("Invalid aggregation phase {phase}: {e}"))? + { + // Logical plans represent complete calls. Keep accepting unspecified phases + // for compatibility with existing DataFusion aggregate and window producers. + AggregationPhase::Unspecified | AggregationPhase::InitialToResult => Ok(()), + phase => not_impl_err!("Unsupported aggregation phase: {}", phase.as_str_name()), + } +} /// Convert Substrait AggregateFunction to DataFusion Expr pub async fn from_substrait_agg_func( @@ -33,6 +47,7 @@ pub async fn from_substrait_agg_func( order_by: Vec, distinct: bool, ) -> datafusion::common::Result> { + validate_aggregation_phase(f.phase)?; let Some(fn_signature) = consumer .get_extensions() .functions diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs index d39b325a54827..bd1a63b7040f5 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use super::aggregate_function::validate_aggregation_phase; use crate::logical_plan::consumer::{ SubstraitConsumer, from_substrait_func_args, from_substrait_rex_vec, from_substrait_sorts, substrait_fun_name, @@ -39,6 +40,7 @@ pub async fn from_window_function( window: &WindowFunction, input_schema: &DFSchema, ) -> datafusion::common::Result { + validate_aggregation_phase(window.phase)?; let Some(fn_signature) = consumer .get_extensions() .functions diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs index e572023f17a92..8af4cb7fa5c39 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -20,11 +20,202 @@ #[cfg(test)] mod tests { use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; - use datafusion::common::Result; + use datafusion::arrow::array::record_batch; + use datafusion::arrow::datatypes as arrow_schema; + use datafusion::common::{Result, ScalarValue, TableReference}; use datafusion::dataframe::DataFrame; use datafusion::prelude::SessionContext; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use insta::assert_snapshot; + use prost::Message; + use serde_json::json; + use substrait::proto::{AggregationPhase, Plan, expression, plan_rel, rel}; + + fn aggregate_phase_plan(phase: i32, rooted: bool) -> Plan { + let i64_type = json!({"i64": {"nullability": "NULLABILITY_REQUIRED"}}); + let output_type = if phase == AggregationPhase::InitialToIntermediate as i32 { + json!({"struct": {"types": [i64_type.clone(), i64_type.clone()], "nullability": "NULLABILITY_REQUIRED"}}) + } else { + json!({"fp64": {"nullability": "NULLABILITY_NULLABLE"}}) + }; + let rel = json!({"aggregate": { + "input": {"read": { + "baseSchema": {"names": ["c0"], "struct": { + "types": [i64_type], "nullability": "NULLABILITY_REQUIRED" + }}, + "namedTable": {"names": ["t_avg"]} + }}, + "measures": [{"measure": { + "functionReference": 1, + "outputType": output_type, + "arguments": [{"value": {"selection": { + "directReference": {"structField": {}}, "rootReference": {} + }}}] + }}] + }}); + let relation = if rooted { + let names = if phase == AggregationPhase::InitialToIntermediate as i32 { + vec!["average", "sum", "count"] + } else { + vec!["average"] + }; + json!({"root": {"input": rel, "names": names}}) + } else { + json!({"rel": rel}) + }; + let mut plan: Plan = serde_json::from_value(json!({ + "extensions": [{"extensionFunction": {"functionAnchor": 1, "name": "avg:i64"}}], + "relations": [relation] + })) + .unwrap(); + let relation = match plan.relations[0].rel_type.as_mut().unwrap() { + plan_rel::RelType::Rel(rel) => rel, + plan_rel::RelType::Root(root) => root.input.as_mut().unwrap(), + }; + let Some(rel::RelType::Aggregate(aggregate)) = relation.rel_type.as_mut() else { + panic!("expected aggregate"); + }; + aggregate.measures[0].measure.as_mut().unwrap().phase = phase; + Plan::decode(plan.encode_to_vec().as_slice()).unwrap() + } + + async fn aggregate_phase_context() -> Result { + let ctx = SessionContext::new(); + ctx.sql("CREATE TABLE t_avg AS SELECT column1 AS c0 FROM (VALUES (1::BIGINT), (2::BIGINT))") + .await? + .collect() + .await?; + Ok(ctx) + } + + #[tokio::test] + async fn aggregate_supported_phases() -> Result<()> { + let ctx = aggregate_phase_context().await?; + for phase in [ + AggregationPhase::Unspecified, + AggregationPhase::InitialToResult, + ] { + for rooted in [false, true] { + let proto = aggregate_phase_plan(phase as i32, rooted); + let plan = from_substrait_plan(&ctx.state(), &proto).await?; + let batches = DataFrame::new(ctx.state(), plan).collect().await?; + assert_eq!( + ScalarValue::try_from_array(batches[0].column(0), 0)?, + ScalarValue::Float64(Some(1.5)) + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn aggregate_unsupported_phases() -> Result<()> { + let ctx = aggregate_phase_context().await?; + for phase in [ + AggregationPhase::InitialToIntermediate, + AggregationPhase::IntermediateToIntermediate, + AggregationPhase::IntermediateToResult, + ] { + for rooted in [false, true] { + let proto = aggregate_phase_plan(phase as i32, rooted); + let err = from_substrait_plan(&ctx.state(), &proto).await.unwrap_err(); + assert!( + err.to_string().contains(&format!( + "Unsupported aggregation phase: {}", + phase.as_str_name() + )), + "{err}" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn aggregate_invalid_phase() -> Result<()> { + let ctx = aggregate_phase_context().await?; + for phase in [-1, 12345] { + let proto = aggregate_phase_plan(phase, false); + let err = from_substrait_plan(&ctx.state(), &proto).await.unwrap_err(); + assert!( + err.to_string() + .contains(&format!("Invalid aggregation phase {phase}")), + "{err}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn window_aggregation_phases() -> Result<()> { + let original = + read_json("tests/testdata/test_plans/select_window_count.substrait.json"); + let ctx = SessionContext::new(); + ctx.register_batch( + TableReference::bare("DATA"), + record_batch!( + ("D", Int32, [1, 2, 3]), + ("PART", Int32, [1, 1, 1]), + ("ORD", Int32, [1, 2, 3]) + )?, + )?; + for phase in [ + AggregationPhase::Unspecified as i32, + AggregationPhase::InitialToResult as i32, + AggregationPhase::InitialToIntermediate as i32, + AggregationPhase::IntermediateToIntermediate as i32, + AggregationPhase::IntermediateToResult as i32, + 12345, + ] { + let mut proto = original.clone(); + let Some(plan_rel::RelType::Root(root)) = + proto.relations[0].rel_type.as_mut() + else { + panic!("expected root"); + }; + let Some(rel::RelType::Project(project)) = + root.input.as_mut().unwrap().rel_type.as_mut() + else { + panic!("expected projection"); + }; + let Some(expression::RexType::WindowFunction(window)) = + project.expressions[0].rex_type.as_mut() + else { + panic!("expected window function"); + }; + window.phase = phase; + let proto = Plan::decode(proto.encode_to_vec().as_slice()).unwrap(); + let result = from_substrait_plan(&ctx.state(), &proto).await; + if matches!( + AggregationPhase::try_from(phase), + Ok(AggregationPhase::Unspecified | AggregationPhase::InitialToResult) + ) { + let batches = DataFrame::new(ctx.state(), result?).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+-----------+", + "| LEAD_EXPR |", + "+-----------+", + "| 2 |", + "| 3 |", + "| 3 |", + "+-----------+" + ], + &batches + ); + } else { + let err = result.unwrap_err(); + let expected = match AggregationPhase::try_from(phase) { + Ok(phase) => { + format!("Unsupported aggregation phase: {}", phase.as_str_name()) + } + Err(_) => format!("Invalid aggregation phase {phase}"), + }; + assert!(err.to_string().contains(&expected), "{err}"); + } + } + Ok(()) + } #[tokio::test] async fn no_grouping_set() -> Result<()> { From f86799b9ec3cd5bb82949047b145762b062e677f Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sun, 13 Sep 2026 23:22:34 -0700 Subject: [PATCH 2/2] docs: clarify unspecified Substrait aggregate phase compatibility --- .../src/logical_plan/consumer/expr/aggregate_function.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs index 32944a4ae1930..9e020325aee1b 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs @@ -31,8 +31,11 @@ pub(super) fn validate_aggregation_phase(phase: i32) -> datafusion::common::Resu match AggregationPhase::try_from(phase) .map_err(|e| plan_datafusion_err!("Invalid aggregation phase {phase}: {e}"))? { - // Logical plans represent complete calls. Keep accepting unspecified phases - // for compatibility with existing DataFusion aggregate and window producers. + // Substrait defines UNSPECIFIED as INTERMEDIATE_TO_RESULT. Accept it as a + // complete call only for compatibility with existing DataFusion-produced + // aggregate and window plans. This exception also accepts unspecified + // intermediate-state calls from other producers; their intent cannot be + // distinguished here. Explicit intermediate phases remain unsupported. AggregationPhase::Unspecified | AggregationPhase::InitialToResult => Ok(()), phase => not_impl_err!("Unsupported aggregation phase: {}", phase.as_str_name()), }