From eceb73259c8a96235358816c441c6fea1f79344a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 19:45:25 -0600 Subject: [PATCH] perf: optimize levenshtein in datafusion-functions --- .../functions/src/string/levenshtein.rs | 196 ++++++++++++++++-- 1 file changed, 181 insertions(+), 15 deletions(-) diff --git a/datafusion/functions/src/string/levenshtein.rs b/datafusion/functions/src/string/levenshtein.rs index 38fa8fa878de9..859d2cead80f1 100644 --- a/datafusion/functions/src/string/levenshtein.rs +++ b/datafusion/functions/src/string/levenshtein.rs @@ -121,6 +121,113 @@ impl ScalarUDFImpl for LevenshteinFunc { } } +/// Longest pattern, in bytes, that the bit-parallel path can encode in a single +/// machine word. That path is ASCII-only, so this is a character count too. +const MYERS_MAX_PATTERN_LEN: usize = u64::BITS as usize; + +/// Pattern length up to which resetting the match table character by character +/// beats zeroing the whole of it. Above this, a vectorized fill is cheaper than +/// the scattered stores. +const MYERS_SPARSE_RESET_LEN: usize = 32; + +/// Scratch space reused across rows so that no single row has to allocate. +struct LevenshteinScratch { + /// Dynamic-programming row for the character-wise fallback. + cache: Vec, + /// Bitmask of the positions at which each ASCII character occurs in the + /// Myers pattern. Characters absent from the pattern match nowhere, hence 0. + /// Left all-zero between rows so that each row can fill it in directly. + peq: [u64; 128], +} + +impl LevenshteinScratch { + fn new() -> Self { + Self { + cache: Vec::new(), + peq: [0; 128], + } + } +} + +/// Levenshtein distance between `a` and `b`. +/// +/// Uses Myers' bit-parallel algorithm when both inputs are ASCII and the shorter +/// one fits in a single 64-bit word, which computes a whole column of the +/// dynamic-programming matrix per text character instead of one cell at a time. +/// Longer or non-ASCII inputs fall back to the character-wise implementation and +/// stay quadratic. +#[inline] +fn levenshtein_distance(a: &str, b: &str, scratch: &mut LevenshteinScratch) -> usize { + // The distance is symmetric, so the shorter side can always be the pattern. + let (pattern, text) = if a.len() <= b.len() { (a, b) } else { (b, a) }; + + if pattern.len() <= MYERS_MAX_PATTERN_LEN && pattern.is_ascii() && text.is_ascii() { + myers_distance(pattern.as_bytes(), text.as_bytes(), &mut scratch.peq) + } else { + // The fallback sizes its buffer from the second argument, so give it the + // shorter side. + datafusion_strsim::levenshtein_with_buffer(text, pattern, &mut scratch.cache) + } +} + +/// Myers' bit-parallel Levenshtein distance. Both inputs must be ASCII and +/// `pattern` must be at most [`MYERS_MAX_PATTERN_LEN`] bytes long. +/// +/// `vp`/`vn` hold the vertical deltas (+1 / -1) of the current matrix column as +/// bitmasks, one bit per pattern character, and `score` tracks the value of that +/// column's last cell. +/// +/// `peq` must be all-zero on entry and is left all-zero on return. +fn myers_distance(pattern: &[u8], text: &[u8], peq: &mut [u64; 128]) -> usize { + debug_assert!(pattern.is_ascii() && text.is_ascii()); + debug_assert!(pattern.len() <= MYERS_MAX_PATTERN_LEN); + debug_assert!(peq.iter().all(|&mask| mask == 0)); + + let m = pattern.len(); + if m == 0 { + return text.len(); + } + + for (i, &c) in pattern.iter().enumerate() { + peq[c as usize] |= 1 << i; + } + + let last_bit = 1u64 << (m - 1); + let mut vp = u64::MAX; + let mut vn = 0u64; + let mut score = m; + + for &c in text { + let eq = peq[c as usize]; + let x = eq | vn; + let d0 = (vp.wrapping_add(x & vp) ^ vp) | x; + let hn = vp & d0; + let hp = vn | !(vp | d0); + + if hp & last_bit != 0 { + score += 1; + } else if hn & last_bit != 0 { + score -= 1; + } + + let hp = (hp << 1) | 1; + let hn = hn << 1; + vp = hn | !(d0 | hp); + vn = hp & d0; + } + + // Restore the all-zero invariant for the next row. + if m <= MYERS_SPARSE_RESET_LEN { + for &c in pattern { + peq[c as usize] = 0; + } + } else { + peq.fill(0); + } + + score +} + ///Returns the Levenshtein distance between the two given strings. /// LEVENSHTEIN('kitten', 'sitting') = 3 fn levenshtein(args: &[ArrayRef]) -> Result { @@ -147,17 +254,16 @@ fn levenshtein(args: &[ArrayRef]) -> Result { let str1_array = as_string_view_array(&str1)?; let str2_array = as_string_view_array(&str2)?; - // Reusable buffer to avoid allocating for each row - let mut cache = Vec::new(); + // Reusable scratch space to avoid allocating for each row + let mut scratch = LevenshteinScratch::new(); let result = str1_array .iter() .zip(str2_array.iter()) .map(|(string1, string2)| match (string1, string2) { (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein_with_buffer( - string1, string2, &mut cache, - ) as i32) + Some(levenshtein_distance(string1, string2, &mut scratch) + as i32) } _ => None, }) @@ -168,17 +274,16 @@ fn levenshtein(args: &[ArrayRef]) -> Result { let str1_array = as_generic_string_array::(&str1)?; let str2_array = as_generic_string_array::(&str2)?; - // Reusable buffer to avoid allocating for each row - let mut cache = Vec::new(); + // Reusable scratch space to avoid allocating for each row + let mut scratch = LevenshteinScratch::new(); let result = str1_array .iter() .zip(str2_array.iter()) .map(|(string1, string2)| match (string1, string2) { (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein_with_buffer( - string1, string2, &mut cache, - ) as i32) + Some(levenshtein_distance(string1, string2, &mut scratch) + as i32) } _ => None, }) @@ -189,17 +294,16 @@ fn levenshtein(args: &[ArrayRef]) -> Result { let str1_array = as_generic_string_array::(&str1)?; let str2_array = as_generic_string_array::(&str2)?; - // Reusable buffer to avoid allocating for each row - let mut cache = Vec::new(); + // Reusable scratch space to avoid allocating for each row + let mut scratch = LevenshteinScratch::new(); let result = str1_array .iter() .zip(str2_array.iter()) .map(|(string1, string2)| match (string1, string2) { (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein_with_buffer( - string1, string2, &mut cache, - ) as i64) + Some(levenshtein_distance(string1, string2, &mut scratch) + as i64) } _ => None, }) @@ -241,4 +345,66 @@ mod tests { Ok(()) } + + /// The bit-parallel path must agree with the reference implementation for + /// every input it accepts, including empty strings and pattern lengths at + /// the 64-character word boundary. + #[test] + fn myers_matches_reference() { + let alphabets: [&[u8]; 3] = [b"ab", b"abcdefg", b"kitensg -0"]; + let mut state = 0x2545_f491_4f6c_dd1du64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let mut scratch = LevenshteinScratch::new(); + + for alphabet in alphabets { + for len1 in [0usize, 1, 2, 5, 16, 63, 64, 65, 100] { + for len2 in [0usize, 1, 3, 8, 32, 64, 65, 120] { + for _ in 0..8 { + let make = |n: usize, next: &mut dyn FnMut() -> u64| { + (0..n) + .map(|_| { + alphabet[(next() % alphabet.len() as u64) as usize] + as char + }) + .collect::() + }; + let a = make(len1, &mut next); + let b = make(len2, &mut next); + assert_eq!( + levenshtein_distance(&a, &b, &mut scratch), + datafusion_strsim::levenshtein(&a, &b), + "levenshtein({a:?}, {b:?})" + ); + } + } + } + } + } + + #[test] + fn myers_known_distances() { + let mut scratch = LevenshteinScratch::new(); + for (a, b, expected) in [ + ("kitten", "sitting", 3), + ("", "", 0), + ("", "abc", 3), + ("abc", "", 3), + ("abc", "abc", 0), + ("flaw", "lawn", 2), + // Non-ASCII falls back to the character-wise implementation. + ("café", "cafe", 1), + ("😀abc", "abc", 1), + ] { + assert_eq!( + levenshtein_distance(a, b, &mut scratch), + expected, + "{a} {b}" + ); + } + } }