Set Substrait output type for functions

Author: wlhjason

datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs

@@ -18,6 +18,7 @@
 use crate::logical_plan::producer::{
     SubstraitProducer, to_substrait_literal_expr, to_substrait_type,
 };
+use datafusion::arrow::datatypes::Field;
 use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err};
 use datafusion::logical_expr::{
     Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, expr,
@@ -31,21 +32,24 @@ pub fn from_scalar_function(
     fun: &expr::ScalarFunction,
     schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    from_function(producer, fun.name(), &fun.args, schema)
+    let (_, output_field) = Expr::ScalarFunction(fun.clone()).to_field(schema)?;
+    from_function(producer, fun.name(), &fun.args, &output_field, schema)
 }
 
 pub fn from_higher_order_function(
     producer: &mut impl SubstraitProducer,
     fun: &expr::HigherOrderFunction,
     schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    from_function(producer, fun.name(), &fun.args, schema)
+    let (_, output_field) = Expr::HigherOrderFunction(fun.clone()).to_field(schema)?;
+    from_function(producer, fun.name(), &fun.args, &output_field, schema)
 }
 
 fn from_function(
     producer: &mut impl SubstraitProducer,
     name: &str,
     args: &[Expr],
+    output_field: &Field,
     schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
     let mut arguments: Vec<FunctionArgument> = vec![];
@@ -56,14 +60,19 @@ fn from_function(
     }
 
     let arguments = custom_argument_handler(name, arguments);
+    let output_type = to_substrait_type(
+        producer,
+        output_field.data_type(),
+        output_field.is_nullable(),
+    )?;
 
     let function_anchor = producer.register_function(name.to_string());
     #[expect(deprecated)]
     Ok(Expression {
         rex_type: Some(RexType::ScalarFunction(ScalarFunction {
             function_reference: function_anchor,
             arguments,
-            output_type: None,
+            output_type: Some(output_type),
             options: vec![],
             args: vec![],
         })),