This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 424cfd8eb7 fix: reject unsupported Substrait aggregation phases
(#25045)
424cfd8eb7 is described below
commit 424cfd8eb7e5d73229227324f02b39179babeeb1
Author: Goutam Adwant <[email protected]>
AuthorDate: Thu Sep 24 07:29:44 2026 +0000
fix: reject unsupported Substrait aggregation phases (#25045)
## Which issue does this PR close?
- Closes #24967.
## Rationale for this change
The Substrait consumer treats explicit intermediate aggregate phases as
complete calls. This can silently return final values when a plan
requests intermediate state, or report an unrelated root-schema naming
error.
## What changes are included in this PR?
- Validate phases in aggregate and window expressions before translating
their arguments.
- Accept `INITIAL_TO_RESULT` and retain `UNSPECIFIED` for compatibility
with existing DataFusion-produced plans.
- Reject explicit intermediate phases and unknown protobuf enum values
with clear errors.
- Document the compatibility limitation: Substrait defines `UNSPECIFIED`
as `INTERMEDIATE_TO_RESULT`, but the consumer cannot distinguish legacy
DataFusion complete calls from unspecified intermediate-state calls
produced elsewhere. Those calls remain accepted. The producer change in
#25146 is complementary and is not duplicated here.
## What is the testing strategy for this PR?
- Reproduced the original behavior: an `INITIAL_TO_INTERMEDIATE` average
over values 1 and 2 returned 1.5 instead of intermediate state.
- Tests cover supported phases, every explicit unsupported phase, rooted
and unrooted aggregates, unknown binary-protobuf enum values, and actual
window output.
- The full Substrait integration target passes: 213 tests passed, with
six existing tests ignored.
- The extended workspace suite passes 11,267 Rust tests, with eight
existing tests ignored, and all 511 SQL logic-test files using an
explicit four-thread limit. An initial default-concurrency run failed
one ordered-aggregate spill test with a test-memory-pool exhaustion; the
complete limited-concurrency rerun passed without source changes.
- `cargo clippy --all-targets --all-features -- -D warnings` and the
complete `./dev/rust_lint.sh` pass, including strict workspace
documentation checks.
## Are there any user-facing changes?
Plans with explicit unsupported aggregate or window phases now fail
instead of being interpreted as complete calls. Existing
unspecified-phase plans remain accepted, including the ambiguity
described above. No public Rust API changes are included.
Intermediate-state execution and the separate AVG output-type mismatch
are not addressed here.
---
.../consumer/expr/aggregate_function.rs | 22 ++-
.../logical_plan/consumer/expr/window_function.rs | 2 +
.../substrait/tests/cases/aggregation_tests.rs | 193 ++++++++++++++++++++-
3 files changed, 214 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 096eef7ae3..9e020325ae 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,28 @@
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}"))?
+ {
+ // 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()),
+ }
+}
/// Convert Substrait AggregateFunction to DataFusion Expr
pub async fn from_substrait_agg_func(
@@ -33,6 +50,7 @@ pub async fn from_substrait_agg_func(
order_by: Vec<SortExpr>,
distinct: bool,
) -> datafusion::common::Result<Arc<Expr>> {
+ 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 d39b325a54..bd1a63b704 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<Expr> {
+ 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 e572023f17..8af4cb7fa5 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<SessionContext> {
+ 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<()> {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
