From 18de0bff108dc63775dfcc3a65de89ce2c2da57b Mon Sep 17 00:00:00 2001 From: veeupup Date: Mon, 27 Nov 2023 22:00:29 +0800 Subject: [PATCH 1/2] fix array_append with null array Signed-off-by: veeupup --- datafusion/expr/src/built_in_function.rs | 18 +++++++- .../physical-expr/src/array_expressions.rs | 41 ++++++++----------- datafusion/sqllogictest/test_files/array.slt | 16 ++++---- 3 files changed, 43 insertions(+), 32 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index c511c752b4d7..4c6fc740227b 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -536,7 +536,23 @@ impl BuiltinScalarFunction { let data_type = get_base_type(&input_expr_types[0])?; Ok(data_type) } - BuiltinScalarFunction::ArrayAppend => Ok(input_expr_types[0].clone()), + BuiltinScalarFunction::ArrayAppend => { + if let DataType::List(field) = input_expr_types[0].clone() { + if field.data_type().equals_datatype(&DataType::Null) { + Ok(DataType::List(Arc::new(Field::new( + "item", + input_expr_types[1].clone(), + true, + )))) + } else { + Ok(input_expr_types[0].clone()) + } + } else { + plan_err!( + "The {self} function can only accept list as the first argument" + ) + } + } BuiltinScalarFunction::ArrayConcat => { let mut expr_type = Null; let mut max_dims = 0; diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 8968bcf2ea4e..89a4744f748c 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -227,16 +227,25 @@ fn compute_array_dims(arr: Option) -> Result>>> } fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> { - let data_type = args[0].data_type(); - if !args.iter().all(|arg| { - arg.data_type().equals_datatype(data_type) - || arg.data_type().equals_datatype(&DataType::Null) - }) { - let types = args.iter().map(|arg| arg.data_type()).collect::>(); - return plan_err!("{name} received incompatible types: '{types:?}'."); + let mut data_types = args + .iter() + .map(|arg| arg.data_type()) + .collect::>(); + match data_types.len() { + 1 => Ok(()), + 2 => { + if data_types.remove(&DataType::Null) { + Ok(()) + } else { + let types = args.iter().map(|arg| arg.data_type()).collect::>(); + plan_err!("{name} received incompatible types: '{types:?}'.") + } + } + _ => { + let types = args.iter().map(|arg| arg.data_type()).collect::>(); + plan_err!("{name} received incompatible types: '{types:?}'.") + } } - - Ok(()) } macro_rules! call_array_function { @@ -2951,20 +2960,6 @@ mod tests { assert_eq!(result, &UInt64Array::from_value(2, 1)); } - #[test] - fn test_check_invalid_datatypes() { - let data = vec![Some(vec![Some(1), Some(2), Some(3)])]; - let list_array = - Arc::new(ListArray::from_iter_primitive::(data)) as ArrayRef; - let int64_array = Arc::new(StringArray::from(vec![Some("string")])) as ArrayRef; - - let args = [list_array.clone(), int64_array.clone()]; - - let array = array_append(&args); - - assert_eq!(array.unwrap_err().strip_backtrace(), "Error during planning: array_append received incompatible types: '[Int64, Utf8]'."); - } - fn return_array() -> ArrayRef { // Returns: [1, 2, 3, 4] let args = [ diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index d33555509e6c..0687dcb3746b 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -1056,16 +1056,16 @@ select make_array(['a','b'], null); # TODO: array_append with NULLs # array_append scalar function #1 -# query ? -# select array_append(make_array(), 4); -# ---- -# [4] +query ? +select array_append(make_array(), 4); +---- +[4] # array_append scalar function #2 -# query ?? -# select array_append(make_array(), make_array()), array_append(make_array(), make_array(4)); -# ---- -# [[]] [[4]] +query ?? +select array_append(make_array(), make_array()), array_append(make_array(), make_array(4)); +---- +[[]] [[4]] # array_append scalar function #3 query ??? From b0ccf19162f7b4288af8b2df20aed63c2bae103a Mon Sep 17 00:00:00 2001 From: veeupup Date: Wed, 29 Nov 2023 21:50:05 +0800 Subject: [PATCH 2/2] fix comments Signed-off-by: veeupup --- datafusion/expr/src/built_in_function.rs | 4 +--- datafusion/physical-expr/src/array_expressions.rs | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index 4c6fc740227b..78e7f4f826da 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -548,9 +548,7 @@ impl BuiltinScalarFunction { Ok(input_expr_types[0].clone()) } } else { - plan_err!( - "The {self} function can only accept list as the first argument" - ) + plan_err!("The {self} function expects LIST as the first argument") } } BuiltinScalarFunction::ArrayConcat => { diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 89a4744f748c..f6896d0cb6d9 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -227,14 +227,14 @@ fn compute_array_dims(arr: Option) -> Result>>> } fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> { - let mut data_types = args + let data_types = args .iter() .map(|arg| arg.data_type()) .collect::>(); match data_types.len() { 1 => Ok(()), 2 => { - if data_types.remove(&DataType::Null) { + if data_types.contains(&DataType::Null) { Ok(()) } else { let types = args.iter().map(|arg| arg.data_type()).collect::>();