Fix get_type for higher-order array functions #13756

findepi · 2024-12-13T13:32:23Z

Which issue does this PR close?

Rationale for this change

Fix a bug, see issue #13755
TL;DR: fix incorrect result of ExprSchemable::get_type for an array function invoked on array of list

What changes are included in this PR?

Just the fix

Are these changes tested?

unit test

Are there any user-facing changes?

yes

findepi · 2024-12-13T13:33:26Z

datafusion/functions-nested/src/extract.rs

+        assert_eq!(
+            ExprSchemable::get_type(&udf_expr, &schema).unwrap(),
+            complex_type
+        );


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

datafusion/functions-nested/src/flatten.rs

The fix is covered by recursive flatten test case in array.slt

jayzhan211 · 2024-12-13T14:11:52Z

datafusion/expr/src/type_coercion/functions.rs

+        }
+    }
+
+    fn recursive_array(array_type: &DataType) -> Option<DataType> {


Can we extend the existing array function for nested array instead of creating another signature for nested array

I don't know how to do this, please advise!
But this function should go away with #13757.

But this function should go away with #13757.

I don't understand -- if the goal is to remove recursive flattening, should we be adding new code to support it 🤔

the pre-existing array signature implied recursively array-infication (replacing FixedLengthList with List, recursively), didn't imply flattening.

the recursive type normalization matters for flatten only, cause it (currently) operates recursively and otherwise would need to gain code to handle FixedLengthList inputs

the recursive array-ification was useless for other array functions and was made non-recursive.
to compensate for this change, new RecursiveArray signature was added for flatten case.

jayzhan211 · 2024-12-13T14:12:48Z

datafusion/functions-nested/src/extract.rs

+    use std::collections::HashMap;
+
+    #[test]
+    fn test_array_element_return_type() {


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

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?

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.

alamb

THanks @findepi and @jayzhan211

From what I can see the point of this PR is to make array_element_udf have different type resolution rules (non recursive), which seems reasonable

However, as you both mention I don't seem to be able to trigger the problem from SQL (element access seems to work correctly): (e.g. the [[20]] isn't flattened in on main:

> create table t as values ([[[10]], [[20]]]);
0 row(s) fetched.
Elapsed 0.007 seconds.

> explain select column1[2] from t;
+---------------+---------------------------------------------------------------------------+
| plan_type     | plan                                                                      |
+---------------+---------------------------------------------------------------------------+
| logical_plan  | Projection: array_element(t.column1, Int64(2))                            |
|               |   TableScan: t projection=[column1]                                       |
| physical_plan | ProjectionExec: expr=[array_element(column1@0, 2) as t.column1[Int64(2)]] |
|               |   MemoryExec: partitions=1, partition_sizes=[1]                           |
|               |                                                                           |
+---------------+---------------------------------------------------------------------------+
2 row(s) fetched.
Elapsed 0.001 seconds.

> select column1[2] from t;
+---------------------+
| t.column1[Int64(2)] |
+---------------------+
| [[20]]              |
+---------------------+
1 row(s) fetched.
Elapsed 0.001 seconds.

And the type seems good too list(list(int))

> select arrow_typeof(column1[2]) from t;
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| arrow_typeof(t.column1[Int64(2)])                                                                                                                                                                                           |
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| List(Field { name: "item", data_type: List(Field { name: "item", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }) |
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

So this problem is quite strange. How is it working today (without this change) 🤔

alamb · 2024-12-13T20:02:36Z

datafusion/expr/src/type_coercion/functions.rs

+        }
+    }
+
+    fn recursive_array(array_type: &DataType) -> Option<DataType> {


But this function should go away with #13757.

I don't understand -- if the goal is to remove recursive flattening, should we be adding new code to support it 🤔

findepi · 2024-12-13T20:40:49Z

So this problem is quite strange. How is it working today (without this change) 🤔

i believe the bug -- if we agree this is a bug -- is compensated by other factors.
For example, at early planning stage it's totally OK to change expression types.
Later, such a change triggers schema change assertion.

I found this bug in case where array_element was inserted into the plan as a result of ScalarUDFImpl::simplify. At this stage it's "loose typing" is no longer OK.

@alamb @jayzhan211 can you please review the attached unit test?
Does it look sound, ie should it pass?
Does it pass for you without other changes from the PR?

alamb · 2024-12-13T21:19:23Z

I am checking this out in more detail

alamb

I am still digging. This is so weird

I messed with the test and it seems like the failure only happens when the complex type is a FixedSizeList for some reason..

alamb · 2024-12-13T21:16:00Z

datafusion/expr/src/type_coercion/functions.rs

    fn array(array_type: &DataType) -> Option<DataType> {
+        match array_type {


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

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

alamb · 2024-12-13T21:24:24Z

datafusion/functions-nested/src/extract.rs

+
+    #[test]
+    fn test_array_element_return_type() {
+        let complex_type = DataType::FixedSizeList(


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

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) 🤔

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

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

correct, #13756 (comment)

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

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.

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.

findepi · 2024-12-13T21:37:02Z

I messed with the test and it seems like the failure only happens when the complex type is a FixedSizeList for some reason..

because coerced_fixed_size_list_to_list called here is recursive

datafusion/datafusion/expr/src/type_coercion/functions.rs

Line 422 in 55e56c4

let array_type = coerced_fixed_size_list_to_list(array_type);

jayzhan211 · 2024-12-14T01:33:15Z

ExprSchemable::get_type for ScalarFunction is basically asking the return_type for the function. Given that we coerce fixed size list to list, the return type of array_element(fixed size list) makes sense to be list. Therefore, I think the unit test is expected to fail since it is coerced to List

findepi · 2024-12-14T14:39:26Z

Given that we coerce fixed size list to list, the return type of array_element(fixed size list) makes sense to be list.

in the unit test, we ask for array_element(list(fixed size list)) and we expect the return type to be fixed size list.
in the fix, we make so that array_element(list(T)) always returns T.

Fix get_type for higher-order array functions

6903259

github-actions bot added the logical-expr Logical plan and expressions label Dec 13, 2024

findepi commented Dec 13, 2024

View reviewed changes

datafusion/functions-nested/src/flatten.rs Show resolved Hide resolved

Fix recursive flatten

6d81418

The fix is covered by recursive flatten test case in array.slt

findepi force-pushed the findepi/array-get-type branch from 1bd311a to 6d81418 Compare December 13, 2024 13:55

jayzhan211 reviewed Dec 13, 2024

View reviewed changes

findepi requested review from alamb, jayzhan211 and jonahgao December 13, 2024 15:39

alamb reviewed Dec 13, 2024

View reviewed changes

Restore "keep LargeList" in Array signature

038a015

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix get_type for higher-order array functions #13756

Fix get_type for higher-order array functions #13756

findepi commented Dec 13, 2024

findepi Dec 13, 2024

jayzhan211 Dec 13, 2024

findepi Dec 13, 2024

alamb Dec 13, 2024

findepi Dec 13, 2024

jayzhan211 Dec 13, 2024

findepi Dec 13, 2024

findepi Dec 13, 2024

alamb left a comment

alamb Dec 13, 2024

findepi commented Dec 13, 2024

alamb commented Dec 13, 2024

alamb left a comment

alamb Dec 13, 2024

findepi Dec 13, 2024

alamb Dec 13, 2024

alamb Dec 13, 2024

alamb Dec 13, 2024

findepi Dec 13, 2024

findepi Dec 13, 2024

jayzhan211 Dec 14, 2024

findepi Dec 14, 2024

findepi commented Dec 13, 2024

jayzhan211 commented Dec 14, 2024

findepi commented Dec 14, 2024

		fn array(array_type: &DataType) -> Option<DataType> {
		match array_type {

	// 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,
	)
	)
	})?;

Fix get_type for higher-order array functions #13756

Are you sure you want to change the base?

Fix get_type for higher-order array functions #13756

Conversation

findepi commented Dec 13, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

findepi commented Dec 13, 2024

alamb commented Dec 13, 2024

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

findepi commented Dec 13, 2024

jayzhan211 commented Dec 14, 2024

findepi commented Dec 14, 2024