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

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: yusinnmao

native/spark-expr/src/string_funcs/levenshtein.rs

@@ -21,7 +21,7 @@
 //! matching Apache Spark's `levenshtein(str1, str2)` semantics.
 
 use arrow::array::{as_string_array, Array, ArrayRef, Int32Array};
-use datafusion::common::{DataFusionError, Result};
+use datafusion::common::{DataFusionError, Result, ScalarValue};
 use datafusion::physical_plan::ColumnarValue;
 use std::sync::Arc;
 
@@ -78,19 +78,40 @@ fn levenshtein_distance(s: &str, t: &str) -> i32 {
 
 /// Spark-compatible levenshtein scalar function.
 ///
-/// Accepts two string arguments and returns an Int32 array of edit distances.
-/// NULL inputs produce NULL outputs.
+/// Accepts two or three arguments:
+/// - `levenshtein(str1, str2)` → edit distance
+/// - `levenshtein(str1, str2, threshold)` → edit distance if <= threshold, else -1
+///
+/// NULL inputs produce NULL outputs. NULL threshold produces NULL output.
 pub fn spark_levenshtein(args: &[ColumnarValue]) -> Result<ColumnarValue> {
-    if args.len() != 2 {
+    if args.len() < 2 || args.len() > 3 {
         return Err(DataFusionError::Internal(format!(
-            "levenshtein requires exactly 2 arguments, got {}",
+            "levenshtein requires 2 or 3 arguments, got {}",
             args.len()
         )));
     }
 
+    // Extract optional threshold (3rd argument must be a scalar Int32)
+    let threshold: Option<i32> = if args.len() == 3 {
+        match &args[2] {
+            ColumnarValue::Scalar(ScalarValue::Int32(t)) => match t {
+                Some(val) => Some(*val),
+                None => return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None))),
+            },
+            _ => {
+                return Err(DataFusionError::Internal(
+                    "levenshtein threshold must be an Int32 scalar".to_string(),
+                ));
+            }
+        }
+    } else {
+        None
+    };
+
     // Expand scalars to arrays for uniform processing
     let len = args
         .iter()
+        .take(2)
         .find_map(|arg| match arg {
             ColumnarValue::Array(a) => Some(a.len()),
             _ => None,
@@ -107,7 +128,13 @@ pub fn spark_levenshtein(args: &[ColumnarValue]) -> Result<ColumnarValue> {
         .iter()
         .zip(right_arr.iter())
         .map(|(l, r)| match (l, r) {
-            (Some(l), Some(r)) => Some(levenshtein_distance(l, r)),
+            (Some(l), Some(r)) => {
+                let dist = levenshtein_distance(l, r);
+                match threshold {
+                    Some(t) if dist > t => Some(-1),
+                    _ => Some(dist),
+                }
+            }
             _ => None, // NULL propagation
         })
         .collect();
@@ -161,4 +188,43 @@ mod tests {
             _ => 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::Scalar(ScalarValue::Int32(None)) => {} // NULL threshold -> NULL
+            _ => panic!("Expected NULL scalar result for NULL threshold"),
+        }
+    }
 }

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -166,7 +166,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[GetJsonObject] -> CometGetJsonObject,
       classOf[InitCap] -> CometInitCap,
       classOf[Length] -> CometLength,
-      classOf[Levenshtein] -> CometScalarFunction("levenshtein"),
+      classOf[Levenshtein] -> CometLevenshtein,
       classOf[Like] -> CometLike,
       classOf[Lower] -> CometLower,
       classOf[OctetLength] -> CometScalarFunction("octet_length"),

spark/src/main/scala/org/apache/comet/serde/strings.scala

@@ -21,7 +21,7 @@ package org.apache.comet.serde
 
 import java.util.Locale
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Concat, ConcatWs, Expression, GetJsonObject, If, InitCap, IsNull, Left, Length, Like, Literal, Lower, RegExpReplace, Right, RLike, StringLPad, StringRepeat, StringRPad, StringSplit, Substring, Upper}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Concat, ConcatWs, Expression, GetJsonObject, If, InitCap, IsNull, Left, Length, Levenshtein, Like, Literal, Lower, RegExpReplace, Right, RLike, StringLPad, StringRepeat, StringRPad, StringSplit, Substring, Upper}
 import org.apache.spark.sql.types.{BinaryType, DataTypes, LongType, StringType}
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -84,6 +84,21 @@ object CometLength extends CometScalarFunction[Length]("length") {
   }
 }
 
+object CometLevenshtein extends CometScalarFunction[Levenshtein]("levenshtein") {
+
+  override def getUnsupportedReasons(): Seq[String] = Seq(
+    "Non-default collation (non-UTF8_BINARY) is not supported")
+
+  override def getSupportLevel(expr: Levenshtein): SupportLevel = {
+    expr.children.headOption match {
+      case Some(child) if QueryPlanSerde.isStringCollationType(child.dataType) =>
+        Unsupported(Some(
+          "Levenshtein with non-default collation is not supported"))
+      case _ => Compatible()
+    }
+  }
+}
+
 object CometInitCap extends CometScalarFunction[InitCap]("initcap") {
 
   override def getIncompatibleReasons(): Seq[String] = Seq(

spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql

@@ -44,3 +44,19 @@ SELECT levenshtein(s1, s1) FROM test_levenshtein
 -- unicode characters
 query
 SELECT levenshtein('café', 'cafe'), levenshtein('你好', '你坏')
+
+-- 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)

spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala

@@ -509,4 +509,21 @@ class CometStringExpressionSuite extends CometTestBase {
     }
   }
 
+  test("levenshtein with threshold") {
+    val data = Seq(
+      ("kitten", "sitting"),
+      ("frog", "fog"),
+      ("abc", "abc"),
+      ("hello", "world"))
+
+    withParquetTable(data, "tbl") {
+      checkSparkAnswerAndOperator(
+        "SELECT levenshtein(_1, _2, 2) FROM tbl")
+      checkSparkAnswerAndOperator(
+        "SELECT levenshtein(_1, _2, 0) FROM tbl")
+      checkSparkAnswerAndOperator(
+        "SELECT levenshtein(_1, _2, 10) FROM tbl")
+    }
+  }
+
 }

spark/src/test/scala/org/apache/spark/sql/benchmark/CometStringExpressionBenchmark.scala

@@ -59,6 +59,9 @@ object CometStringExpressionBenchmark extends CometBenchmarkBase {
     StringExprConfig("length", "select length(c1) from parquetV1Table"),
     StringExprConfig("levenshtein", "select levenshtein(c1, 'test') 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"),