jayzhan211 commented on code in PR #9108:
URL: https://github.com/apache/arrow-datafusion/pull/9108#discussion_r1499278949
##########
datafusion/expr/src/signature.rs:
##########
@@ -123,35 +126,211 @@ pub enum TypeSignature {
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ArrayFunctionSignature {
/// Specialized Signature for ArrayAppend and similar functions
- /// The first argument should be List/LargeList, and the second argument
should be non-list or list.
+ /// 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 Array functions of the form (List/LargeList,
Index)
- ArrayAndIndex,
+ /// 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(bool),
/// Specialized Signature for Array functions of the form (List/LargeList,
Element, Optional Index)
ArrayAndElementAndOptionalIndex,
+ /// 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.
+ /// If `allow_null` is true, the function also accepts a single argument
of type Null.
+ Array(bool),
+}
+
+impl ArrayFunctionSignature {
+ /// Arguments to ArrayFunctionSignature
+ /// `current_types` - The data types of the arguments
+ /// `allow_null_coercion` - Whether null type coercion is allowed
+ /// Returns the valid types for the function signature
+ pub fn get_type_signature(
+ &self,
+ current_types: &[DataType],
+ ) -> Result<Vec<Vec<DataType>>> {
+ fn array_element_and_optional_index(
+ current_types: &[DataType],
+ ) -> Result<Vec<Vec<DataType>>> {
+ // make sure there's 2 or 3 arguments
+ if !(current_types.len() == 2 || current_types.len() == 3) {
+ return Ok(vec![vec![]]);
+ }
+
+ let first_two_types = ¤t_types[0..2];
+ let mut valid_types =
+ array_append_or_prepend_valid_types(first_two_types, true,
true)?;
+
+ // Early return if there are only 2 arguments
+ if current_types.len() == 2 {
+ return Ok(valid_types);
+ }
+
+ let valid_types_with_index = valid_types
+ .iter()
+ .map(|t| {
+ let mut t = t.clone();
+ t.push(DataType::Int64);
+ t
+ })
+ .collect::<Vec<_>>();
+
+ valid_types.extend(valid_types_with_index);
+
+ Ok(valid_types)
+ }
+
+ fn array_append_or_prepend_valid_types(
+ current_types: &[DataType],
+ is_append: bool,
+ allow_null_coercion: bool,
+ ) -> Result<Vec<Vec<DataType>>> {
+ if current_types.len() != 2 {
+ return Ok(vec![vec![]]);
+ }
+
+ let (array_type, elem_type) = if is_append {
+ (¤t_types[0], ¤t_types[1])
+ } else {
+ (¤t_types[1], ¤t_types[0])
+ };
+
+ // We follow Postgres on `array_append(Null, T)`, which is not
valid.
+ if array_type.eq(&DataType::Null) {
+ if allow_null_coercion {
+ return Ok(vec![vec![array_type.clone(),
elem_type.clone()]]);
+ } else {
+ return Ok(vec![vec![]]);
+ }
+ }
+
+ // We need to find the coerced base type, mainly for cases like:
+ // `array_append(List(null), i64)` -> `List(i64)`
+ let array_base_type =
datafusion_common::utils::base_type(array_type);
+ let elem_base_type =
datafusion_common::utils::base_type(elem_type);
+ let new_base_type = comparison_coercion(&array_base_type,
&elem_base_type);
+
+ let new_base_type = new_base_type.ok_or_else(|| {
+ internal_datafusion_err!(
+ "Coercion from {array_base_type:?} to {elem_base_type:?}
not supported."
+ )
+ })?;
+
+ let new_array_type =
+ datafusion_common::utils::coerced_type_with_base_type_only(
+ array_type,
+ &new_base_type,
+ );
+
+ match new_array_type {
+ DataType::List(ref field)
+ | DataType::LargeList(ref field)
+ | DataType::FixedSizeList(ref field, _) => {
+ let new_elem_type = field.data_type();
+ if is_append {
+ Ok(vec![vec![new_array_type.clone(),
new_elem_type.clone()]])
+ } else {
+ Ok(vec![vec![new_elem_type.to_owned(),
new_array_type.clone()]])
+ }
+ }
+ _ => Ok(vec![vec![]]),
+ }
+ }
+ fn array_and_index(
+ current_types: &[DataType],
+ allow_null_coercion: bool,
+ ) -> Result<Vec<Vec<DataType>>> {
+ 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(_)
+ | DataType::FixedSizeList(_, _) => {
+ let array_type =
coerced_fixed_size_list_to_list(array_type);
+ Ok(vec![vec![array_type, DataType::Int64]])
Review Comment:
I still think it can be further simplied if there is a function that process
`coerced_fixed_size_list_to_list` and return `array_type`. Then can be used
both in `array_and_index` and `array`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
