fix: Set Substrait output_type on window functions and LIKE

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<Expression> {
+    make_substrait_like_expr(producer, like, schema)
+}
+
+fn make_substrait_like_expr(
+    producer: &mut impl SubstraitProducer,
+    like: &Like,
+    schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
     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<char>,
-    schema: &DFSchemaRef,
-) -> datafusion::common::Result<Expression> {
-    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(())
+    }
 }

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,47 +81,36 @@ 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<FunctionArgument>,
-    partitions: Vec<Expression>,
-    sorts: Vec<SortField>,
-    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
             },
             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() {