This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 756ddd7e14 feat: Support Spark levenshtein expression in native
execution (#4105)
756ddd7e14 is described below
commit 756ddd7e1457bf97d784c3f04810fcee766140a4
Author: Myx778 <[email protected]>
AuthorDate: Sat Jul 25 01:30:55 2026 +0800
feat: Support Spark levenshtein expression in native execution (#4105)
* feat: Support Spark levenshtein expression in native execution
Implements the Levenshtein edit distance function as a native Comet
scalar UDF, enabling Spark's `levenshtein(str1, str2)` to run via
DataFusion instead of falling back to JVM.
Implementation:
- Rust kernel: O(min(m,n)) space DP algorithm with Unicode char-level
distance computation and proper NULL propagation
- Scala serde: Register via CometScalarFunction("levenshtein") in
QueryPlanSerde stringExpressions map
- Tests: Basic, NULL handling, and Unicode test cases
Closes #3084
* style: fix Spotless formatting violations in test code
- Reformat Seq literal to match scalafmt maxColumn=98 constraint
- Break long sql() call into multi-line format
- Reformat unicode test data as multi-line Seq
* fix: address CI lint failures
- Fix Clippy needless_range_loop: use iter_mut().enumerate() instead
of indexing loop for initializing prev array
- Fix Scalafix unused import: remove Levenshtein from strings.scala
imports (it's used via wildcard import in QueryPlanSerde.scala)
* test: add SLT tests for levenshtein expression
Add SQL logic test file covering:
- Column arguments with NULL propagation
- Column + literal combinations
- Literal + literal edge cases
- Identical string comparison
- Unicode character-level distance
* Add levenshtein to benchmark and update expressions.md
- Add levenshtein benchmark entry in CometStringExpressionBenchmark
- Update expressions.md to list Levenshtein as supported
* Support three-argument levenshtein with threshold and add collation check
- Implement threshold semantics: levenshtein(str1, str2, threshold) returns
-1 when distance exceeds threshold, matching Spark behavior
- Add CometLevenshtein serde with getSupportLevel that falls back for
non-default collations (Spark 4 StringTypeWithCollation)
- Add Rust tests for threshold and NULL threshold cases
- Add SLT tests for threshold variants
- Add Scala integration test for threshold
- Add levenshtein_threshold to benchmark
* Fix Spotless formatting in levenshtein threshold test
* Fix Spotless formatting in strings.scala and test suite
* Use scalarFunctionExprToProtoWithReturnType for levenshtein
Override convert() to use scalarFunctionExprToProtoWithReturnType
with IntegerType, so the native planner skips the DataFusion
registry lookup and does not conflict with DataFusion's built-in
2-arg levenshtein function when 3 args (threshold) are passed.
Also fix Spotless: remove unnecessary string interpolation prefix.
* Fix Spark 3.4 compatibility: skip levenshtein threshold tests on Spark <
3.5
The three-argument levenshtein(str1, str2, threshold) is only available
in Spark 3.5+. This moves threshold-related SQL tests to a separate file
with MinSparkVersion: 3.5 and adds assume(isSpark35Plus) to the unit test.
* Remove duplicate levenshtein benchmark entry
* Support threshold as column (not just literal) in levenshtein
- Refactored spark_levenshtein to handle threshold as either
ColumnarValue::Scalar or ColumnarValue::Array using into_array()
- NULL threshold in a column produces NULL result for that row
- Added Rust unit tests for array threshold, nulls, and negative values
- Added Scala integration tests with threshold as column reference
- Added SQL file tests with threshold column and NULL scenarios
* Fix unused import warning: move ScalarValue to test module
* Fix Spotless formatting in threshold column tests
* Run cargo fmt
---------
Co-authored-by: yusinnmao <[email protected]>
Co-authored-by: Cursor <[email protected]>
---
native/spark-expr/src/comet_scalar_funcs.rs | 4 +
native/spark-expr/src/string_funcs/levenshtein.rs | 485 +++++++++++++++++++++
native/spark-expr/src/string_funcs/mod.rs | 2 +
.../scala/org/apache/comet/serde/strings.scala | 27 +-
.../sql-tests/expressions/string/levenshtein.sql | 30 +-
.../expressions/string/levenshtein_threshold.sql | 60 +++
.../scala/org/apache/comet/CometCodegenSuite.scala | 14 +-
.../benchmark/CometStringExpressionBenchmark.scala | 4 +
8 files changed, 609 insertions(+), 17 deletions(-)
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs
b/native/spark-expr/src/comet_scalar_funcs.rs
index 98d78e0e67..8e913dabab 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -233,6 +233,10 @@ pub fn create_comet_physical_fun_with_eval_mode(
"next_day" => Ok(Arc::new(ScalarUDF::new_from_impl(SparkNextDay::new(
fail_on_error,
)))),
+ "levenshtein" => {
+ let func = Arc::new(crate::string_funcs::spark_levenshtein);
+ make_comet_scalar_udf!("levenshtein", func, without data_type)
+ }
_ => registry.udf(fun_name).map_err(|e| {
DataFusionError::Execution(format!(
"Function {fun_name} not found in the registry: {e}",
diff --git a/native/spark-expr/src/string_funcs/levenshtein.rs
b/native/spark-expr/src/string_funcs/levenshtein.rs
new file mode 100644
index 0000000000..19a81f92ff
--- /dev/null
+++ b/native/spark-expr/src/string_funcs/levenshtein.rs
@@ -0,0 +1,485 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Levenshtein distance expression implementation.
+//!
+//! Computes the Levenshtein edit distance between two strings,
+//! matching Apache Spark's `levenshtein(str1, str2)` semantics.
+
+use arrow::array::{Array, ArrayRef, GenericStringArray, Int32Array,
OffsetSizeTrait};
+use arrow::datatypes::DataType;
+use datafusion::common::{cast::as_generic_string_array, DataFusionError,
Result};
+use datafusion::physical_plan::ColumnarValue;
+use std::sync::Arc;
+
+/// 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 {
+ 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];
+
+ // Initialize base case: distance from empty string
+ for (i, val) in prev.iter_mut().enumerate() {
+ *val = i as i32;
+ }
+
+ 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
+ }
+ std::mem::swap(&mut prev, &mut 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;
+ }
+
+ let s_chars: Vec<char> = s.chars().collect();
+ let t_chars: Vec<char> = t.chars().collect();
+ let (shorter, longer) = if s_chars.len() <= t_chars.len() {
+ (s_chars, t_chars)
+ } else {
+ (t_chars, s_chars)
+ };
+ 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);
+ let mut prev = vec![out_of_band; m + 1];
+ let mut curr = vec![out_of_band; m + 1];
+ for (i, value) in prev.iter_mut().enumerate().take(m.min(threshold) + 1) {
+ *value = i;
+ }
+
+ for j in 1..=n {
+ let start = 1.max(j.saturating_sub(threshold));
+ let end = m.min(j.saturating_add(threshold));
+ if start > end {
+ return -1;
+ }
+
+ curr[0] = if j <= threshold { j } else { out_of_band };
+ curr[start - 1] = if start == 1 { curr[0] } else { out_of_band };
+ for i in start..=end {
+ let cost = usize::from(shorter[i - 1] != longer[j - 1]);
+ curr[i] = prev[i]
+ .saturating_add(1)
+ .min(curr[i - 1].saturating_add(1))
+ .min(prev[i - 1].saturating_add(cost));
+ }
+ if end < m {
+ curr[end + 1] = out_of_band;
+ }
+ std::mem::swap(&mut prev, &mut curr);
+ }
+
+ if prev[m] <= threshold {
+ prev[m] as i32
+ } else {
+ -1
+ }
+}
+
+fn evaluate_levenshtein<LeftOffset, RightOffset>(
+ left: &GenericStringArray<LeftOffset>,
+ right: &GenericStringArray<RightOffset>,
+ threshold: Option<&Int32Array>,
+) -> Int32Array
+where
+ LeftOffset: OffsetSizeTrait,
+ RightOffset: OffsetSizeTrait,
+{
+ left.iter()
+ .zip(right.iter())
+ .enumerate()
+ .map(|(i, (left_value, right_value))| {
+ if threshold.is_some_and(|values| values.is_null(i)) {
+ return None;
+ }
+
+ match (left_value, right_value) {
+ (Some(left_value), Some(right_value)) => Some(match threshold {
+ Some(values) => levenshtein_distance_with_threshold(
+ left_value,
+ right_value,
+ values.value(i),
+ ),
+ None => levenshtein_distance(left_value, right_value),
+ }),
+ _ => None,
+ }
+ })
+ .collect()
+}
+
+fn evaluate_string_arrays(
+ left: &ArrayRef,
+ right: &ArrayRef,
+ threshold: Option<&Int32Array>,
+) -> Result<Int32Array> {
+ match (left.data_type(), right.data_type()) {
+ (DataType::Utf8, DataType::Utf8) => Ok(evaluate_levenshtein(
+ as_generic_string_array::<i32>(left.as_ref())?,
+ as_generic_string_array::<i32>(right.as_ref())?,
+ threshold,
+ )),
+ (DataType::Utf8, DataType::LargeUtf8) => Ok(evaluate_levenshtein(
+ as_generic_string_array::<i32>(left.as_ref())?,
+ as_generic_string_array::<i64>(right.as_ref())?,
+ threshold,
+ )),
+ (DataType::LargeUtf8, DataType::Utf8) => Ok(evaluate_levenshtein(
+ as_generic_string_array::<i64>(left.as_ref())?,
+ as_generic_string_array::<i32>(right.as_ref())?,
+ threshold,
+ )),
+ (DataType::LargeUtf8, DataType::LargeUtf8) => Ok(evaluate_levenshtein(
+ as_generic_string_array::<i64>(left.as_ref())?,
+ as_generic_string_array::<i64>(right.as_ref())?,
+ threshold,
+ )),
+ (left_type, right_type) => Err(DataFusionError::Execution(format!(
+ "levenshtein expects Utf8 or LargeUtf8 arguments, got
{left_type:?} and {right_type:?}"
+ ))),
+ }
+}
+
+/// Spark-compatible levenshtein scalar function.
+///
+/// Accepts two or three arguments:
+/// - `levenshtein(str1, str2)` → edit distance
+/// - `levenshtein(str1, str2, threshold)` → edit distance if <= threshold,
else -1
+///
+/// The threshold argument can be either a scalar or a column (array).
+/// NULL inputs produce NULL outputs. NULL threshold produces NULL output for
that row.
+pub fn spark_levenshtein(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() < 2 || args.len() > 3 {
+ return Err(DataFusionError::Internal(format!(
+ "levenshtein requires 2 or 3 arguments, got {}",
+ args.len()
+ )));
+ }
+
+ // Determine array length from any array argument
+ let len = args
+ .iter()
+ .find_map(|arg| match arg {
+ ColumnarValue::Array(a) => Some(a.len()),
+ _ => None,
+ })
+ .unwrap_or(1);
+
+ let left = args[0].clone().into_array(len)?;
+ let right = args[1].clone().into_array(len)?;
+ if left.len() != len || right.len() != len {
+ return Err(DataFusionError::Internal(
+ "levenshtein arguments must have the same length".to_string(),
+ ));
+ }
+
+ // Handle the optional threshold argument (scalar or array)
+ let threshold_array = if args.len() == 3 {
+ let threshold_array = args[2].clone().into_array(len)?;
+ if threshold_array.len() != len {
+ return Err(DataFusionError::Internal(
+ "levenshtein threshold must have the same length as string
arguments".to_string(),
+ ));
+ }
+ Some(threshold_array)
+ } else {
+ None
+ };
+ let threshold = threshold_array
+ .as_ref()
+ .map(|array| {
+ array.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
+ DataFusionError::Internal("levenshtein threshold must be
Int32".to_string())
+ })
+ })
+ .transpose()?;
+
+ let result = evaluate_string_arrays(&left, &right, threshold)?;
+ Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{LargeStringArray, StringArray};
+ use datafusion::common::ScalarValue;
+
+ #[test]
+ fn test_levenshtein_basic() {
+ assert_eq!(levenshtein_distance("", ""), 0);
+ assert_eq!(levenshtein_distance("abc", ""), 3);
+ assert_eq!(levenshtein_distance("", "abc"), 3);
+ assert_eq!(levenshtein_distance("abc", "abc"), 0);
+ assert_eq!(levenshtein_distance("kitten", "sitting"), 3);
+ assert_eq!(levenshtein_distance("frog", "fog"), 1);
+ }
+
+ #[test]
+ fn test_levenshtein_unicode() {
+ // Spark counts character-level (not byte-level) edit distance
+ assert_eq!(levenshtein_distance("你好", "你坏"), 1);
+ assert_eq!(levenshtein_distance("abc", "äbc"), 1);
+ }
+
+ #[test]
+ fn test_spark_levenshtein_nulls() {
+ let left = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("abc"),
+ None,
+ Some("hello"),
+ ])));
+ let right = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("adc"),
+ Some("test"),
+ None,
+ ])));
+
+ let result = spark_levenshtein(&[left, right]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.value(0), 1); // abc -> adc = 1
+ assert!(int_arr.is_null(1)); // NULL -> test = NULL
+ assert!(int_arr.is_null(2)); // hello -> NULL = NULL
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_large_utf8() {
+ let left =
+
ColumnarValue::Array(Arc::new(LargeStringArray::from(vec![Some("kitten"),
None])));
+ let right = ColumnarValue::Array(Arc::new(LargeStringArray::from(vec![
+ Some("sitting"),
+ None,
+ ])));
+
+ let result = spark_levenshtein(&[left, right]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.value(0), 3);
+ assert!(int_arr.is_null(1));
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_with_threshold() {
+ let left = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("kitten"),
+ Some("abc"),
+ Some("frog"),
+ ])));
+ let right = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("sitting"),
+ Some("adc"),
+ Some("fog"),
+ ])));
+ let threshold = ColumnarValue::Scalar(ScalarValue::Int32(Some(2)));
+
+ let result = spark_levenshtein(&[left, right, threshold]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.value(0), -1); // kitten->sitting=3 > 2,
return -1
+ assert_eq!(int_arr.value(1), 1); // abc->adc=1 <= 2, return 1
+ assert_eq!(int_arr.value(2), 1); // frog->fog=1 <= 2, return 1
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_null_threshold() {
+ let left =
ColumnarValue::Array(Arc::new(StringArray::from(vec![Some("abc")])));
+ let right =
ColumnarValue::Array(Arc::new(StringArray::from(vec![Some("adc")])));
+ let threshold = ColumnarValue::Scalar(ScalarValue::Int32(None));
+
+ let result = spark_levenshtein(&[left, right, threshold]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.len(), 1);
+ assert!(int_arr.is_null(0)); // NULL threshold -> NULL result
+ }
+ _ => panic!("Expected array result with NULL for NULL threshold"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_threshold_as_array() {
+ // threshold is a column (array) with per-row values
+ let left = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("kitten"),
+ Some("frog"),
+ Some("abc"),
+ Some("hello"),
+ ])));
+ let right = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("sitting"),
+ Some("fog"),
+ Some("abc"),
+ Some("world"),
+ ])));
+ // Per-row thresholds: 2, 5, 0, 3
+ let threshold = ColumnarValue::Array(Arc::new(Int32Array::from(vec![
+ Some(2),
+ Some(5),
+ Some(0),
+ Some(3),
+ ])));
+
+ let result = spark_levenshtein(&[left, right, threshold]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.value(0), -1); // kitten->sitting=3 > 2,
return -1
+ assert_eq!(int_arr.value(1), 1); // frog->fog=1 <= 5, return 1
+ assert_eq!(int_arr.value(2), 0); // abc->abc=0 <= 0, return 0
+ assert_eq!(int_arr.value(3), -1); // hello->world=4 > 3,
return -1
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_threshold_array_with_nulls() {
+ // threshold array where some values are NULL
+ let left = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("abc"),
+ Some("hello"),
+ Some("frog"),
+ ])));
+ let right = ColumnarValue::Array(Arc::new(StringArray::from(vec![
+ Some("adc"),
+ Some("world"),
+ Some("fog"),
+ ])));
+ let threshold = ColumnarValue::Array(Arc::new(Int32Array::from(vec![
+ Some(2),
+ None, // NULL threshold for this row
+ Some(0),
+ ])));
+
+ let result = spark_levenshtein(&[left, right, threshold]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(int_arr.value(0), 1); // abc->adc=1 <= 2, return 1
+ assert!(int_arr.is_null(1)); // NULL threshold -> NULL
+ assert_eq!(int_arr.value(2), -1); // frog->fog=1 > 0, return -1
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_spark_levenshtein_threshold_negative() {
+ // Negative threshold means distance always exceeds threshold → return
-1
+ let left =
+ ColumnarValue::Array(Arc::new(StringArray::from(vec![Some("abc"),
Some("abc")])));
+ let right =
+ ColumnarValue::Array(Arc::new(StringArray::from(vec![Some("abc"),
Some("adc")])));
+ let threshold =
ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(-1), Some(-5)])));
+
+ let result = spark_levenshtein(&[left, right, threshold]).unwrap();
+ match result {
+ ColumnarValue::Array(arr) => {
+ let int_arr =
arr.as_any().downcast_ref::<Int32Array>().unwrap();
+ // dist=0 > -1 is true, so return -1
+ assert_eq!(int_arr.value(0), -1);
+ // dist=1 > -5 is true, so return -1
+ assert_eq!(int_arr.value(1), -1);
+ }
+ _ => panic!("Expected array result"),
+ }
+ }
+
+ #[test]
+ fn test_levenshtein_threshold_banded() {
+ assert_eq!(
+ levenshtein_distance_with_threshold("kitten", "sitting", 2),
+ -1
+ );
+ assert_eq!(
+ levenshtein_distance_with_threshold("kitten", "sitting", 3),
+ 3
+ );
+ assert_eq!(
+ levenshtein_distance_with_threshold("abcdefghij", "abxdefghij", 1),
+ 1
+ );
+ assert_eq!(
+ levenshtein_distance_with_threshold("short", "a much longer
string", 2),
+ -1
+ );
+ }
+}
diff --git a/native/spark-expr/src/string_funcs/mod.rs
b/native/spark-expr/src/string_funcs/mod.rs
index ca425fa0d1..5dd422a04e 100644
--- a/native/spark-expr/src/string_funcs/mod.rs
+++ b/native/spark-expr/src/string_funcs/mod.rs
@@ -18,6 +18,7 @@
mod base64;
mod contains;
mod get_json_object;
+mod levenshtein;
mod regexp_extract;
mod regexp_extract_all;
mod regexp_extract_common;
@@ -26,6 +27,7 @@ mod split;
pub use base64::spark_base64;
pub use contains::SparkContains;
pub use get_json_object::spark_get_json_object;
+pub use levenshtein::spark_levenshtein;
pub use regexp_extract::spark_regexp_extract;
pub use regexp_extract_all::spark_regexp_extract_all;
pub use split::{spark_split, spark_split_sql};
diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala
b/spark/src/main/scala/org/apache/comet/serde/strings.scala
index 1e124dbd42..1482d3e79e 100644
--- a/spark/src/main/scala/org/apache/comet/serde/strings.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala
@@ -20,7 +20,7 @@
package org.apache.comet.serde
import org.apache.spark.sql.catalyst.expressions.{Attribute, Base64,
BitLength, Cast, Concat, ConcatWs, Contains, Elt, Empty2Null, EndsWith,
Expression, FindInSet, FormatNumber, FormatString, GetJsonObject, InitCap,
Left, Length, Levenshtein, Like, Literal, Lower, Mask, OctetLength, Overlay,
RegExpExtract, RegExpExtractAll, RegExpInStr, RegExpReplace, Right, RLike,
SoundEx, StartsWith, StringLocate, StringLPad, StringRepeat, StringReplace,
StringRPad, StringSplit, StringTranslate, Substr [...]
-import org.apache.spark.sql.types.{BinaryType, DataTypes, LongType, StringType}
+import org.apache.spark.sql.types.{BinaryType, DataTypes, IntegerType,
LongType, StringType}
import org.apache.comet.CometConf
import org.apache.comet.serde.ExprOuterClass.Expr
@@ -120,6 +120,29 @@ object CometStringTranslate extends
CometScalarFunction[StringTranslate]("transl
Some(incompatReason))
}
+object CometLevenshtein extends CometExpressionSerde[Levenshtein] {
+
+ override def getUnsupportedReasons(): Seq[String] = Seq(
+ "Non-default collation (non-UTF8_BINARY) is not supported")
+
+ override def getSupportLevel(expr: Levenshtein): SupportLevel =
+ if (expr.children.exists(child =>
QueryPlanSerde.isStringCollationType(child.dataType))) {
+ Unsupported(Some("Levenshtein with non-default collation is not
supported"))
+ } else {
+ Compatible()
+ }
+
+ override def convert(
+ expr: Levenshtein,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[Expr] = {
+ val childExprs = expr.children.map(exprToProtoInternal(_, inputs, binding))
+ val optExpr =
+ scalarFunctionExprToProtoWithReturnType("levenshtein", IntegerType,
false, childExprs: _*)
+ optExprWithFallbackReason(optExpr, expr, expr.children: _*)
+ }
+}
+
object CometInitCap extends CometScalarFunction[InitCap]("initcap") with
NativeOptInAvailable {
override def getIncompatibleReasons(): Seq[String] =
@@ -655,8 +678,6 @@ object CometGetJsonObject extends
CometCodegenDispatch[GetJsonObject] with Nativ
// Expressions routed through the JVM codegen dispatcher: no native
implementation, so Spark's own
// doGenCode runs inside the Comet pipeline, matching Spark exactly.
-object CometLevenshtein extends CometCodegenDispatch[Levenshtein]
-
object CometElt extends CometCodegenDispatch[Elt]
object CometFindInSet extends CometCodegenDispatch[FindInSet]
diff --git
a/spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql
b/spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql
index 62c52c70e1..6ae8225cbe 100644
--- a/spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql
+++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql
@@ -15,17 +15,33 @@
-- specific language governing permissions and limitations
-- under the License.
--- Routes levenshtein through the codegen dispatcher so behavior matches Spark
exactly.
-
statement
-CREATE TABLE test_levenshtein(a string, b string) USING parquet
+CREATE TABLE test_levenshtein(s1 string, s2 string) USING parquet
statement
-INSERT INTO test_levenshtein VALUES ('kitten', 'sitting'), ('', 'abc'),
('abc', 'abc'), ('flaw', 'lawn'), (NULL, 'x'), ('x', NULL), (NULL, NULL)
+INSERT INTO test_levenshtein VALUES ('kitten', 'sitting'), ('frog', 'fog'),
('abc', 'abc'), ('', 'hello'), ('hello', ''), ('', ''), (NULL, 'test'),
('hello', NULL), (NULL, NULL)
+
+-- column arguments
+query
+SELECT levenshtein(s1, s2) FROM test_levenshtein
+
+-- column + literal
+query
+SELECT levenshtein(s1, 'abc') FROM test_levenshtein
+-- literal + column
query
-SELECT a, b, levenshtein(a, b) FROM test_levenshtein
+SELECT levenshtein('kitten', s2) FROM test_levenshtein
--- literal arguments
+-- literal + literal
query
-SELECT levenshtein('kitten', 'sitting'), levenshtein('', ''),
levenshtein(NULL, 'x')
+SELECT levenshtein('kitten', 'sitting'), levenshtein('frog', 'fog'),
levenshtein('', ''), levenshtein(NULL, 'a')
+
+-- identical strings
+query
+SELECT levenshtein(s1, s1) FROM test_levenshtein
+
+-- unicode characters
+query
+SELECT levenshtein('café', 'cafe'), levenshtein('你好', '你坏')
+
diff --git
a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql
b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql
new file mode 100644
index 0000000000..a0756a3403
--- /dev/null
+++
b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql
@@ -0,0 +1,60 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- MinSparkVersion: 3.5
+
+statement
+CREATE TABLE test_levenshtein(s1 string, s2 string) USING parquet
+
+statement
+INSERT INTO test_levenshtein VALUES ('kitten', 'sitting'), ('frog', 'fog'),
('abc', 'abc'), ('', 'hello'), ('hello', ''), ('', ''), (NULL, 'test'),
('hello', NULL), (NULL, NULL)
+
+-- three argument version with threshold
+query
+SELECT levenshtein('kitten', 'sitting', 2), levenshtein('kitten', 'sitting',
3), levenshtein('kitten', 'sitting', 4)
+
+-- threshold with column arguments
+query
+SELECT levenshtein(s1, s2, 2) FROM test_levenshtein
+
+-- threshold edge cases
+query
+SELECT levenshtein('abc', 'abc', 0), levenshtein('abc', 'adc', 0),
levenshtein('', '', 0)
+
+-- threshold with NULL
+query
+SELECT levenshtein('abc', 'adc', NULL), levenshtein(NULL, 'test', 2)
+
+-- threshold as column
+statement
+CREATE TABLE test_levenshtein_col(s1 string, s2 string, threshold int) USING
parquet
+
+statement
+INSERT INTO test_levenshtein_col VALUES ('kitten', 'sitting', 2), ('frog',
'fog', 5), ('abc', 'abc', 0), ('hello', 'world', 3)
+
+query
+SELECT levenshtein(s1, s2, threshold) FROM test_levenshtein_col
+
+-- threshold as column with NULLs
+statement
+CREATE TABLE test_levenshtein_col_nulls(s1 string, s2 string, threshold int)
USING parquet
+
+statement
+INSERT INTO test_levenshtein_col_nulls VALUES ('abc', 'adc', 2), ('hello',
'world', NULL), (NULL, 'test', 3), ('frog', 'fog', -1)
+
+query
+SELECT levenshtein(s1, s2, threshold) FROM test_levenshtein_col_nulls
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index aa610d0751..068e698fa6 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -132,12 +132,12 @@ class CometCodegenSuite
}
test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") {
- // With the opt-in flag on, `hypot` and `levenshtein` (both
`CometCodegenDispatch`) roll up
- // into one `[COMET-INFO: JVM codegen dispatcher: hypot, levenshtein]`
line on the
+ // With the opt-in flag on, `hypot` and `nanvl` (both
`CometCodegenDispatch`) roll up
+ // into one `[COMET-INFO: JVM codegen dispatcher: hypot, nanvl]` line on
the
// `CometProject`. With the flag off (default), no such line appears.
withTable("t") {
- sql("CREATE TABLE t (a DOUBLE, b DOUBLE, s1 STRING, s2 STRING) USING
parquet")
- sql("INSERT INTO t VALUES (3.0, 4.0, 'kitten', 'sitting')")
+ sql("CREATE TABLE t (a DOUBLE, b DOUBLE) USING parquet")
+ sql("INSERT INTO t VALUES (3.0, 4.0)")
withSQLConf(
CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true",
@@ -145,7 +145,7 @@ class CometCodegenSuite
CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true",
CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key ->
CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) {
- val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t")
+ val df = sql("SELECT hypot(a, b), nanvl(a, b) FROM t")
checkSparkAnswerAndOperator(df)
val explain =
new
ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan)
@@ -154,7 +154,7 @@ class CometCodegenSuite
s"expected a [COMET-INFO: segment, got:\n$explain")
// Names appear alphabetically via `.distinct.sorted` in
rollUpInfoMessages.
assert(
- explain.contains("JVM codegen dispatcher: hypot, levenshtein"),
+ explain.contains("JVM codegen dispatcher: hypot, nanvl"),
s"expected combined codegen-dispatch info, got:\n$explain")
}
@@ -164,7 +164,7 @@ class CometCodegenSuite
CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true",
CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key ->
CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) {
- val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t")
+ val df = sql("SELECT hypot(a, b), nanvl(a, b) FROM t")
checkSparkAnswerAndOperator(df)
val explain =
new
ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan)
diff --git
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
index 653ca6fd0b..2b862f9e60 100644
---
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
+++
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala
@@ -57,6 +57,10 @@ object CometStringExpressionBenchmark extends
CometBenchmarkBase {
StringExprConfig("initCap", "select initCap(c1) from parquetV1Table"),
StringExprConfig("instr", "select instr(c1, '123') from parquetV1Table"),
StringExprConfig("length", "select length(c1) from parquetV1Table"),
+ StringExprConfig("levenshtein", "select levenshtein(c1, 'test') from
parquetV1Table"),
+ StringExprConfig(
+ "levenshtein_threshold",
+ "select levenshtein(c1, 'test', 3) from parquetV1Table"),
StringExprConfig("like", "select c1 like '%123%' from parquetV1Table"),
StringExprConfig("lower", "select lower(c1) from parquetV1Table"),
StringExprConfig("lpad", "select lpad(c1, 150, 'x') from parquetV1Table"),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
