Skip to content

Commit 62f2943

Browse files
fix: Set Substrait output_type on window functions and LIKE (#25367)
## Which issue does this PR close? - Closes #25366. ## Rationale for this change The producer left `output_type` unset on window function calls and on `LIKE` / `ILIKE`, including the `not` that wraps a negated one. Substrait documents both `Expression.WindowFunction.output_type` and `Expression.ScalarFunction.output_type` as: > Must be set to the return type of the function, exactly as derived using the declaration in the extension. A consumer that reads the field rejects such a call. substrait-java 0.103.0 refuses both plans, from `ProtoTypeConverter.from:125`: | plan produced for | error before this PR | | --- | --- | | `SELECT sum(i) OVER (ORDER BY i) FROM t` | `UnsupportedOperationException: Type is not set` at `ProtoExpressionConverter.fromWindowFunction:406` | | `SELECT i FROM t WHERE CAST(i AS VARCHAR) LIKE '1%'` | `UnsupportedOperationException: Type is not set` at `ProtoExpressionConverter:201`, via `ProtoRelConverter.newFilter` | A DataFusion round trip cannot catch this: the consumer reads `output_type` only in `consumer/expr/cast.rs`, so producer and consumer agree on the omission. #15831 and #20597 set this field for binary and unary expressions, `from_function` and higher order functions, and #25049 and #25090 cover `AggregateFunction`. These were the remaining expression kinds. ## What changes are included in this PR? Both types come from the expression itself, so they match what DataFusion derives rather than being restated in the producer: - `from_window_function` derives the type from `Expr::WindowFunction(..).to_field(schema)` and writes it on the call it builds. `make_substrait_window_function` had one caller and eight parameters once the type was added, so its body now sits in `from_window_function` and the helper is gone. - `from_like` derives the type from `Expr::Like(..).to_field(schema)` and sets it on the `like` call and on the `not` wrapper of a negated one, which yields that same type. `make_substrait_like_expr` had `from_like` as its only caller, so, as with the window helper, its body now sits in `from_like`. ## What is the testing strategy for this PR? - Two unit tests next to the existing `binary_expr_output_type`: `window_function_output_type` asserts `sum(i)` over a nullable `i64` carries a nullable `i64`, and `like_output_type` asserts a `LIKE` over a nullable input carries a nullable boolean, on the `like` call, on the `not` of a negated one, and on the `like` nested inside that `not`. - Each half was reverted on its own to check the tests pin it: without the window change only `window_function_output_type` fails, without the `LIKE` change only `like_output_type` fails. - `cargo test -p datafusion-substrait` passes (60 unit, 210 integration with the 6 that were already ignored, 3 doc tests), and `./ci/scripts/rust_clippy.sh`, the workspace clippy CI runs, is clean. - Emitted protobuf before and after, read directly rather than through a consumer: | plan | before | after | | --- | --- | --- | | window `sum` | `output_type` unset | set | | `like` | unset | set | | `not` around a negated `like` | unset | set | - With the same plans, substrait-java 0.103.0 no longer raises `Type is not set`. Both now get through `ProtoPlanConverter`, while the plans built from `main` still fail there. Each then stops further along for reasons that have nothing to do with this field: substrait-spark has no visitor for a window function that appears in a project expression, and Spark's `like` takes two arguments while the call we emit passes three, the third being the escape character. That second one is #25442, which emits `LIKE` in the two-argument form the extension defines. ## Are there any user-facing changes? No API changes. Plans produced for window functions and `LIKE` now carry the expression's type, which consumers that require it will accept.
1 parent 6e319bd commit 62f2943

2 files changed

Lines changed: 138 additions & 55 deletions

File tree

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

Lines changed: 58 additions & 26 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;
@@ -245,31 +246,18 @@ pub fn from_like(
245246
escape_char,
246247
case_insensitive,
247248
} = 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 {
249+
let function_anchor = if *case_insensitive {
269250
producer.register_function("ilike".to_string())
270251
} else {
271252
producer.register_function("like".to_string())
272253
};
254+
// Substrait documents `output_type` as "Must be set to the return type of
255+
// the function, exactly as derived using the declaration in the extension",
256+
// and a consumer that reads it rejects the call when it is unset. The type
257+
// comes from the expression itself so that it matches what DataFusion
258+
// derives, rather than being restated here.
259+
let (_, output_field) = Expr::Like(like.clone()).to_field(schema)?;
260+
let output_type = to_substrait_type_from_field(producer, &output_field)?;
273261
let expr = producer.handle_expr(expr, schema)?;
274262
let pattern = producer.handle_expr(pattern, schema)?;
275263
let escape_char = to_substrait_literal_expr(
@@ -293,13 +281,13 @@ fn make_substrait_like_expr(
293281
rex_type: Some(RexType::ScalarFunction(ScalarFunction {
294282
function_reference: function_anchor,
295283
arguments,
296-
output_type: None,
284+
output_type: Some(output_type.clone()),
297285
args: vec![],
298286
options: vec![],
299287
})),
300288
};
301289

302-
if negated {
290+
if *negated {
303291
let function_anchor = producer.register_function("not".to_string());
304292

305293
#[expect(deprecated)]
@@ -309,7 +297,8 @@ fn make_substrait_like_expr(
309297
arguments: vec![FunctionArgument {
310298
arg_type: Some(ArgType::Value(substrait_like)),
311299
}],
312-
output_type: None,
300+
// `not` yields the type its argument does.
301+
output_type: Some(output_type),
313302
args: vec![],
314303
options: vec![],
315304
})),
@@ -452,12 +441,14 @@ mod tests {
452441
use crate::logical_plan::producer::{
453442
DefaultSubstraitProducer, SubstraitProducer, to_substrait_type,
454443
};
455-
use datafusion::arrow::datatypes::DataType;
444+
use datafusion::arrow::datatypes::{DataType, Field, Schema};
456445
use datafusion::common::{DFSchema, DFSchemaRef};
457446
use datafusion::execution::SessionStateBuilder;
458-
use datafusion::prelude::lit;
447+
use datafusion::logical_expr::{Expr, Like};
448+
use datafusion::prelude::{col, lit};
459449
use substrait::proto::Expression;
460450
use substrait::proto::expression::{RexType, ScalarFunction};
451+
use substrait::proto::function_argument::ArgType;
461452

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

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

Lines changed: 80 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,67 @@ 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::count::count_all_window;
184+
use datafusion::functions_aggregate::sum::sum_udaf;
185+
use datafusion::logical_expr::expr::WindowFunction;
186+
use datafusion::prelude::col;
187+
use substrait::proto::expression::RexType;
188+
189+
/// Substrait requires the return type on a window function call, and a
190+
/// consumer that reads it rejects the call when it is unset.
191+
#[test]
192+
fn window_function_output_type() -> datafusion::common::Result<()> {
193+
let state = SessionStateBuilder::default().build();
194+
let schema =
195+
DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new(
196+
"i",
197+
DataType::Int64,
198+
true,
199+
)]))?);
200+
let mut producer = DefaultSubstraitProducer::new(&state);
201+
202+
let window_fn = WindowFunction::new(sum_udaf(), vec![col("i")]);
203+
let expr = Expr::WindowFunction(Box::new(window_fn));
204+
let substrait_expr = producer.handle_expr(&expr, &schema)?;
205+
206+
let Some(RexType::WindowFunction(window)) = substrait_expr.rex_type else {
207+
panic!("Substrait WindowFunction expected")
208+
};
209+
// `sum` over a nullable i64 yields a nullable i64.
210+
let expected = to_substrait_type(&mut producer, &DataType::Int64, true)?;
211+
assert_eq!(window.output_type, Some(expected));
212+
Ok(())
213+
}
214+
215+
/// `COUNT(*) OVER ()` always yields a non-null count, so this guards
216+
/// against a regression where `output_type` is emitted as nullable
217+
/// unconditionally.
218+
#[test]
219+
fn count_window_function_output_type_is_non_nullable()
220+
-> datafusion::common::Result<()> {
221+
let state = SessionStateBuilder::default().build();
222+
let schema =
223+
DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![Field::new(
224+
"i",
225+
DataType::Int64,
226+
true,
227+
)]))?);
228+
let mut producer = DefaultSubstraitProducer::new(&state);
229+
230+
let substrait_expr = producer.handle_expr(&count_all_window(), &schema)?;
231+
232+
let Some(RexType::WindowFunction(window)) = substrait_expr.rex_type else {
233+
panic!("Substrait WindowFunction expected")
234+
};
235+
// `count(*)` never produces a null, regardless of input nullability.
236+
let expected = to_substrait_type(&mut producer, &DataType::Int64, false)?;
237+
assert_eq!(window.output_type, Some(expected));
238+
Ok(())
239+
}
189240

190241
#[test]
191242
fn window_frame_offsets() {

0 commit comments

Comments
 (0)