datafusion_functions/string/
levenshtein.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
18use std::sync::Arc;
19
20use arrow::array::{ArrayRef, Int32Array, Int64Array, OffsetSizeTrait};
21use arrow::datatypes::DataType;
22
23use crate::utils::{make_scalar_function, utf8_to_int_type};
24use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
25use datafusion_common::types::logical_string;
26use datafusion_common::utils::datafusion_strsim;
27use datafusion_common::utils::take_function_args;
28use datafusion_common::{Result, exec_err};
29use datafusion_expr::type_coercion::binary::{
30    binary_to_string_coercion, string_coercion,
31};
32use datafusion_expr::{
33    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
34    TypeSignatureClass, Volatility,
35};
36use datafusion_macros::user_doc;
37
38#[user_doc(
39    doc_section(label = "String Functions"),
40    description = "Returns the [`Levenshtein distance`](https://en.wikipedia.org/wiki/Levenshtein_distance) between the two given strings.",
41    syntax_example = "levenshtein(str1, str2)",
42    sql_example = r#"```sql
43> select levenshtein('kitten', 'sitting');
44+---------------------------------------------+
45| levenshtein(Utf8("kitten"),Utf8("sitting")) |
46+---------------------------------------------+
47| 3                                           |
48+---------------------------------------------+
49```"#,
50    argument(
51        name = "str1",
52        description = "String expression to compute Levenshtein distance with str2."
53    ),
54    argument(
55        name = "str2",
56        description = "String expression to compute Levenshtein distance with str1."
57    )
58)]
59#[derive(Debug, PartialEq, Eq, Hash)]
60pub struct LevenshteinFunc {
61    signature: Signature,
62}
63
64impl Default for LevenshteinFunc {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl LevenshteinFunc {
71    pub fn new() -> Self {
72        Self {
73            signature: Signature::coercible(
74                vec![
75                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
76                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
77                ],
78                Volatility::Immutable,
79            ),
80        }
81    }
82}
83
84impl ScalarUDFImpl for LevenshteinFunc {
85    fn name(&self) -> &str {
86        "levenshtein"
87    }
88
89    fn signature(&self) -> &Signature {
90        &self.signature
91    }
92
93    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
94        if let Some(coercion_data_type) = string_coercion(&arg_types[0], &arg_types[1])
95            .or_else(|| binary_to_string_coercion(&arg_types[0], &arg_types[1]))
96        {
97            utf8_to_int_type(&coercion_data_type, "levenshtein")
98        } else {
99            exec_err!(
100                "Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View"
101            )
102        }
103    }
104
105    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
106        match args.args[0].data_type() {
107            DataType::Utf8View | DataType::Utf8 => {
108                make_scalar_function(levenshtein::<i32>, vec![])(&args.args)
109            }
110            DataType::LargeUtf8 => {
111                make_scalar_function(levenshtein::<i64>, vec![])(&args.args)
112            }
113            other => {
114                exec_err!("Unsupported data type {other:?} for function levenshtein")
115            }
116        }
117    }
118
119    fn documentation(&self) -> Option<&Documentation> {
120        self.doc()
121    }
122}
123
124///Returns the Levenshtein distance between the two given strings.
125/// LEVENSHTEIN('kitten', 'sitting') = 3
126fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
127    let [str1, str2] = take_function_args("levenshtein", args)?;
128
129    if let Some(coercion_data_type) =
130        string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| {
131            binary_to_string_coercion(args[0].data_type(), args[1].data_type())
132        })
133    {
134        let str1 = if str1.data_type() == &coercion_data_type {
135            Arc::clone(str1)
136        } else {
137            arrow::compute::kernels::cast::cast(&str1, &coercion_data_type)?
138        };
139        let str2 = if str2.data_type() == &coercion_data_type {
140            Arc::clone(str2)
141        } else {
142            arrow::compute::kernels::cast::cast(&str2, &coercion_data_type)?
143        };
144
145        match coercion_data_type {
146            DataType::Utf8View => {
147                let str1_array = as_string_view_array(&str1)?;
148                let str2_array = as_string_view_array(&str2)?;
149
150                // Reusable buffer to avoid allocating for each row
151                let mut cache = Vec::new();
152
153                let result = str1_array
154                    .iter()
155                    .zip(str2_array.iter())
156                    .map(|(string1, string2)| match (string1, string2) {
157                        (Some(string1), Some(string2)) => {
158                            Some(datafusion_strsim::levenshtein_with_buffer(
159                                string1, string2, &mut cache,
160                            ) as i32)
161                        }
162                        _ => None,
163                    })
164                    .collect::<Int32Array>();
165                Ok(Arc::new(result) as ArrayRef)
166            }
167            DataType::Utf8 => {
168                let str1_array = as_generic_string_array::<T>(&str1)?;
169                let str2_array = as_generic_string_array::<T>(&str2)?;
170
171                // Reusable buffer to avoid allocating for each row
172                let mut cache = Vec::new();
173
174                let result = str1_array
175                    .iter()
176                    .zip(str2_array.iter())
177                    .map(|(string1, string2)| match (string1, string2) {
178                        (Some(string1), Some(string2)) => {
179                            Some(datafusion_strsim::levenshtein_with_buffer(
180                                string1, string2, &mut cache,
181                            ) as i32)
182                        }
183                        _ => None,
184                    })
185                    .collect::<Int32Array>();
186                Ok(Arc::new(result) as ArrayRef)
187            }
188            DataType::LargeUtf8 => {
189                let str1_array = as_generic_string_array::<T>(&str1)?;
190                let str2_array = as_generic_string_array::<T>(&str2)?;
191
192                // Reusable buffer to avoid allocating for each row
193                let mut cache = Vec::new();
194
195                let result = str1_array
196                    .iter()
197                    .zip(str2_array.iter())
198                    .map(|(string1, string2)| match (string1, string2) {
199                        (Some(string1), Some(string2)) => {
200                            Some(datafusion_strsim::levenshtein_with_buffer(
201                                string1, string2, &mut cache,
202                            ) as i64)
203                        }
204                        _ => None,
205                    })
206                    .collect::<Int64Array>();
207                Ok(Arc::new(result) as ArrayRef)
208            }
209            other => {
210                exec_err!(
211                    "levenshtein was called with {other} datatype arguments. It requires Utf8View, Utf8 or LargeUtf8."
212                )
213            }
214        }
215    } else {
216        exec_err!(
217            "Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View"
218        )
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use arrow::array::StringArray;
225
226    use datafusion_common::cast::as_int32_array;
227
228    use super::*;
229
230    #[test]
231    fn to_levenshtein() -> Result<()> {
232        let string1_array =
233            Arc::new(StringArray::from(vec!["123", "abc", "xyz", "kitten"]));
234        let string2_array =
235            Arc::new(StringArray::from(vec!["321", "def", "zyx", "sitting"]));
236        let res = levenshtein::<i32>(&[string1_array, string2_array]).unwrap();
237        let result =
238            as_int32_array(&res).expect("failed to initialized function levenshtein");
239        let expected = Int32Array::from(vec![2, 3, 2, 3]);
240        assert_eq!(&expected, result);
241
242        Ok(())
243    }
244}
datafusion_functions/string/levenshtein.rs

datafusion_functions/string/
levenshtein.rs
