sunchao commented on code in PR #5042:
URL: https://github.com/apache/datafusion-comet/pull/5042#discussion_r4104969344
##########
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:
[P2] Preserve throughput for longer string inputs. With 128 non-null rows
containing `"a".repeat(512)` and `"b".repeat(512)`, the two-argument
`spark_levenshtein` returns the correct distance but consistently takes about
66 ms versus 54 ms at base, approximately 23% slower. At 128 characters, the
slowdown is roughly 30–38%. This increases processing time for longer-text
comparisons despite the short-string gains, and the current benchmarks only
exercise very short strings. Please recover the longer-input kernel performance
and add these benchmark shapes. Validating scratch-row lengths before the inner
loop removed most of the slowdown in a disposable variant, providing a concrete
direction to investigate bounds-check elimination.
Evidence: Compiled unchanged base/head modules side by side with rustc
1.98.1, `-O -C overflow-checks=no -C codegen-units=1`, linking the checkout’s
Arrow 58.4.0/DataFusion 54.1.0 debug dependency artifacts. Constructed two
128-row StringArrays outside the timer, invoked `spark_levenshtein` without a
threshold, and consumed results with `black_box`. Four alternating-order
samples in each of two separate CPU-pinned processes measured 512-character
ASCII batches at 53.619–54.068 ms for base and 65.939–68.380 ms for head.
Independent extracted-kernel measurements with thin LTO also reproduced the
regression. Reproduction: `/tmp/comet-5042-9e28012a-review/bench_api.rs`;
results: `bench-api-1.log` and `bench-api-2.log` in that directory. Both module
copies were byte-verified against the requested Git objects.
--
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]
