From 9b22063e00fddb57de0a26eb9729f7ccf223166c Mon Sep 17 00:00:00 2001 From: Weijun-H Date: Tue, 13 Feb 2024 18:56:21 +0800 Subject: [PATCH] Fix array_element function signature --- datafusion/expr/src/built_in_function.rs | 2 +- datafusion/expr/src/signature.rs | 92 ++++++++++--------- .../expr/src/type_coercion/functions.rs | 4 +- .../physical-expr/src/array_expressions.rs | 1 + datafusion/sqllogictest/test_files/array.slt | 4 +- 5 files changed, 53 insertions(+), 50 deletions(-) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index e57a52bcf2464..6fa41495f0f70 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -954,7 +954,7 @@ impl BuiltinScalarFunction { Signature::array(true, self.volatility()) } BuiltinScalarFunction::ArrayElement => { - Signature::array_and_index(false, self.volatility()) + Signature::array_and_index(true, self.volatility()) } BuiltinScalarFunction::ArrayExcept => Signature::any(2, self.volatility()), BuiltinScalarFunction::Flatten => Signature::array(true, self.volatility()), diff --git a/datafusion/expr/src/signature.rs b/datafusion/expr/src/signature.rs index 3b478322e7971..0aff2db60cda3 100644 --- a/datafusion/expr/src/signature.rs +++ b/datafusion/expr/src/signature.rs @@ -122,29 +122,32 @@ pub enum TypeSignature { /// is `OneOf(vec![Any(0), VariadicAny])`. OneOf(Vec), /// Specifies Signatures for array functions - /// Boolean value specifies whether null type coercion is allowed - ArraySignature(ArrayFunctionSignature, bool), + ArraySignature(ArrayFunctionSignature), } #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ArrayFunctionSignature { /// Specialized Signature for ArrayAppend and similar functions + /// If `allow_null` is true, the function also accepts a single argument of type Null. /// The first argument should be List/LargeList/FixedSizedList, and the second argument should be non-list or list. /// The second argument's list dimension should be one dimension less than the first argument's list dimension. /// List dimension of the List/LargeList is equivalent to the number of List. /// List dimension of the non-list is 0. - ArrayAndElement, + ArrayAndElement(bool), /// Specialized Signature for ArrayPrepend and similar functions + /// If `allow_null` is true, the function also accepts a single argument of type Null. /// The first argument should be non-list or list, and the second argument should be List/LargeList. /// The first argument's list dimension should be one dimension less than the second argument's list dimension. - ElementAndArray, + ElementAndArray(bool), /// Specialized Signature for ArrayEqual and similar functions + /// If `allow_null` is true, the function also accepts a single argument of type Null. /// The first argument should be List/LargeList/FixedSizedList, and the second argument should be Int64. - ArrayAndIndex, + ArrayAndIndex(bool), /// Specialized Signature for ArrayEmpty and similar functions /// The function takes a single argument that must be a List/LargeList/FixedSizeList /// or something that can be coerced to one of those types. - Array, + /// If `allow_null` is true, the function also accepts a single argument of type Null. + Array(bool), } impl ArrayFunctionSignature { @@ -155,7 +158,6 @@ impl ArrayFunctionSignature { pub fn get_type_signature( &self, current_types: &[DataType], - allow_null_coercion: bool, ) -> Result>> { fn array_append_or_prepend_valid_types( current_types: &[DataType], @@ -208,13 +210,20 @@ impl ArrayFunctionSignature { _ => Ok(vec![vec![]]), } } - fn array_and_index(current_types: &[DataType]) -> Result>> { + fn array_and_index( + current_types: &[DataType], + allow_null_coercion: bool, + ) -> Result>> { if current_types.len() != 2 { return Ok(vec![vec![]]); } let array_type = ¤t_types[0]; + if array_type.eq(&DataType::Null) && !allow_null_coercion { + return Ok(vec![vec![]]); + } + match array_type { DataType::List(_) | DataType::LargeList(_) @@ -222,6 +231,7 @@ impl ArrayFunctionSignature { let array_type = coerced_fixed_size_list_to_list(array_type); Ok(vec![vec![array_type, DataType::Int64]]) } + DataType::Null => Ok(vec![vec![array_type.clone(), DataType::Int64]]), _ => Ok(vec![vec![]]), } } @@ -253,22 +263,18 @@ impl ArrayFunctionSignature { } } match self { - ArrayFunctionSignature::ArrayAndElement => { - array_append_or_prepend_valid_types( - current_types, - true, - allow_null_coercion, - ) + ArrayFunctionSignature::ArrayAndElement(allow_null) => { + array_append_or_prepend_valid_types(current_types, true, *allow_null) } - ArrayFunctionSignature::ElementAndArray => { - array_append_or_prepend_valid_types( - current_types, - false, - allow_null_coercion, - ) + ArrayFunctionSignature::ElementAndArray(allow_null) => { + array_append_or_prepend_valid_types(current_types, false, *allow_null) + } + ArrayFunctionSignature::ArrayAndIndex(allow_null) => { + array_and_index(current_types, *allow_null) + } + ArrayFunctionSignature::Array(allow_null) => { + array(current_types, *allow_null) } - ArrayFunctionSignature::ArrayAndIndex => array_and_index(current_types), - ArrayFunctionSignature::Array => array(current_types, allow_null_coercion), } } } @@ -276,17 +282,17 @@ impl ArrayFunctionSignature { impl std::fmt::Display for ArrayFunctionSignature { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - ArrayFunctionSignature::ArrayAndElement => { - write!(f, "array, element") + ArrayFunctionSignature::ArrayAndElement(allow_null) => { + write!(f, "ArrayAndElement({})", *allow_null) } - ArrayFunctionSignature::ElementAndArray => { - write!(f, "element, array") + ArrayFunctionSignature::ElementAndArray(allow_null) => { + write!(f, "ElementAndArray({})", *allow_null) } - ArrayFunctionSignature::ArrayAndIndex => { - write!(f, "array, index") + ArrayFunctionSignature::ArrayAndIndex(allow_null) => { + write!(f, "ArrayAndIndex({})", *allow_null) } - ArrayFunctionSignature::Array => { - write!(f, "array") + ArrayFunctionSignature::Array(allow_null) => { + write!(f, "Array({})", *allow_null) } } } @@ -320,7 +326,7 @@ impl TypeSignature { TypeSignature::OneOf(sigs) => { sigs.iter().flat_map(|s| s.to_string_repr()).collect() } - TypeSignature::ArraySignature(array_signature, _) => { + TypeSignature::ArraySignature(array_signature) => { vec![array_signature.to_string()] } } @@ -425,42 +431,38 @@ impl Signature { } } /// Specialized Signature for ArrayAppend and similar functions - pub fn array_and_element(allow_null_coercion: bool, volatility: Volatility) -> Self { + pub fn array_and_element(allow_null: bool, volatility: Volatility) -> Self { Signature { type_signature: TypeSignature::ArraySignature( - ArrayFunctionSignature::ArrayAndElement, - allow_null_coercion, + ArrayFunctionSignature::ArrayAndElement(allow_null), ), volatility, } } /// Specialized Signature for ArrayPrepend and similar functions - pub fn element_and_array(allow_null_coercion: bool, volatility: Volatility) -> Self { + pub fn element_and_array(allow_null: bool, volatility: Volatility) -> Self { Signature { type_signature: TypeSignature::ArraySignature( - ArrayFunctionSignature::ElementAndArray, - allow_null_coercion, + ArrayFunctionSignature::ElementAndArray(allow_null), ), volatility, } } /// Specialized Signature for ArrayElement and similar functions - pub fn array_and_index(allow_null_coercion: bool, volatility: Volatility) -> Self { + pub fn array_and_index(allow_null: bool, volatility: Volatility) -> Self { Signature { type_signature: TypeSignature::ArraySignature( - ArrayFunctionSignature::ArrayAndIndex, - allow_null_coercion, + ArrayFunctionSignature::ArrayAndIndex(allow_null), ), volatility, } } /// Specialized Signature for ArrayEmpty and similar functions - pub fn array(allow_null_coercion: bool, volatility: Volatility) -> Self { + pub fn array(allow_null: bool, volatility: Volatility) -> Self { Signature { - type_signature: TypeSignature::ArraySignature( - ArrayFunctionSignature::Array, - allow_null_coercion, - ), + type_signature: TypeSignature::ArraySignature(ArrayFunctionSignature::Array( + allow_null, + )), volatility, } } diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index b0054aa28e161..4e5a2f1b69552 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -110,8 +110,8 @@ fn get_valid_types( } TypeSignature::Exact(valid_types) => vec![valid_types.clone()], - TypeSignature::ArraySignature(ref function_signature, allow_null_coercion) => { - function_signature.get_type_signature(current_types, *allow_null_coercion)? + TypeSignature::ArraySignature(ref function_signature) => { + function_signature.get_type_signature(current_types)? } TypeSignature::Any(number) => { diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 15cd24f2f073d..50c70ccfdb111 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -433,6 +433,7 @@ pub fn array_element(args: &[ArrayRef]) -> Result { let indexes = as_int64_array(&args[1])?; general_array_element::(array, indexes) } + DataType::Null => Ok(args[0].clone()), _ => exec_err!( "array_element does not support type: {:?}", args[0].data_type() diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index 0b8870a771c2d..789238b5723b4 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -1084,7 +1084,7 @@ from arrays_values_without_nulls; ## array_element (aliases: array_extract, list_extract, list_element) # array_element error -query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(array, index\) +query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_element\(Int64, Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_element\(ArrayAndIndex\(true\)\) select array_element(1, 2); # array_element with null @@ -4267,7 +4267,7 @@ NULL 10 ## array_dims (aliases: `list_dims`) # array dims error -query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(array\) +query error DataFusion error: Error during planning: No function matches the given name and argument types 'array_dims\(Int64\)'. You might need to add explicit type casts.\n\tCandidate functions:\n\tarray_dims\(Array\(false\)\) select array_dims(1); # array_dims scalar function