diff --git a/datafusion/functions-nested/src/distance.rs b/datafusion/functions-nested/src/distance.rs index edf1806b66c2d..c9aec816676a7 100644 --- a/datafusion/functions-nested/src/distance.rs +++ b/datafusion/functions-nested/src/distance.rs @@ -18,9 +18,7 @@ //! [ScalarUDFImpl] definitions for array_distance function. use crate::utils::make_scalar_function; -use arrow::array::{ - Array, ArrayRef, Float64Array, LargeListArray, ListArray, OffsetSizeTrait, -}; +use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; use arrow::datatypes::{ DataType, DataType::{FixedSizeList, LargeList, List, Null}, @@ -35,7 +33,6 @@ use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; -use datafusion_functions::downcast_arg; use datafusion_macros::user_doc; use itertools::Itertools; use std::sync::Arc; @@ -44,13 +41,13 @@ make_udf_expr_and_func!( ArrayDistance, array_distance, array, - "returns the Euclidean distance between two numeric arrays.", + "returns the Euclidean distance between two one-dimensional numeric arrays.", array_distance_udf ); #[user_doc( doc_section(label = "Array Functions"), - description = "Returns the Euclidean distance between two input arrays of equal length.", + description = "Returns the Euclidean distance between two one-dimensional input arrays of equal length.", syntax_example = "array_distance(array1, array2)", sql_example = r#"```sql > select array_distance([1, 2], [1, 4]); @@ -106,16 +103,30 @@ impl ScalarUDFImpl for ArrayDistance { fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let [_, _] = take_function_args(self.name(), arg_types)?; let coercion = Some(&ListCoercion::FixedSizedListToList); - let arg_types = arg_types.iter().map(|arg_type| { - if matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + let arg_types = arg_types.iter().map(|arg_type| match arg_type { + Null => Ok(coerced_type_with_base_type_only( + arg_type, + &DataType::Float64, + coercion, + )), + List(field) | LargeList(field) | FixedSizeList(field, _) => { + // Distance between nested lists is not supported + if matches!( + field.data_type(), + List(_) | LargeList(_) | FixedSizeList(..) + ) { + return plan_err!( + "{} only supports one-dimensional arrays, got {arg_type}", + self.name() + ); + } Ok(coerced_type_with_base_type_only( arg_type, &DataType::Float64, coercion, )) - } else { - plan_err!("{} does not support type {arg_type}", self.name()) } + _ => plan_err!("{} does not support type {arg_type}", self.name()), }); arg_types.try_collect() @@ -172,43 +183,6 @@ fn compute_array_distance( None => return Ok(None), }; - let mut value1 = value1; - let mut value2 = value2; - - loop { - match value1.data_type() { - List(_) => { - if downcast_arg!(value1, ListArray).null_count() > 0 { - return Ok(None); - } - value1 = downcast_arg!(value1, ListArray).value(0); - } - LargeList(_) => { - if downcast_arg!(value1, LargeListArray).null_count() > 0 { - return Ok(None); - } - value1 = downcast_arg!(value1, LargeListArray).value(0); - } - _ => break, - } - - match value2.data_type() { - List(_) => { - if downcast_arg!(value2, ListArray).null_count() > 0 { - return Ok(None); - } - value2 = downcast_arg!(value2, ListArray).value(0); - } - LargeList(_) => { - if downcast_arg!(value2, LargeListArray).null_count() > 0 { - return Ok(None); - } - value2 = downcast_arg!(value2, LargeListArray).value(0); - } - _ => break, - } - } - // Check for NULL values inside the arrays if value1.null_count() != 0 || value2.null_count() != 0 { return Ok(None); diff --git a/datafusion/sqllogictest/test_files/array/array_length.slt b/datafusion/sqllogictest/test_files/array/array_length.slt index 1bb5382339854..7741d815bc234 100644 --- a/datafusion/sqllogictest/test_files/array/array_length.slt +++ b/datafusion/sqllogictest/test_files/array/array_length.slt @@ -159,21 +159,6 @@ select array_distance([2], [3]), list_distance([1], [2]), list_distance([1], [-2 query error select list_distance([1], [1, 2]); -query R -select array_distance([[1, 1]], [1, 2]); ----- -1 - -query R -select array_distance([[1, 1]], [[1, 2]]); ----- -1 - -query R -select array_distance([[1, 1]], [[1, 2]]); ----- -1 - query RR select array_distance([1, 1, 0, 0], [2, 2, 1, 1]), list_distance([1, 2, 3], [1, 2, 3]); ---- @@ -204,6 +189,40 @@ select list_distance([1, 2, 3], [1, 2, 3]) AS distance; ---- 0 +# array_distance with null outer arrays +query RR +select + array_distance(arrow_cast(NULL, 'List(Float64)'), [1, 2]), + array_distance([1, 2], arrow_cast(NULL, 'List(Float64)')); +---- +NULL NULL + +# invalid argument count and types +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance(); + +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance([1]); + +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance([1], [2], [3]); + +query error array_distance does not support type Int64 +select array_distance(1, [1]); + +query error array_distance does not support types +select array_distance([1], arrow_cast([1], 'LargeList(Float64)')); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 1]], [1, 2]); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 1]], [[1, 2]]); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 2], [100, 100]], [[1, 4], [0, 0]]); + + # array_distance with columns query RRR select array_distance(column1, column2), array_distance(column1, column3), array_distance(column1, column4) from arrays_distance_table;