sunchao commented on code in PR #5042:
URL: https://github.com/apache/datafusion-comet/pull/5042#discussion_r3876125094
##########
native/spark-expr/src/string_funcs/levenshtein.rs:
##########
@@ -26,67 +26,195 @@ use datafusion::common::{cast::as_generic_string_array,
DataFusionError, Result}
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
+const MAX_RETAINED_CAPACITY: usize = 1024;
+
+// Thread-local scratch buffers to avoid heap allocations in the row
processing loop
+thread_local! {
+ static LEVENSHTEIN_SCRATCH: std::cell::RefCell<(Vec<i32>, Vec<i32>)> =
+ std::cell::RefCell::new((Vec::with_capacity(64),
Vec::with_capacity(64)));
+}
+
+#[inline]
+fn prepare_scratch(buf: &mut Vec<i32>, len: usize, default_val: i32) {
+ if buf.capacity() > MAX_RETAINED_CAPACITY && len <= MAX_RETAINED_CAPACITY {
+ *buf = Vec::with_capacity(MAX_RETAINED_CAPACITY);
+ }
+ buf.clear();
+ buf.resize(len, default_val);
+}
+
/// Computes the Levenshtein edit distance between two UTF-8 strings.
-///
-/// This uses the standard dynamic programming algorithm with O(min(m,n))
space.
fn levenshtein_distance(s: &str, t: &str) -> i32 {
+ // Fast path for ASCII strings: operate directly on raw bytes without
vector allocations
+ if s.is_ascii() && t.is_ascii() {
+ let s_bytes = s.as_bytes();
+ let t_bytes = t.as_bytes();
+ let m = s_bytes.len();
+ let n = t_bytes.len();
+
+ if m == 0 {
+ return n as i32;
+ }
+ if n == 0 {
+ return m as i32;
+ }
+
+ let (s_bytes, t_bytes, m, n) = if m > n {
+ (t_bytes, s_bytes, n, m)
+ } else {
+ (s_bytes, t_bytes, m, n)
+ };
+
+ return LEVENSHTEIN_SCRATCH.with(|scratch| {
+ let mut borrow = scratch.borrow_mut();
+ let (prev, curr) = &mut *borrow;
+
+ prepare_scratch(prev, m + 1, 0);
+ prepare_scratch(curr, m + 1, 0);
+
+ for (i, val) in prev.iter_mut().enumerate() {
+ *val = i as i32;
+ }
+
+ for (j, &t_byte) in t_bytes.iter().enumerate().take(n) {
+ curr[0] = (j + 1) as i32;
+ for i in 1..=m {
+ let cost = if s_bytes[i - 1] == t_byte { 0 } else { 1 };
+ curr[i] = (prev[i] + 1).min(curr[i - 1] + 1).min(prev[i -
1] + cost);
+ }
+ std::mem::swap(prev, curr);
+ }
+
+ prev[m]
+ });
+ }
+
+ // General Unicode path for non-ASCII strings
let s_chars: Vec<char> = s.chars().collect();
let t_chars: Vec<char> = t.chars().collect();
let m = s_chars.len();
let n = t_chars.len();
- // Optimization: if one string is empty, distance is the length of the
other
if m == 0 {
return n as i32;
}
if n == 0 {
return m as i32;
}
- // Use the shorter string for the "column" to minimize space usage
let (s_chars, t_chars, m, n) = if m > n {
(t_chars, s_chars, n, m)
} else {
(s_chars, t_chars, m, n)
};
- // Previous and current row of distances
- let mut prev = vec![0i32; m + 1];
- let mut curr = vec![0i32; m + 1];
+ LEVENSHTEIN_SCRATCH.with(|scratch| {
+ let mut borrow = scratch.borrow_mut();
+ let (prev, curr) = &mut *borrow;
- // Initialize base case: distance from empty string
- for (i, val) in prev.iter_mut().enumerate() {
- *val = i as i32;
- }
+ prev.resize(m + 1, 0);
+ curr.resize(m + 1, 0);
- for j in 1..=n {
- curr[0] = j as i32;
- for i in 1..=m {
- let cost = if s_chars[i - 1] == t_chars[j - 1] {
- 0
- } else {
- 1
- };
- curr[i] = (prev[i] + 1) // deletion
- .min(curr[i - 1] + 1) // insertion
- .min(prev[i - 1] + cost); // substitution
+ for (i, val) in prev.iter_mut().enumerate() {
+ *val = i as i32;
}
- std::mem::swap(&mut prev, &mut curr);
- }
- prev[m]
+ for j in 1..=n {
+ curr[0] = j as i32;
+ for i in 1..=m {
+ let cost = if s_chars[i - 1] == t_chars[j - 1] {
+ 0
+ } else {
+ 1
+ };
+ curr[i] = (prev[i] + 1).min(curr[i - 1] + 1).min(prev[i - 1] +
cost);
+ }
+ std::mem::swap(prev, curr);
+ }
+
+ prev[m]
+ })
}
/// Computes the Levenshtein distance up to `threshold` using a diagonal band.
-///
-/// Spark's three-argument form uses the threshold to avoid evaluating cells
that cannot
-/// contribute to a result within the requested distance. This keeps the
complexity at
-/// O(threshold * max(m, n)) when the threshold is small rather than always
using O(m * n).
fn levenshtein_distance_with_threshold(s: &str, t: &str, threshold: i32) ->
i32 {
if threshold < 0 {
return -1;
}
+ // Fast path for ASCII strings
+ if s.is_ascii() && t.is_ascii() {
+ let s_bytes = s.as_bytes();
+ let t_bytes = t.as_bytes();
+ let (shorter, longer) = if s_bytes.len() <= t_bytes.len() {
+ (s_bytes, t_bytes)
+ } else {
+ (t_bytes, s_bytes)
+ };
+ let m = shorter.len();
+ let n = longer.len();
+ let threshold = threshold as usize;
+
+ if n - m > threshold {
+ return -1;
+ }
+ if m == 0 {
+ return if n <= threshold { n as i32 } else { -1 };
+ }
+
+ let out_of_band = n.saturating_add(1) as i32;
+
+ return LEVENSHTEIN_SCRATCH.with(|scratch| {
+ let mut borrow = scratch.borrow_mut();
+ let (prev, curr) = &mut *borrow;
+
+ prepare_scratch(prev, m + 1, out_of_band);
+ prepare_scratch(curr, m + 1, out_of_band);
+
+ for (i, value) in
prev.iter_mut().enumerate().take(m.min(threshold) + 1) {
+ *value = i as i32;
+ }
+
+ for (j_idx, &t_char) in longer.iter().enumerate() {
+ let j = j_idx + 1;
+ let start = j.saturating_sub(threshold).max(1);
+ let end = (j + threshold).min(m);
+
+ if start > end {
+ return -1;
+ }
+
+ if start > 1 {
+ curr[start - 1] = out_of_band;
+ }
+
+ if j <= threshold {
+ curr[0] = j as i32;
+ }
Review Comment:
[P1] Reset `curr[0]` when it leaves the threshold band
When `j == threshold + 1`, `start == 1`, so neither branch resets `curr[0]`;
its old value admits an invalid insertion path. On `8456bc5`,
`spark_levenshtein` returns `1` for (`"a"`, `"bb"`, `1`) instead of `-1`
(distance `2`), causing false matches. This reproduces through the public Rust
API for scalar, array, and mixed Utf8/LargeUtf8 inputs; the base returns `-1`.
Please set `curr[0]` to `out_of_band` when `j > threshold` and add this
regression case. Validation is Rust component/API execution, not an end-to-end
Spark run.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
