kazantsev-maksim commented on code in PR #5042:
URL: https://github.com/apache/datafusion-comet/pull/5042#discussion_r4116142778
##########
native/spark-expr/src/string_funcs/levenshtein.rs:
##########
@@ -26,10 +26,89 @@ use datafusion::common::{cast::as_generic_string_array,
DataFusionError, Result}
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
+/// Maximum retained scratch buffer capacity (1024 elements * 4 bytes = 4 KB).
+/// Inputs requiring larger buffers bypass TLS to avoid unbounded memory
retention.
+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)));
+}
+
+/// Executes a closure using scratch buffers.
+///
+/// For sizes up to `MAX_RETAINED_CAPACITY`, reuses TLS buffers (bounded to
+/// at most `2 * MAX_RETAINED_CAPACITY * 4` bytes per worker thread).
+/// For oversized rows, allocates temporary vectors in the call scope so the
+/// TLS buffers never grow beyond the cap.
+#[inline]
+fn with_scratch_buffers<F, R>(len: usize, default_val: i32, f: F) -> R
+where
+ F: FnOnce(&mut Vec<i32>, &mut Vec<i32>) -> R,
+{
+ if len > MAX_RETAINED_CAPACITY {
+ let mut prev = vec![default_val; len];
+ let mut curr = vec![default_val; len];
+ f(&mut prev, &mut curr)
+ } else {
+ LEVENSHTEIN_SCRATCH.with(|scratch| {
+ let mut borrow = scratch.borrow_mut();
+ let (prev, curr) = &mut *borrow;
+
+ prev.clear();
+ prev.resize(len, default_val);
+ curr.clear();
+ curr.resize(len, default_val);
+
+ f(prev, curr)
+ })
+ }
+}
+
/// 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 with_scratch_buffers(m + 1, 0, |prev, curr| {
+ 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);
Review Comment:
Thanks for the detailed finding and the pointer on bounds-check elimination!
I have updated the PR to address the `[P2]` longer-input performance
regression:
1. **Eliminated hot-loop bounds checks:** Added explicit bounds assertions
(`assert!(prev.len() > m && curr.len() > m)` and `assert!(s_bytes.len() >= m)`)
directly before the inner loop, enabling LLVM to prove slice boundaries and
eliminate dynamic panic checks inside the hot DP loop.
2. **Register caching for DP row:** Kept `left` (`curr[i-1]`) and `diag`
(`prev[i-1]`) in registers across inner iterations, reducing memory traffic to
a single read and single write per cell.
3. **Extended benchmarks:** Added `ascii-128` and `ascii-512` dataset shapes
with 128 rows to `benches/levenshtein.rs` to permanently cover longer inputs
and prevent future regressions.
### Benchmark Results (Criterion)
### Long ASCII — 128 chars ("a" * 128 / "b" * 128)
| Batch | Profile | Median time | Change vs prev head |
| --- | --- | --- | --- |
| 128 | no_nulls | 3.267 ms | -48.8% |
| 128 | sparse | 2.946 ms | -49.3% |
| 128 | all_null | 584.15 ns | -5.5% |
### Long ASCII — 512 chars ("a" * 512 / "b" * 512, P2 verification)
| Batch | Profile | Median time | Change vs prev head |
| --- | --- | --- | --- |
| 128 | no_nulls | 52.47 ms | -51.4% *(base: ~54 ms)* |
| 128 | sparse | 47.15 ms | -51.1% |
| 128 | all_null | 584.71 ns | -4.4% |
### ASCII (kitten / sitting)
| Batch | Profile | Median time | Change vs prev head |
| --- | --- | --- | --- |
| 8192 | no_nulls | 665.19 µs | -68.9% |
| 8192 | sparse | 620.54 µs | -68.7% |
| 8192 | all_null | 23.98 µs | -22.2% |
| 65536 | no_nulls | 5.262 ms | -69.8% |
| 65536 | sparse | 4.856 ms | -68.9% |
| 65536 | all_null | 180.04 µs | -26.9% |
| 524288 | no_nulls | 42.80 ms | -68.8% |
| 524288 | sparse | 39.17 ms | -69.2% |
| 524288 | all_null | 1.444 ms | -26.7% |
### Non-ASCII (naïve / smörgås)
| Batch | Profile | Median time | Change vs prev head |
| --- | --- | --- | --- |
| 8192 | no_nulls | 1.374 ms | -28.8% |
| 8192 | sparse | 1.264 ms | -28.5% |
| 8192 | all_null | 23.72 µs | -21.5% |
| 65536 | no_nulls | 10.98 ms | -28.3% |
| 65536 | sparse | 10.08 ms | -28.3% |
| 65536 | all_null | 176.43 µs | -27.8% |
| 524288 | no_nulls | 87.91 ms | -27.8% |
| 524288 | sparse | 80.66 ms | -27.8% |
| 524288 | all_null | 1.454 ms | -24.1% |
### Mixed (cafe / café)
| Batch | Profile | Median time | Change vs prev head |
| --- | --- | --- |
| 8192 | no_nulls | 597.78 µs | -34.5% |
| 8192 | sparse | 563.21 µs | -34.1% |
| 8192 | all_null | 23.88 µs | -18.4% |
| 65536 | no_nulls | 4.806 ms | -34.6% |
| 65536 | sparse | 4.497 ms | -34.0% |
| 65536 | all_null | 179.48 µs | -26.8% |
| 524288 | no_nulls | 37.94 ms | -35.6% |
| 524288 | sparse | 35.57 ms | -35.2% |
| 524288 | all_null | 1.414 ms | -27.9% |
--
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]
