Skip to content

Commit dd81646

Browse files
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.
1 parent 86ba5e0 commit dd81646

2 files changed

Lines changed: 119 additions & 57 deletions

File tree

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

Lines changed: 66 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
use crate::logical_plan::producer::{
1919
SubstraitProducer, to_substrait_literal_expr, to_substrait_type,
20+
to_substrait_type_from_field,
2021
};
2122
use datafusion::arrow::datatypes::DataType;
2223
use datafusion::common::datatype::FieldExt;
@@ -237,6 +238,14 @@ pub fn from_like(
237238
producer: &mut impl SubstraitProducer,
238239
like: &Like,
239240
schema: &DFSchemaRef,
241+
) -> datafusion::common::Result<Expression> {
242+
make_substrait_like_expr(producer, like, schema)
243+
}
244+
245+
fn make_substrait_like_expr(
246+
producer: &mut impl SubstraitProducer,
247+
like: &Like,
248+
schema: &DFSchemaRef,
240249
) -> datafusion::common::Result<Expression> {
241250
let Like {
242251
negated,
@@ -245,31 +254,18 @@ pub fn from_like(
245254
escape_char,
246255
case_insensitive,
247256
} = like;
248-
make_substrait_like_expr(
249-
producer,
250-
*case_insensitive,
251-
*negated,
252-
expr,
253-
pattern,
254-
*escape_char,
255-
schema,
256-
)
257-
}
258-
259-
fn make_substrait_like_expr(
260-
producer: &mut impl SubstraitProducer,
261-
ignore_case: bool,
262-
negated: bool,
263-
expr: &Expr,
264-
pattern: &Expr,
265-
escape_char: Option<char>,
266-
schema: &DFSchemaRef,
267-
) -> datafusion::common::Result<Expression> {
268-
let function_anchor = if ignore_case {
257+
let function_anchor = if *case_insensitive {
269258
producer.register_function("ilike".to_string())
270259
} else {
271260
producer.register_function("like".to_string())
272261
};
262+
// Substrait documents `output_type` as "Must be set to the return type of
263+
// the function, exactly as derived using the declaration in the extension",
264+
// and a consumer that reads it rejects the call when it is unset. The type
265+
// comes from the expression itself so that it matches what DataFusion
266+
// derives, rather than being restated here.
267+
let (_, output_field) = Expr::Like(like.clone()).to_field(schema)?;
268+
let output_type = to_substrait_type_from_field(producer, &output_field)?;
273269
let expr = producer.handle_expr(expr, schema)?;
274270
let pattern = producer.handle_expr(pattern, schema)?;
275271
let escape_char = to_substrait_literal_expr(
@@ -293,13 +289,13 @@ fn make_substrait_like_expr(
293289
rex_type: Some(RexType::ScalarFunction(ScalarFunction {
294290
function_reference: function_anchor,
295291
arguments,
296-
output_type: None,
292+
output_type: Some(output_type.clone()),
297293
args: vec![],
298294
options: vec![],
299295
})),
300296
};
301297

302-
if negated {
298+
if *negated {
303299
let function_anchor = producer.register_function("not".to_string());
304300

305301
#[expect(deprecated)]
@@ -309,7 +305,8 @@ fn make_substrait_like_expr(
309305
arguments: vec![FunctionArgument {
310306
arg_type: Some(ArgType::Value(substrait_like)),
311307
}],
312-
output_type: None,
308+
// `not` yields the type its argument does.
309+
output_type: Some(output_type),
313310
args: vec![],
314311
options: vec![],
315312
})),
@@ -318,8 +315,6 @@ fn make_substrait_like_expr(
318315
Ok(substrait_like)
319316
}
320317
}
321-
322-
/// Util to generate substrait [RexType::ScalarFunction] with one argument
323318
fn to_substrait_unary_scalar_fn(
324319
producer: &mut impl SubstraitProducer,
325320
fn_name: &str,
@@ -452,12 +447,14 @@ mod tests {
452447
use crate::logical_plan::producer::{
453448
DefaultSubstraitProducer, SubstraitProducer, to_substrait_type,
454449
};
455-
use datafusion::arrow::datatypes::DataType;
450+
use datafusion::arrow::datatypes::{DataType, Field, Schema};
456451
use datafusion::common::{DFSchema, DFSchemaRef};
457452
use datafusion::execution::SessionStateBuilder;
458-
use datafusion::prelude::lit;
453+
use datafusion::logical_expr::{Expr, Like};
454+
use datafusion::prelude::{col, lit};
459455
use substrait::proto::Expression;
460456
use substrait::proto::expression::{RexType, ScalarFunction};
457+
use substrait::proto::function_argument::ArgType;
461458

462459
#[tokio::test]
463460
async fn binary_expr_output_type() -> datafusion::common::Result<()> {
@@ -479,4 +476,45 @@ mod tests {
479476
panic!("Substrait ScalarFunction expected")
480477
}
481478
}
479+
480+
/// The `like` call carries the type the expression yields, and so does the
481+
/// `not` that wraps a negated one.
482+
#[tokio::test]
483+
async fn like_output_type() -> datafusion::common::Result<()> {
484+
let state = SessionStateBuilder::default().build();
485+
let schema =
486+
DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new(
487+
"s",
488+
DataType::Utf8,
489+
true,
490+
)]))?);
491+
let mut producer = DefaultSubstraitProducer::new(&state);
492+
// A `LIKE` over a nullable input yields a nullable boolean.
493+
let expected = to_substrait_type(&mut producer, &DataType::Boolean, true)?;
494+
495+
let like = Like::new(false, Box::new(col("s")), Box::new(lit("a%")), None, false);
496+
let substrait_expr = producer.handle_expr(&Expr::Like(like), &schema)?;
497+
let Some(RexType::ScalarFunction(like_fn)) = substrait_expr.rex_type else {
498+
panic!("Substrait ScalarFunction expected")
499+
};
500+
assert_eq!(like_fn.output_type, Some(expected.clone()));
501+
502+
let negated =
503+
Like::new(true, Box::new(col("s")), Box::new(lit("a%")), None, false);
504+
let substrait_expr = producer.handle_expr(&Expr::Like(negated), &schema)?;
505+
let Some(RexType::ScalarFunction(not_fn)) = substrait_expr.rex_type else {
506+
panic!("Substrait ScalarFunction expected")
507+
};
508+
assert_eq!(not_fn.output_type, Some(expected.clone()));
509+
510+
// The `not` wraps the `like`, which carries the type as well.
511+
let Some(ArgType::Value(Expression {
512+
rex_type: Some(RexType::ScalarFunction(inner)),
513+
})) = not_fn.arguments[0].arg_type.clone()
514+
else {
515+
panic!("Substrait ScalarFunction expected inside `not`")
516+
};
517+
assert_eq!(inner.output_type, Some(expected));
518+
Ok(())
519+
}
482520
}

‎datafusion/substrait/src/logical_plan/producer/expr/window_function.rs‎

Lines changed: 53 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
// under the License.
1717

1818
use crate::logical_plan::producer::SubstraitProducer;
19+
use crate::logical_plan::producer::to_substrait_type_from_field;
1920
use crate::logical_plan::producer::utils::substrait_sort_field;
2021
use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err};
2122
use datafusion::logical_expr::expr::{WindowFunction, WindowFunctionParams};
23+
use datafusion::logical_expr::{Expr, ExprSchemable};
2224
use datafusion::logical_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits};
2325
use substrait::proto::aggregate_function::AggregationInvocation;
2426
use substrait::proto::expression::RexType;
@@ -27,7 +29,7 @@ use substrait::proto::expression::window_function::bound as SubstraitBound;
2729
use substrait::proto::expression::window_function::bound::Kind as BoundKind;
2830
use substrait::proto::expression::window_function::{Bound, BoundsType};
2931
use substrait::proto::function_argument::ArgType;
30-
use substrait::proto::{Expression, FunctionArgument, SortField};
32+
use substrait::proto::{Expression, FunctionArgument};
3133

3234
pub fn from_window_function(
3335
producer: &mut impl SubstraitProducer,
@@ -79,47 +81,36 @@ pub fn from_window_function(
7981
// window frame
8082
let bounds = to_substrait_bounds(window_frame)?;
8183
let bound_type = to_substrait_bound_type(window_frame)?;
82-
Ok(make_substrait_window_function(
83-
function_anchor,
84-
arguments,
85-
partition_by,
86-
order_by,
87-
bounds,
88-
bound_type,
89-
*distinct,
90-
))
91-
}
84+
// Substrait documents `output_type` as "Must be set to the return type of
85+
// the function, exactly as derived using the declaration in the extension",
86+
// and a consumer that reads it rejects the call when it is unset. The type
87+
// comes from the expression itself so that it matches what DataFusion
88+
// derives, rather than being restated here.
89+
let (_, output_field) =
90+
Expr::WindowFunction(Box::new(window_fn.clone())).to_field(schema)?;
91+
let output_type = to_substrait_type_from_field(producer, &output_field)?;
9292

93-
fn make_substrait_window_function(
94-
function_reference: u32,
95-
arguments: Vec<FunctionArgument>,
96-
partitions: Vec<Expression>,
97-
sorts: Vec<SortField>,
98-
bounds: (Bound, Bound),
99-
bounds_type: BoundsType,
100-
distinct: bool,
101-
) -> Expression {
10293
#[expect(deprecated)]
103-
Expression {
94+
Ok(Expression {
10495
rex_type: Some(RexType::WindowFunction(SubstraitWindowFunction {
105-
function_reference,
96+
function_reference: function_anchor,
10697
arguments,
107-
partitions,
108-
sorts,
98+
partitions: partition_by,
99+
sorts: order_by,
109100
options: vec![],
110-
output_type: None,
101+
output_type: Some(output_type),
111102
phase: 0, // default to AGGREGATION_PHASE_UNSPECIFIED
112-
invocation: if distinct {
103+
invocation: if *distinct {
113104
AggregationInvocation::Distinct as i32
114105
} else {
115106
AggregationInvocation::All as i32
116107
},
117108
lower_bound: Some(bounds.0),
118109
upper_bound: Some(bounds.1),
119110
args: vec![],
120-
bounds_type: bounds_type as i32,
111+
bounds_type: bound_type as i32,
121112
})),
122-
}
113+
})
123114
}
124115

125116
fn to_substrait_bound_type(
@@ -185,7 +176,40 @@ fn to_substrait_bound_offset(value: &ScalarValue) -> datafusion::common::Result<
185176
#[cfg(test)]
186177
mod tests {
187178
use super::*;
188-
use datafusion::common::assert_contains;
179+
use crate::logical_plan::producer::{DefaultSubstraitProducer, to_substrait_type};
180+
use datafusion::arrow::datatypes::{DataType, Field, Schema};
181+
use datafusion::common::{DFSchema, assert_contains};
182+
use datafusion::execution::SessionStateBuilder;
183+
use datafusion::functions_aggregate::sum::sum_udaf;
184+
use datafusion::logical_expr::expr::WindowFunction;
185+
use datafusion::prelude::col;
186+
use substrait::proto::expression::RexType;
187+
188+
/// Substrait requires the return type on a window function call, and a
189+
/// consumer that reads it rejects the call when it is unset.
190+
#[test]
191+
fn window_function_output_type() -> datafusion::common::Result<()> {
192+
let state = SessionStateBuilder::default().build();
193+
let schema =
194+
DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new(
195+
"i",
196+
DataType::Int64,
197+
true,
198+
)]))?);
199+
let mut producer = DefaultSubstraitProducer::new(&state);
200+
201+
let window_fn = WindowFunction::new(sum_udaf(), vec![col("i")]);
202+
let expr = Expr::WindowFunction(Box::new(window_fn));
203+
let substrait_expr = producer.handle_expr(&expr, &schema)?;
204+
205+
let Some(RexType::WindowFunction(window)) = substrait_expr.rex_type else {
206+
panic!("Substrait WindowFunction expected")
207+
};
208+
// `sum` over a nullable i64 yields a nullable i64.
209+
let expected = to_substrait_type(&mut producer, &DataType::Int64, true)?;
210+
assert_eq!(window.output_type, Some(expected));
211+
Ok(())
212+
}
189213

190214
#[test]
191215
fn window_frame_offsets() {

0 commit comments

Comments
 (0)