From be39554b78dd04d9b7d302d6305c23e1f50e2786 Mon Sep 17 00:00:00 2001 From: naman Date: Wed, 16 Sep 2026 14:38:52 +0530 Subject: [PATCH] fix: Set Substrait output_type on window functions and LIKE Substrait documents output_type on both ScalarFunction and WindowFunction as the return type of the function, and a consumer that reads it rejects a call that leaves it unset: substrait-java refuses a window function plan and a LIKE plan with "Type is not set". Both types are derived from the expression itself, so they match what DataFusion derives: from_window_function from Expr::WindowFunction, and make_substrait_like_expr from Expr::Like, which also covers the not() wrapper a negated LIKE emits. make_substrait_window_function had a single caller and its body now sits in from_window_function, and make_substrait_like_expr takes the Like rather than its five fields, so neither grows an extra parameter. A DataFusion round trip cannot catch this, because the consumer reads output_type only when it converts a cast. --- .../producer/expr/scalar_function.rs | 94 +++++++++++++------ .../producer/expr/window_function.rs | 82 ++++++++++------ 2 files changed, 119 insertions(+), 57 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index 75720395aae7c..9e531a1ae549f 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs @@ -17,6 +17,7 @@ use crate::logical_plan::producer::{ SubstraitProducer, to_substrait_literal_expr, to_substrait_type, + to_substrait_type_from_field, }; use datafusion::arrow::datatypes::DataType; use datafusion::common::datatype::FieldExt; @@ -237,6 +238,14 @@ pub fn from_like( producer: &mut impl SubstraitProducer, like: &Like, schema: &DFSchemaRef, +) -> datafusion::common::Result { + make_substrait_like_expr(producer, like, schema) +} + +fn make_substrait_like_expr( + producer: &mut impl SubstraitProducer, + like: &Like, + schema: &DFSchemaRef, ) -> datafusion::common::Result { let Like { negated, @@ -245,31 +254,18 @@ pub fn from_like( escape_char, case_insensitive, } = like; - make_substrait_like_expr( - producer, - *case_insensitive, - *negated, - expr, - pattern, - *escape_char, - schema, - ) -} - -fn make_substrait_like_expr( - producer: &mut impl SubstraitProducer, - ignore_case: bool, - negated: bool, - expr: &Expr, - pattern: &Expr, - escape_char: Option, - schema: &DFSchemaRef, -) -> datafusion::common::Result { - let function_anchor = if ignore_case { + let function_anchor = if *case_insensitive { producer.register_function("ilike".to_string()) } else { producer.register_function("like".to_string()) }; + // Substrait documents `output_type` as "Must be set to the return type of + // the function, exactly as derived using the declaration in the extension", + // and a consumer that reads it rejects the call when it is unset. The type + // comes from the expression itself so that it matches what DataFusion + // derives, rather than being restated here. + let (_, output_field) = Expr::Like(like.clone()).to_field(schema)?; + let output_type = to_substrait_type_from_field(producer, &output_field)?; let expr = producer.handle_expr(expr, schema)?; let pattern = producer.handle_expr(pattern, schema)?; let escape_char = to_substrait_literal_expr( @@ -293,13 +289,13 @@ fn make_substrait_like_expr( rex_type: Some(RexType::ScalarFunction(ScalarFunction { function_reference: function_anchor, arguments, - output_type: None, + output_type: Some(output_type.clone()), args: vec![], options: vec![], })), }; - if negated { + if *negated { let function_anchor = producer.register_function("not".to_string()); #[expect(deprecated)] @@ -309,7 +305,8 @@ fn make_substrait_like_expr( arguments: vec![FunctionArgument { arg_type: Some(ArgType::Value(substrait_like)), }], - output_type: None, + // `not` yields the type its argument does. + output_type: Some(output_type), args: vec![], options: vec![], })), @@ -318,8 +315,6 @@ fn make_substrait_like_expr( Ok(substrait_like) } } - -/// Util to generate substrait [RexType::ScalarFunction] with one argument fn to_substrait_unary_scalar_fn( producer: &mut impl SubstraitProducer, fn_name: &str, @@ -452,12 +447,14 @@ mod tests { use crate::logical_plan::producer::{ DefaultSubstraitProducer, SubstraitProducer, to_substrait_type, }; - use datafusion::arrow::datatypes::DataType; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::execution::SessionStateBuilder; - use datafusion::prelude::lit; + use datafusion::logical_expr::{Expr, Like}; + use datafusion::prelude::{col, lit}; use substrait::proto::Expression; use substrait::proto::expression::{RexType, ScalarFunction}; + use substrait::proto::function_argument::ArgType; #[tokio::test] async fn binary_expr_output_type() -> datafusion::common::Result<()> { @@ -479,4 +476,45 @@ mod tests { panic!("Substrait ScalarFunction expected") } } + + /// The `like` call carries the type the expression yields, and so does the + /// `not` that wraps a negated one. + #[tokio::test] + async fn like_output_type() -> datafusion::common::Result<()> { + let state = SessionStateBuilder::default().build(); + let schema = + DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new( + "s", + DataType::Utf8, + true, + )]))?); + let mut producer = DefaultSubstraitProducer::new(&state); + // A `LIKE` over a nullable input yields a nullable boolean. + let expected = to_substrait_type(&mut producer, &DataType::Boolean, true)?; + + let like = Like::new(false, Box::new(col("s")), Box::new(lit("a%")), None, false); + let substrait_expr = producer.handle_expr(&Expr::Like(like), &schema)?; + let Some(RexType::ScalarFunction(like_fn)) = substrait_expr.rex_type else { + panic!("Substrait ScalarFunction expected") + }; + assert_eq!(like_fn.output_type, Some(expected.clone())); + + let negated = + Like::new(true, Box::new(col("s")), Box::new(lit("a%")), None, false); + let substrait_expr = producer.handle_expr(&Expr::Like(negated), &schema)?; + let Some(RexType::ScalarFunction(not_fn)) = substrait_expr.rex_type else { + panic!("Substrait ScalarFunction expected") + }; + assert_eq!(not_fn.output_type, Some(expected.clone())); + + // The `not` wraps the `like`, which carries the type as well. + let Some(ArgType::Value(Expression { + rex_type: Some(RexType::ScalarFunction(inner)), + })) = not_fn.arguments[0].arg_type.clone() + else { + panic!("Substrait ScalarFunction expected inside `not`") + }; + assert_eq!(inner.output_type, Some(expected)); + Ok(()) + } } 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..79dc5f461b38f 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/window_function.rs @@ -16,9 +16,11 @@ // under the License. use crate::logical_plan::producer::SubstraitProducer; +use crate::logical_plan::producer::to_substrait_type_from_field; 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::{Expr, ExprSchemable}; use datafusion::logical_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits}; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::expression::RexType; @@ -27,7 +29,7 @@ use substrait::proto::expression::window_function::bound as SubstraitBound; use substrait::proto::expression::window_function::bound::Kind as BoundKind; use substrait::proto::expression::window_function::{Bound, BoundsType}; use substrait::proto::function_argument::ArgType; -use substrait::proto::{Expression, FunctionArgument, SortField}; +use substrait::proto::{Expression, FunctionArgument}; pub fn from_window_function( producer: &mut impl SubstraitProducer, @@ -79,37 +81,26 @@ pub fn from_window_function( // window frame let bounds = to_substrait_bounds(window_frame)?; let bound_type = to_substrait_bound_type(window_frame)?; - Ok(make_substrait_window_function( - function_anchor, - arguments, - partition_by, - order_by, - bounds, - bound_type, - *distinct, - )) -} + // Substrait documents `output_type` as "Must be set to the return type of + // the function, exactly as derived using the declaration in the extension", + // and a consumer that reads it rejects the call when it is unset. The type + // comes from the expression itself so that it matches what DataFusion + // derives, rather than being restated here. + let (_, output_field) = + Expr::WindowFunction(Box::new(window_fn.clone())).to_field(schema)?; + let output_type = to_substrait_type_from_field(producer, &output_field)?; -fn make_substrait_window_function( - function_reference: u32, - arguments: Vec, - partitions: Vec, - sorts: Vec, - bounds: (Bound, Bound), - bounds_type: BoundsType, - distinct: bool, -) -> Expression { #[expect(deprecated)] - Expression { + Ok(Expression { rex_type: Some(RexType::WindowFunction(SubstraitWindowFunction { - function_reference, + function_reference: function_anchor, arguments, - partitions, - sorts, + partitions: partition_by, + sorts: order_by, options: vec![], - output_type: None, + output_type: Some(output_type), phase: 0, // default to AGGREGATION_PHASE_UNSPECIFIED - invocation: if distinct { + invocation: if *distinct { AggregationInvocation::Distinct as i32 } else { AggregationInvocation::All as i32 @@ -117,9 +108,9 @@ fn make_substrait_window_function( lower_bound: Some(bounds.0), upper_bound: Some(bounds.1), args: vec![], - bounds_type: bounds_type as i32, + bounds_type: bound_type as i32, })), - } + }) } fn to_substrait_bound_type( @@ -185,7 +176,40 @@ fn to_substrait_bound_offset(value: &ScalarValue) -> datafusion::common::Result< #[cfg(test)] mod tests { use super::*; - use datafusion::common::assert_contains; + use crate::logical_plan::producer::{DefaultSubstraitProducer, to_substrait_type}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{DFSchema, assert_contains}; + use datafusion::execution::SessionStateBuilder; + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::logical_expr::expr::WindowFunction; + use datafusion::prelude::col; + use substrait::proto::expression::RexType; + + /// Substrait requires the return type on a window function call, and a + /// consumer that reads it rejects the call when it is unset. + #[test] + fn window_function_output_type() -> datafusion::common::Result<()> { + let state = SessionStateBuilder::default().build(); + let schema = + DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new( + "i", + DataType::Int64, + true, + )]))?); + let mut producer = DefaultSubstraitProducer::new(&state); + + let window_fn = WindowFunction::new(sum_udaf(), vec![col("i")]); + let expr = Expr::WindowFunction(Box::new(window_fn)); + let substrait_expr = producer.handle_expr(&expr, &schema)?; + + let Some(RexType::WindowFunction(window)) = substrait_expr.rex_type else { + panic!("Substrait WindowFunction expected") + }; + // `sum` over a nullable i64 yields a nullable i64. + let expected = to_substrait_type(&mut producer, &DataType::Int64, true)?; + assert_eq!(window.output_type, Some(expected)); + Ok(()) + } #[test] fn window_frame_offsets() {