Skip to content

Commit 6d7d41f

Browse files
committed
Set Substrait output type for unary expressions
1 parent 69e2164 commit 6d7d41f

1 file changed

Lines changed: 9 additions & 2 deletions

File tree

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,13 @@ pub fn from_unary_expr(
181181
Expr::Negative(arg) => ("negate", arg),
182182
expr => not_impl_err!("Unsupported expression: {expr:?}")?,
183183
};
184-
to_substrait_unary_scalar_fn(producer, fn_name, arg, schema)
184+
let (_, output_field) = expr.to_field(schema)?;
185+
let output_type = to_substrait_type(
186+
producer,
187+
output_field.data_type(),
188+
output_field.is_nullable(),
189+
)?;
190+
to_substrait_unary_scalar_fn(producer, fn_name, arg, schema, &output_type)
185191
}
186192

187193
pub fn from_binary_expr(
@@ -299,6 +305,7 @@ fn to_substrait_unary_scalar_fn(
299305
fn_name: &str,
300306
arg: &Expr,
301307
schema: &DFSchemaRef,
308+
output_type: &Type,
302309
) -> datafusion::common::Result<Expression> {
303310
let function_anchor = producer.register_function(fn_name.to_string());
304311
let substrait_expr = producer.handle_expr(arg, schema)?;
@@ -309,7 +316,7 @@ fn to_substrait_unary_scalar_fn(
309316
arguments: vec![FunctionArgument {
310317
arg_type: Some(ArgType::Value(substrait_expr)),
311318
}],
312-
output_type: None,
319+
output_type: Some(output_type.clone()),
313320
options: vec![],
314321
..Default::default()
315322
})),

0 commit comments

Comments
 (0)