kazantsev-maksim commented on code in PR #5042:
URL: https://github.com/apache/datafusion-comet/pull/5042#discussion_r3883017684
##########
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:
### Criterion Benchmark Results (vs `main` Baseline)
| Batch Size | Data Profile | Time (mean) | Change vs `main` | Status |
| :--- | :--- | :--- | :--- | :--- |
| **8,192** | `no_nulls` | 1.28 ms | **-35.60%** | Improved |
| | `sparse` | 1.20 ms | **-36.68%** | Improved |
| | `all_null` | 24.04 µs | **-21.52%** | Improved |
| **65,536** | `no_nulls` | 10.11 ms | **-39.54%** | Improved |
| | `sparse` | 9.11 ms | **-40.06%** | Improved |
| | `all_null` | 182.18 µs | **-26.20%** | Improved |
| **524,288** | `no_nulls` | 78.79 ms | **-42.06%** | Improved |
| | `sparse` | 72.63 ms | **-42.06%** | Improved |
| | `all_null` | 1.42 ms | **-26.96%** | Improved |
--
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]
