datafusion_functions_nested/
distance.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [ScalarUDFImpl] definitions for array_distance function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
22use arrow::datatypes::{
23    DataType,
24    DataType::{FixedSizeList, LargeList, List, Null},
25};
26use datafusion_common::cast::{
27    as_float32_array, as_float64_array, as_generic_list_array, as_int32_array,
28    as_int64_array,
29};
30use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
31use datafusion_common::{Result, exec_err, plan_err, utils::take_function_args};
32use datafusion_expr::{
33    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
34    Volatility,
35};
36use datafusion_macros::user_doc;
37use itertools::Itertools;
38use std::sync::Arc;
39
40make_udf_expr_and_func!(
41    ArrayDistance,
42    array_distance,
43    array,
44    "returns the Euclidean distance between two one-dimensional numeric arrays.",
45    array_distance_udf
46);
47
48#[user_doc(
49    doc_section(label = "Array Functions"),
50    description = "Returns the Euclidean distance between two one-dimensional input arrays of equal length.",
51    syntax_example = "array_distance(array1, array2)",
52    sql_example = r#"```sql
53> select array_distance([1, 2], [1, 4]);
54+------------------------------------+
55| array_distance(List([1,2], [1,4])) |
56+------------------------------------+
57| 2.0                                |
58+------------------------------------+
59```"#,
60    argument(
61        name = "array1",
62        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
63    ),
64    argument(
65        name = "array2",
66        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
67    )
68)]
69#[derive(Debug, PartialEq, Eq, Hash)]
70pub struct ArrayDistance {
71    signature: Signature,
72    aliases: Vec<String>,
73}
74
75impl Default for ArrayDistance {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl ArrayDistance {
82    pub fn new() -> Self {
83        Self {
84            signature: Signature::user_defined(Volatility::Immutable),
85            aliases: vec!["list_distance".to_string()],
86        }
87    }
88}
89
90impl ScalarUDFImpl for ArrayDistance {
91    fn name(&self) -> &str {
92        "array_distance"
93    }
94
95    fn signature(&self) -> &Signature {
96        &self.signature
97    }
98
99    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
100        Ok(DataType::Float64)
101    }
102
103    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
104        let [_, _] = take_function_args(self.name(), arg_types)?;
105        let coercion = Some(&ListCoercion::FixedSizedListToList);
106        let arg_types = arg_types.iter().map(|arg_type| match arg_type {
107            Null => Ok(coerced_type_with_base_type_only(
108                arg_type,
109                &DataType::Float64,
110                coercion,
111            )),
112            List(field) | LargeList(field) | FixedSizeList(field, _) => {
113                // Distance between nested lists is not supported
114                if matches!(
115                    field.data_type(),
116                    List(_) | LargeList(_) | FixedSizeList(..)
117                ) {
118                    return plan_err!(
119                        "{} only supports one-dimensional arrays, got {arg_type}",
120                        self.name()
121                    );
122                }
123                Ok(coerced_type_with_base_type_only(
124                    arg_type,
125                    &DataType::Float64,
126                    coercion,
127                ))
128            }
129            _ => plan_err!("{} does not support type {arg_type}", self.name()),
130        });
131
132        arg_types.try_collect()
133    }
134
135    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
136        make_scalar_function(array_distance_inner)(&args.args)
137    }
138
139    fn aliases(&self) -> &[String] {
140        &self.aliases
141    }
142
143    fn documentation(&self) -> Option<&Documentation> {
144        self.doc()
145    }
146}
147
148fn array_distance_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
149    let [array1, array2] = take_function_args("array_distance", args)?;
150    match (array1.data_type(), array2.data_type()) {
151        (List(_), List(_)) => general_array_distance::<i32>(args),
152        (LargeList(_), LargeList(_)) => general_array_distance::<i64>(args),
153        (arg_type1, arg_type2) => {
154            exec_err!("array_distance does not support types {arg_type1} and {arg_type2}")
155        }
156    }
157}
158
159fn general_array_distance<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<ArrayRef> {
160    let list_array1 = as_generic_list_array::<O>(&arrays[0])?;
161    let list_array2 = as_generic_list_array::<O>(&arrays[1])?;
162
163    let result = list_array1
164        .iter()
165        .zip(list_array2.iter())
166        .map(|(arr1, arr2)| compute_array_distance(arr1, arr2))
167        .collect::<Result<Float64Array>>()?;
168
169    Ok(Arc::new(result) as ArrayRef)
170}
171
172/// Computes the Euclidean distance between two arrays
173fn compute_array_distance(
174    arr1: Option<ArrayRef>,
175    arr2: Option<ArrayRef>,
176) -> Result<Option<f64>> {
177    let value1 = match arr1 {
178        Some(arr) => arr,
179        None => return Ok(None),
180    };
181    let value2 = match arr2 {
182        Some(arr) => arr,
183        None => return Ok(None),
184    };
185
186    // Check for NULL values inside the arrays
187    if value1.null_count() != 0 || value2.null_count() != 0 {
188        return Ok(None);
189    }
190
191    let values1 = convert_to_f64_array(&value1)?;
192    let values2 = convert_to_f64_array(&value2)?;
193
194    if values1.len() != values2.len() {
195        return exec_err!("Both arrays must have the same length");
196    }
197
198    let sum_squares: f64 = values1
199        .iter()
200        .zip(values2.iter())
201        .map(|(v1, v2)| {
202            let diff = v1.unwrap_or(0.0) - v2.unwrap_or(0.0);
203            diff * diff
204        })
205        .sum();
206
207    Ok(Some(sum_squares.sqrt()))
208}
209
210/// Converts an array of any numeric type to a Float64Array.
211fn convert_to_f64_array(array: &ArrayRef) -> Result<Float64Array> {
212    match array.data_type() {
213        DataType::Float64 => Ok(as_float64_array(array)?.clone()),
214        DataType::Float32 => {
215            let array = as_float32_array(array)?;
216            let converted: Float64Array =
217                array.iter().map(|v| v.map(|v| v as f64)).collect();
218            Ok(converted)
219        }
220        DataType::Int64 => {
221            let array = as_int64_array(array)?;
222            let converted: Float64Array =
223                array.iter().map(|v| v.map(|v| v as f64)).collect();
224            Ok(converted)
225        }
226        DataType::Int32 => {
227            let array = as_int32_array(array)?;
228            let converted: Float64Array =
229                array.iter().map(|v| v.map(|v| v as f64)).collect();
230            Ok(converted)
231        }
232        _ => exec_err!("Unsupported array type for conversion to Float64Array"),
233    }
234}
datafusion_functions_nested/distance.rs

datafusion_functions_nested/
distance.rs
