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

Author: yusinnmao

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

@@ -82,7 +82,8 @@ fn levenshtein_distance(s: &str, t: &str) -> i32 {
 /// - `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.
+/// 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!(
@@ -91,27 +92,9 @@ pub fn spark_levenshtein(args: &[ColumnarValue]) -> Result<ColumnarValue> {
         )));
     }
 
-    // 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
+    // Determine array length from any array argument
     let len = args
         .iter()
-        .take(2)
         .find_map(|arg| match arg {
             ColumnarValue::Array(a) => Some(a.len()),
             _ => None,
@@ -124,22 +107,56 @@ pub fn spark_levenshtein(args: &[ColumnarValue]) -> Result<ColumnarValue> {
     let left_arr = as_string_array(&left);
     let right_arr = as_string_array(&right);
 
-    let result: Int32Array = left_arr
-        .iter()
-        .zip(right_arr.iter())
-        .map(|(l, r)| match (l, r) {
-            (Some(l), Some(r)) => {
-                let dist = levenshtein_distance(l, r);
-                match threshold {
-                    Some(t) if dist > t => Some(-1),
-                    _ => Some(dist),
+    // Handle the optional threshold argument (scalar or array)
+    if args.len() == 3 {
+        let threshold_array = args[2].clone().into_array(len)?;
+        let threshold_arr = threshold_array
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .ok_or_else(|| {
+                DataFusionError::Internal(
+                    "levenshtein threshold must be Int32".to_string(),
+                )
+            })?;
+
+        let result: Int32Array = left_arr
+            .iter()
+            .zip(right_arr.iter())
+            .enumerate()
+            .map(|(i, (l, r))| {
+                // If threshold is NULL for this row, result is NULL
+                if threshold_arr.is_null(i) {
+                    return None;
                 }
-            }
-            _ => None, // NULL propagation
-        })
-        .collect();
+                match (l, r) {
+                    (Some(l), Some(r)) => {
+                        let dist = levenshtein_distance(l, r);
+                        let t = threshold_arr.value(i);
+                        if dist > t {
+                            Some(-1)
+                        } else {
+                            Some(dist)
+                        }
+                    }
+                    _ => None, // NULL propagation
+                }
+            })
+            .collect();
 
-    Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+        Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+    } else {
+        // No threshold: just compute distance
+        let result: Int32Array = left_arr
+            .iter()
+            .zip(right_arr.iter())
+            .map(|(l, r)| match (l, r) {
+                (Some(l), Some(r)) => Some(levenshtein_distance(l, r)),
+                _ => None, // NULL propagation
+            })
+            .collect();
+
+        Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+    }
 }
 
 #[cfg(test)]
@@ -223,8 +240,108 @@ mod tests {
 
         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"),
+            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"),
         }
     }
 }

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

@@ -38,3 +38,23 @@ SELECT levenshtein('abc', 'abc', 0), levenshtein('abc', 'adc', 0), levenshtein('
 -- 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

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

@@ -750,4 +750,30 @@ class CometStringExpressionSuite extends CometTestBase {
     }
   }
 
+  test("levenshtein with threshold as column") {
+    assume(isSpark35Plus, "levenshtein with threshold requires Spark 3.5+")
+    val table = "levenshtein_col_threshold_test"
+    withTable(table) {
+      sql(s"CREATE TABLE $table(s1 STRING, s2 STRING, threshold INT) USING parquet")
+      sql(
+        s"INSERT INTO $table VALUES " +
+          "('kitten', 'sitting', 2), ('frog', 'fog', 5), ('abc', 'abc', 0), ('hello', 'world', 3)")
+      // threshold as column reference
+      checkSparkAnswerAndOperator(s"SELECT levenshtein(s1, s2, threshold) FROM $table")
+    }
+  }
+
+  test("levenshtein with threshold as column with nulls") {
+    assume(isSpark35Plus, "levenshtein with threshold requires Spark 3.5+")
+    val table = "levenshtein_col_threshold_null_test"
+    withTable(table) {
+      sql(s"CREATE TABLE $table(s1 STRING, s2 STRING, threshold INT) USING parquet")
+      sql(
+        s"INSERT INTO $table VALUES " +
+          "('abc', 'adc', 2), ('hello', 'world', NULL), (NULL, 'test', 3), ('frog', 'fog', -1)")
+      // NULL threshold and NULL strings should produce NULL; negative threshold returns -1
+      checkSparkAnswerAndOperator(s"SELECT levenshtein(s1, s2, threshold) FROM $table")
+    }
+  }
+
 }