reuse char buffers

Author: n0r0shi

datafusion/common/src/utils/mod.rs

@@ -794,27 +794,42 @@ pub mod datafusion_strsim {
     pub fn levenshtein_with_threshold(a: &str, b: &str, threshold: i32) -> i32 {
         let mut p = Vec::new();
         let mut d = Vec::new();
-        levenshtein_with_threshold_and_buffer(a, b, threshold, &mut p, &mut d)
+        let mut a_buf = Vec::new();
+        let mut b_buf = Vec::new();
+        levenshtein_with_threshold_and_buffer(
+            a,
+            b,
+            threshold,
+            &mut p,
+            &mut d,
+            &mut a_buf,
+            &mut b_buf,
+        )
     }
 
     /// Calculates the Levenshtein distance with a threshold using reusable buffers.
     /// See [`levenshtein_with_threshold`] for details.
     ///
-    /// The `p` and `d` buffers will be resized as needed and reused across calls.
+    /// The `p`, `d`, `a_buf`, and `b_buf` buffers will be resized as needed
+    /// and reused across calls.
     pub fn levenshtein_with_threshold_and_buffer(
         a: &str,
         b: &str,
         threshold: i32,
         p: &mut Vec<i32>,
         d: &mut Vec<i32>,
+        a_buf: &mut Vec<char>,
+        b_buf: &mut Vec<char>,
     ) -> i32 {
-        let a_chars: Vec<char> = a.chars().collect();
-        let b_chars: Vec<char> = b.chars().collect();
+        a_buf.clear();
+        a_buf.extend(a.chars());
+        b_buf.clear();
+        b_buf.extend(b.chars());
 
-        let (s, t, n, m) = if a_chars.len() <= b_chars.len() {
-            (&a_chars, &b_chars, a_chars.len(), b_chars.len())
+        let (s, t, n, m) = if a_buf.len() <= b_buf.len() {
+            (a_buf.as_slice(), b_buf.as_slice(), a_buf.len(), b_buf.len())
         } else {
-            (&b_chars, &a_chars, b_chars.len(), a_chars.len())
+            (b_buf.as_slice(), a_buf.as_slice(), b_buf.len(), a_buf.len())
         };
         // n <= m is guaranteed
 

datafusion/functions/src/string/levenshtein.rs

@@ -187,6 +187,8 @@ fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
     let mut cache = Vec::new();
     let mut p_buf = Vec::new();
     let mut d_buf = Vec::new();
+    let mut a_buf = Vec::new();
+    let mut b_buf = Vec::new();
 
     match coercion_data_type {
         DataType::Utf8View => {
@@ -208,6 +210,8 @@ fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
                         &mut cache,
                         &mut p_buf,
                         &mut d_buf,
+                        &mut a_buf,
+                        &mut b_buf,
                     )
                 })
                 .collect();
@@ -232,6 +236,8 @@ fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
                         &mut cache,
                         &mut p_buf,
                         &mut d_buf,
+                        &mut a_buf,
+                        &mut b_buf,
                     )
                 })
                 .collect();
@@ -256,6 +262,8 @@ fn levenshtein<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
                         &mut cache,
                         &mut p_buf,
                         &mut d_buf,
+                        &mut a_buf,
+                        &mut b_buf,
                     )
                     .map(|d| d as i64)
                 })
@@ -280,11 +288,13 @@ fn compute_distance(
     cache: &mut Vec<usize>,
     p_buf: &mut Vec<i32>,
     d_buf: &mut Vec<i32>,
+    a_buf: &mut Vec<char>,
+    b_buf: &mut Vec<char>,
 ) -> Option<i32> {
     match (s1, s2) {
         (Some(s1), Some(s2)) => match threshold {
             Some(t) => Some(datafusion_strsim::levenshtein_with_threshold_and_buffer(
-                s1, s2, t, p_buf, d_buf,
+                s1, s2, t, p_buf, d_buf, a_buf, b_buf,
             )),
             None => {
                 Some(datafusion_strsim::levenshtein_with_buffer(s1, s2, cache) as i32)