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

datafusion_functions/string/
levenshtein.rs
