sgrebnov commented on code in PR #15149:
URL: https://github.com/apache/datafusion/pull/15149#discussion_r2296274979
##########
datafusion/expr/src/type_coercion/functions.rs:
##########
@@ -364,98 +366,67 @@ fn get_valid_types(
return Ok(vec![vec![]]);
}
- let array_idx = arguments.iter().enumerate().find_map(|(idx, arg)| {
- if *arg == ArrayFunctionArgument::Array {
- Some(idx)
- } else {
- None
- }
- });
- let Some(array_idx) = array_idx else {
- return Err(internal_datafusion_err!("Function '{function_name}'
expected at least one argument array argument"));
- };
- let Some(array_type) = array(¤t_types[array_idx]) 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 mut new_base_type =
datafusion_common::utils::base_type(&array_type);
- for (current_type, argument_type) in
current_types.iter().zip(arguments.iter()) {
- match argument_type {
- ArrayFunctionArgument::Element | ArrayFunctionArgument::Array
=> {
- new_base_type =
- coerce_array_types(function_name, current_type,
&new_base_type)?;
+ let mut large_list = false;
+ let mut fixed_size = array_coercion !=
Some(&ListCoercion::FixedSizedListToList);
+ let mut list_sizes = Vec::with_capacity(arguments.len());
+ let mut element_types = Vec::with_capacity(arguments.len());
+ for (argument, current_type) in
arguments.iter().zip(current_types.iter()) {
+ match argument {
+ ArrayFunctionArgument::Index | ArrayFunctionArgument::String
=> (),
+ ArrayFunctionArgument::Element => {
+ element_types.push(current_type.clone())
}
- ArrayFunctionArgument::Index | ArrayFunctionArgument::String
=> {}
+ ArrayFunctionArgument::Array => match current_type {
+ DataType::Null => element_types.push(DataType::Null),
+ DataType::List(field) => {
+ element_types.push(field.data_type().clone());
+ fixed_size = false;
+ }
+ DataType::LargeList(field) => {
+ element_types.push(field.data_type().clone());
+ large_list = true;
+ fixed_size = false;
+ }
+ DataType::FixedSizeList(field, size) => {
+ element_types.push(field.data_type().clone());
+ list_sizes.push(*size)
+ }
+ arg_type => {
+ plan_err!("{function_name} does not support type
{arg_type}")?
+ }
+ },
}
}
- let new_array_type =
datafusion_common::utils::coerced_type_with_base_type_only(
- &array_type,
- &new_base_type,
- array_coercion,
- );
- let new_elem_type = match new_array_type {
- DataType::List(ref field)
- | DataType::LargeList(ref field)
- | DataType::FixedSizeList(ref field, _) => field.data_type(),
- _ => return Ok(vec![vec![]]),
+ let Some(element_type) = type_union_resolution(&element_types) else {
+ return Ok(vec![vec![]]);
};
- let mut valid_types = Vec::with_capacity(arguments.len());
- for (current_type, argument_type) in
current_types.iter().zip(arguments.iter()) {
- let valid_type = match argument_type {
- ArrayFunctionArgument::Element => new_elem_type.clone(),
+ if !fixed_size {
+ list_sizes.clear()
+ }
+
+ let mut list_sizes = list_sizes.into_iter();
+ let valid_types = arguments.iter().zip(current_types.iter()).map(
+ |(argument_type, current_type)| match argument_type {
ArrayFunctionArgument::Index => DataType::Int64,
ArrayFunctionArgument::String => DataType::Utf8,
+ ArrayFunctionArgument::Element => element_type.clone(),
ArrayFunctionArgument::Array => {
- let Some(current_type) = array(current_type) else {
- return Ok(vec![vec![]]);
- };
- let new_type =
-
datafusion_common::utils::coerced_type_with_base_type_only(
- ¤t_type,
- &new_base_type,
- array_coercion,
- );
- // All array arguments must be coercible to the same type
- if new_type != new_array_type {
- return Ok(vec![vec![]]);
+ if current_type.is_null() {
+ DataType::Null
+ } else if large_list {
+ DataType::new_large_list(element_type.clone(), true)
+ } else if let Some(size) = list_sizes.next() {
+ DataType::new_fixed_size_list(element_type.clone(),
size, true)
+ } else {
+ DataType::new_list(element_type.clone(), true)
Review Comment:
It seems we are loosing information on original type elements nullability
and just always return nullable:true here
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
