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::{exec_err, Result};
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)]
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!("Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View")
105        }
106    }
107
108    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
109        match args.args[0].data_type() {
110            DataType::Utf8View | DataType::Utf8 => {
111                make_scalar_function(levenshtein::<i32>, vec![])(&args.args)
112            }
113            DataType::LargeUtf8 => {
114                make_scalar_function(levenshtein::<i64>, vec![])(&args.args)
115            }
116            other => {
117                exec_err!("Unsupported data type {other:?} for function levenshtein")
118            }
119        }
120    }
121
122    fn documentation(&self) -> Option<&Documentation> {
123        self.doc()
124    }
125}
126
127///Returns the Levenshtein distance between the two given strings.
128/// LEVENSHTEIN('kitten', 'sitting') = 3
129fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
130    let [str1, str2] = take_function_args("levenshtein", args)?;
131
132    if let Some(coercion_data_type) =
133        string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| {
134            binary_to_string_coercion(args[0].data_type(), args[1].data_type())
135        })
136    {
137        let str1 = if str1.data_type() == &coercion_data_type {
138            Arc::clone(str1)
139        } else {
140            arrow::compute::kernels::cast::cast(&str1, &coercion_data_type)?
141        };
142        let str2 = if str2.data_type() == &coercion_data_type {
143            Arc::clone(str2)
144        } else {
145            arrow::compute::kernels::cast::cast(&str2, &coercion_data_type)?
146        };
147
148        match coercion_data_type {
149            DataType::Utf8View => {
150                let str1_array = as_string_view_array(&str1)?;
151                let str2_array = as_string_view_array(&str2)?;
152                let result = str1_array
153                    .iter()
154                    .zip(str2_array.iter())
155                    .map(|(string1, string2)| match (string1, string2) {
156                        (Some(string1), Some(string2)) => {
157                            Some(datafusion_strsim::levenshtein(string1, string2) as i32)
158                        }
159                        _ => None,
160                    })
161                    .collect::<Int32Array>();
162                Ok(Arc::new(result) as ArrayRef)
163            }
164            DataType::Utf8 => {
165                let str1_array = as_generic_string_array::<T>(&str1)?;
166                let str2_array = as_generic_string_array::<T>(&str2)?;
167                let result = str1_array
168                    .iter()
169                    .zip(str2_array.iter())
170                    .map(|(string1, string2)| match (string1, string2) {
171                        (Some(string1), Some(string2)) => {
172                            Some(datafusion_strsim::levenshtein(string1, string2) as i32)
173                        }
174                        _ => None,
175                    })
176                    .collect::<Int32Array>();
177                Ok(Arc::new(result) as ArrayRef)
178            }
179            DataType::LargeUtf8 => {
180                let str1_array = as_generic_string_array::<T>(&str1)?;
181                let str2_array = as_generic_string_array::<T>(&str2)?;
182                let result = str1_array
183                    .iter()
184                    .zip(str2_array.iter())
185                    .map(|(string1, string2)| match (string1, string2) {
186                        (Some(string1), Some(string2)) => {
187                            Some(datafusion_strsim::levenshtein(string1, string2) as i64)
188                        }
189                        _ => None,
190                    })
191                    .collect::<Int64Array>();
192                Ok(Arc::new(result) as ArrayRef)
193            }
194            other => {
195                exec_err!(
196                    "levenshtein was called with {other} datatype arguments. It requires Utf8View, Utf8 or LargeUtf8."
197                )
198            }
199        }
200    } else {
201        exec_err!("Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View")
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use arrow::array::StringArray;
208
209    use datafusion_common::cast::as_int32_array;
210
211    use super::*;
212
213    #[test]
214    fn to_levenshtein() -> Result<()> {
215        let string1_array =
216            Arc::new(StringArray::from(vec!["123", "abc", "xyz", "kitten"]));
217        let string2_array =
218            Arc::new(StringArray::from(vec!["321", "def", "zyx", "sitting"]));
219        let res = levenshtein::<i32>(&[string1_array, string2_array]).unwrap();
220        let result =
221            as_int32_array(&res).expect("failed to initialized function levenshtein");
222        let expected = Int32Array::from(vec![2, 3, 2, 3]);
223        assert_eq!(&expected, result);
224
225        Ok(())
226    }
227}
datafusion_functions/string/levenshtein.rs

datafusion_functions/string/
levenshtein.rs
