Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix get_type for higher-order array functions #13756

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+114 −3
Open

Fix get_type for higher-order array functions #13756

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6903259
Fix get_type for higher-order array functions
findepi Dec 13, 2024
6d81418
Fix recursive flatten
findepi Dec 13, 2024
038a015
Restore "keep LargeList" in Array signature
findepi Dec 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions datafusion/expr/src/type_coercion/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use arrow::{
use datafusion_common::{
exec_err, internal_datafusion_err, internal_err, plan_err,
types::{LogicalType, NativeType},
utils::{coerced_fixed_size_list_to_list, list_ndims},
utils::list_ndims,
Result,
};
use datafusion_expr_common::{
Expand Down Expand Up @@ -416,11 +416,9 @@ fn get_valid_types(
}
fn array(array_type: &DataType) -> Option<DataType> {
match array_type {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so this says that if the type is a list, keep the type, but if the type is large list / fixed size list then take the field type?

Why doesn't it also take the field type for List 🤔 ? (Aka it doesn't make sense to me that List is treated differently than LargeList and FixedSizeList

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for backwards compat i should keep LargeList so it stays LargeList, will push shortly

Aka it doesn't make sense to me that List is treated differently than LargeList and FixedSizeList

not my invention, it was like this before.
i think the intention is "converge List, LL and FSL into one type... or maybe two types... to keep UDF impl simpler".

i am not attached to this approach, but i think code may be reliant on that

DataType::List(_)
| DataType::LargeList(_)
| DataType::FixedSizeList(_, _) => {
let array_type = coerced_fixed_size_list_to_list(array_type);
Some(array_type)
DataType::List(_) => Some(array_type.clone()),
DataType::LargeList(field) | DataType::FixedSizeList(field, _) => {
Some(DataType::List(Arc::clone(field)))
}
_ => None,
}
Expand Down
81 changes: 81 additions & 0 deletions datafusion/functions-nested/src/extract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -993,3 +993,84 @@ where
let data = mutable.freeze();
Ok(arrow::array::make_array(data))
}

#[cfg(test)]
mod tests {
use super::array_element_udf;
use arrow_schema::{DataType, Field};
use datafusion_common::{Column, DFSchema, ScalarValue};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::{cast, Expr, ExprSchemable};
use std::collections::HashMap;

#[test]
fn test_array_element_return_type() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can add tests in slt file that cover the array signature test cases, so we can avoid creating rust test here.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The rust test allows explicitly exercising various ways of getting expression type.
Before i wrote it, I wasn't even sure whether it's a bug or a feature.

I can add slt test, how would it look like?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did try to write some slt regression tests, but i couldn't expose the bug. Yet, the unit tests proves the bug exists.
I trust you have a better intuition how signature related bug can be exposed in SLT. Please advise.

let complex_type = DataType::FixedSizeList(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When I change this complex type to DataType::List the test passes 🤔

        let complex_type = DataType::List(
            Field::new("some_arbitrary_test_field", DataType::Int32, false).into(),
        );

It also passes when complex_type is a Struct

        let complex_type = DataType::Struct(Fields::from(vec![
            Arc::new(Field::new("some_arbitrary_test_field", DataType::Int32, false)),
        ]));

It seems like there is something about FixedSizeList that is causing issues to me

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Weird, when I remove this line in expr schema the test passes (with FixedSizedList):

diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs
index 3317deafb..50aeb222f 100644
--- a/datafusion/expr/src/expr_schema.rs
+++ b/datafusion/expr/src/expr_schema.rs
@@ -152,6 +152,7 @@ impl ExprSchemable for Expr {
                     .map(|e| e.get_type(schema))
                     .collect::<Result<Vec<_>>>()?;

+
                 // Verify that function is invoked with correct number and type of arguments as defined in `TypeSignature`
                 let new_data_types = data_types_with_scalar_udf(&arg_data_types, func)
                     .map_err(|err| {
@@ -168,7 +169,7 @@ impl ExprSchemable for Expr {

                 // Perform additional function arguments validation (due to limited
                 // expressiveness of `TypeSignature`), then infer return type
-                Ok(func.return_type_from_exprs(args, schema, &new_data_types)?)
+                Ok(func.return_type_from_exprs(args, schema, &arg_data_types)?)
             }
             Expr::WindowFunction(window_function) => self
                 .data_type_and_nullable_with_window_function(schema, window_function)

Which basically says pass the input data types directly to the function call rather than calling data_types_with_scalar_udf first (which claims to type coercion)

datafusion/datafusion/expr/src/expr_schema.rs

Line 171 in 68ead28

Ok(func.return_type_from_exprs(args, schema, &new_data_types)?)

🤔 this looks like it was added in Sep via 1b3608d (before that the input types were passed directly) 🤔

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't seem right to me that ExprSchema is coercing the arguments (implicitly) to me 🤔

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems like there is something about FixedSizeList that is causing issues to me

correct, #13756 (comment)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Weird, when I remove this line in expr schema the test passes (with FixedSizedList):

i did the same, basically removing this block

datafusion/datafusion/expr/src/expr_schema.rs

Lines 155 to 167 in b30c200

// Verify that function is invoked with correct number and type of arguments as defined in `TypeSignature`
let new_data_types = data_types_with_scalar_udf(&arg_data_types, func)
.map_err(|err| {
plan_datafusion_err!(
"{} {}",
err,
utils::generate_signature_error_msg(
func.name(),
func.signature().clone(),
&arg_data_types,
)
)
})?;

it's enough to fix the unit test in this PR
but other things start to fail

It doesn't seem right to me that ExprSchema is coercing the arguments (implicitly) to me 🤔

agreed

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't seem right to me that ExprSchema is coercing the arguments (implicitly) to me 🤔

We need to get the return_type of the function here and the arguments of return_type is the "coerced data type" therefore I think new_data_types is the right choice to me.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in a logical plan, the function arguments should already be of the right coerced type, we should just use them.
they may be not required type on the early plan building phase (when plan is still syntactical, not semantical), which unfortuantely uses the same LogicalPlan and Expr types. #12604 would address that.

Field::new("some_arbitrary_test_field", DataType::Int32, false).into(),
13,
);
let array_type =
DataType::List(Field::new_list_field(complex_type.clone(), true).into());
let index_type = DataType::Int64;

let schema = DFSchema::from_unqualified_fields(
vec![
Field::new("my_array", array_type.clone(), false),
Field::new("my_index", index_type.clone(), false),
]
.into(),
HashMap::default(),
)
.unwrap();

let udf = array_element_udf();

// ScalarUDFImpl::return_type
assert_eq!(
udf.return_type(&[array_type.clone(), index_type.clone()])
.unwrap(),
complex_type
);

// ScalarUDFImpl::return_type_from_exprs with typed exprs
assert_eq!(
udf.return_type_from_exprs(
&[
cast(Expr::Literal(ScalarValue::Null), array_type.clone()),
cast(Expr::Literal(ScalarValue::Null), index_type.clone()),
],
&schema,
&[array_type.clone(), index_type.clone()]
)
.unwrap(),
complex_type
);

// ScalarUDFImpl::return_type_from_exprs with exprs not carrying type
assert_eq!(
udf.return_type_from_exprs(
&[
Expr::Column(Column::new_unqualified("my_array")),
Expr::Column(Column::new_unqualified("my_index")),
],
&schema,
&[array_type.clone(), index_type.clone()]
)
.unwrap(),
complex_type
);

// Via ExprSchemable::get_type (e.g. SimplifyInfo)
let udf_expr = Expr::ScalarFunction(ScalarFunction {
func: array_element_udf(),
args: vec![
Expr::Column(Column::new_unqualified("my_array")),
Expr::Column(Column::new_unqualified("my_index")),
],
});
assert_eq!(
ExprSchemable::get_type(&udf_expr, &schema).unwrap(),
complex_type
);
Comment on lines +1071 to +1074
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This didn't pass before the change. The assertions above did pass.

}
}
Loading