feat: Support Spark levenshtein expression in native execution

[Myx778]

Which issue does this PR close?

Closes #3084

What changes are included in this PR?

Implements Spark's levenshtein(str1, str2) function as a native Comet scalar UDF, enabling it to run via DataFusion instead of falling back to the JVM.

Implementation Details

Rust (native/spark-expr/src/string_funcs/levenshtein.rs)

• Standard DP algorithm with O(min(m,n)) space optimization
• Unicode character-level distance (not byte-level), matching Spark semantics
• Proper NULL propagation: any NULL input → NULL output
• Unit tests for basic cases, Unicode, and NULL handling

Scala Serde (strings.scala + QueryPlanSerde.scala)

• Registered via CometScalarFunction("levenshtein") — leverages existing ScalarFunc proto pathway
• No new protobuf message needed (reuses generic ScalarFunc)

UDF Registration (comet_scalar_funcs.rs)

• Registered spark_levenshtein in the create_comet_physical_fun_with_eval_mode match

Tests (CometStringExpressionSuite.scala)

• test("levenshtein") — basic edit distance computation
• test("levenshtein with nulls") — NULL propagation
• test("levenshtein with unicode") — character-level distance for CJK and emoji

How are these changes tested?

• Rust unit tests: cargo test -p datafusion-comet-spark-expr
• Spark integration tests: CometStringExpressionSuite (3 new test cases)
• All tests use checkSparkAnswerAndOperator to verify result matches Spark AND runs natively

[coderfender]

Thank you for the PR . @Myx778 . Could you please also add SLT tests to match sparks behavior?

[Myx778]

Thanks for the review! I've added SLT tests in spark/src/test/resources/sql-tests/expressions/string/levenshtein.sql.

The test file covers:

• Column arguments with NULL propagation (NULL input → NULL output)
• Column + literal and literal + column combinations
• Literal-only edge cases (empty string, NULL)
• Identical string comparison (distance = 0)
• Unicode character-level distance (café/cafe, 你好/你坏)

[andygrove]

Could you add tests for the three argument version where the third argument is threshold.

@ExpressionDescription(
  usage = """
    _FUNC_(str1, str2[, threshold]) - Returns the Levenshtein distance between the two given strings. If threshold is set and distance more than it, return -1.""",
  examples = """
    Examples:
      > SELECT _FUNC_('kitten', 'sitting');
       3
      > SELECT _FUNC_('kitten', 'sitting', 2);
       -1
  """,
  since = "1.5.0",
  group = "string_funcs")

[Myx778]

Done! I've implemented full three-argument threshold support:

Rust (levenshtein.rs):

• spark_levenshtein now accepts 2 or 3 args
• 3rd arg is extracted as ScalarValue::Int32 threshold
• Returns -1 when distance > threshold, matching Spark semantics
• NULL threshold → NULL result

Tests added:

• Rust unit tests: test_spark_levenshtein_with_threshold and test_spark_levenshtein_null_threshold
• SLT: threshold with literals, columns, edge cases (threshold=0), and NULL threshold
• Scala: "levenshtein with threshold" test with checkSparkAnswerAndOperator
• Benchmark: levenshtein_threshold entry

Commit: 7f37fd8

[andygrove]

Spark 4 lets levenshtein accept StringTypeWithCollation. With a non-default collation like utf8_lcase, Spark expects collation-aware comparison. The PR does no collation check. Could you implement getSupportLevel for this expression so it falls back for unsupported collations in Spark 4.

[Myx778]

Done! Added CometLevenshtein serde object in strings.scala with getSupportLevel that checks for non-default collation:

object CometLevenshtein extends CometScalarFunction[Levenshtein]("levenshtein") {
  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()
    }
  }
}

This uses the existing isStringCollationType from CometTypeShim (which checks collationId != StringType.collationId on Spark 4, no-op on Spark 3.x). Non-default collations will fall back to Spark's implementation.

Updated QueryPlanSerde.scala mapping from generic CometScalarFunction("levenshtein") to CometLevenshtein.

Commit: 7f37fd8

[andygrove]

Could you add a microbenchmark in spark/src/test/scala/org/apache/spark/sql/benchmark/ and share performance numbers?

[andygrove]

Could you update expressions.md to show that this is now supported

[Myx778]

Thanks @andygrove for the review!

Done:

1. Benchmark: Added levenshtein entry to CometStringExpressionBenchmark. I'll run the benchmark and share performance numbers once CI passes.
2. expressions.md: Updated to list Levenshtein as supported.

Commit: 8bdc34a

[Myx778]

Benchmark Results

Ran CometStringExpressionBenchmark on GitHub Codespace (4-core AMD EPYC 7763, 16GB RAM, Ubuntu 22.04, JDK 17.0.18).

levenshtein (2-arg)

OpenJDK 64-Bit Server VM 17.0.18+8-Ubuntu-124.04.1 on Linux 6.8.0-1044-azure
AMD EPYC 7763 64-Core Processor
levenshtein:                              Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
------------------------------------------------------------------------------------------------------------------------
Spark                                                21             24           3          0.0       20689.0       1.0X
Comet (Scan)                                         22             24           2          0.0       21926.5       0.9X
Comet (Scan + Exec)                                  22             25           3          0.0       21696.1       1.0X

levenshtein_threshold (3-arg)

OpenJDK 64-Bit Server VM 17.0.18+8-Ubuntu-124.04.1 on Linux 6.8.0-1044-azure
AMD EPYC 7763 64-Core Processor
levenshtein_threshold:                    Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
------------------------------------------------------------------------------------------------------------------------
Spark                                                16             18           2          0.1       15537.4       1.0X
Comet (Scan)                                         17             20           3          0.1       16607.5       0.9X
Comet (Scan + Exec)                                  17             19           3          0.1       16605.2       0.9X

Analysis

Performance is roughly on par with Spark for both variants (~1.0X). This is expected because:

1. The benchmark dataset is small (1024 rows with short repeated strings), so the overhead of JVM↔Native transitions dominates the actual computation time
2. levenshtein is a character-by-character comparison — the real gains from native execution would show on larger datasets where SIMD auto-vectorization and cache-efficient memory layout have more impact
3. The WARNING: Benchmark plan is NOT fully Comet native! shows the expression falls back through ColumnarToRow, which adds conversion overhead

The key benefit is correctness — ensuring the native path produces identical results to Spark, which is validated by the unit tests and SLT tests.

[andygrove]

Thanks @Myx778 - could you run "make format" to fix CI failures:

--- /__w/datafusion-comet/datafusion-comet/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala
+++ <expected fix>
@@ -493,7 +493,7 @@
       sql(s"CREATE TABLE $table(s1 STRING, s2 STRING) USING parquet")
       sql(
         s"INSERT INTO $table VALUES " +
-          s"('abc', 'adc'), (NULL, 'test'), ('hello', NULL), (NULL, NULL)")
+          "('abc', 'adc'), (NULL, 'test'), ('hello', NULL), (NULL, NULL)")
       checkSparkAnswerAndOperator(s"SELECT levenshtein(s1, s2) FROM $table")
     }
   }

[andygrove]

There are some test failures:

== Exception ==
org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 3405.0 failed 1 times, most recent failure: Lost task 0.0 in stage 3405.0 (TID 5800) (localhost executor driver): org.apache.comet.CometNativeException: Error from DataFusion: Function 'levenshtein' expects 2 arguments but received 3.

[andygrove]

The current CometLevenshtein inherits CometScalarFunction.convert, which does not set return_type on the proto. When the native planner sees no return type it falls back to session_ctx.udf("levenshtein"), which finds DataFusion's built-in 2-arg LevenshteinFunc and fails signature validation as soon as a third arg is present. That is what CometNativeException: Error
from DataFusion: Function 'levenshtein' expects 2 arguments but received 3. is telling us in the 3.5 / 4.0 logs. Could you override convert and use scalarFunctionExprToProtoWithReturnType("levenshtein", IntegerType, false, ...) so the planner skips the registry lookup, the same way CometStringSplit and CometGetJsonObject do?

[Myx778]

Fixed! Changed CometLevenshtein from extending CometScalarFunction to CometExpressionSerde, and overrode convert() to use scalarFunctionExprToProtoWithReturnType:

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: _*)
  optExprWithInfo(optExpr, expr, expr.children: _*)
}

This ensures the proto includes return_type = IntegerType, so the native planner uses our registered Comet UDF directly instead of looking up DataFusion's built-in 2-arg LevenshteinFunc.

Also fixed the Spotless issue (removed unnecessary s prefix).

Commit: f099ffe

[Myx778]

Thanks @andygrove for catching both issues!

Fixed in commit f099ffe:

1. 3-arg test failure: Changed CometLevenshtein to override convert() using scalarFunctionExprToProtoWithReturnType("levenshtein", IntegerType, false, ...). This embeds the return type in the proto so the native planner uses our registered Comet UDF directly, avoiding the conflict with DataFusion's built-in 2-arg LevenshteinFunc.
2. Spotless: Removed unnecessary s string interpolation prefix.

[andygrove]

Thanks for addressing review comments @Myx778. I kicked off CI.

[andygrove]

Could you add a test where the threshold argument is a column rather than a literal? I am not sure that the implementation handles this currently

[Myx778]

Done! The implementation now supports threshold as both scalar and column.

Changes:

• Refactored spark_levenshtein to use into_array() which handles both ColumnarValue::Scalar and ColumnarValue::Array uniformly
• Per-row NULL threshold → NULL result for that row
• Added Scala integration tests (levenshtein with threshold as column + null/negative cases)
• Added SQL file tests with threshold column

Rust unit test results (all 8 passing):

running 8 tests
test string_funcs::levenshtein::tests::test_levenshtein_basic ... ok
test string_funcs::levenshtein::tests::test_levenshtein_unicode ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_nulls ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_null_threshold ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_with_threshold ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_threshold_as_array ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_threshold_array_with_nulls ... ok
test string_funcs::levenshtein::tests::test_spark_levenshtein_threshold_negative ... ok

test result: ok. 8 passed; 0 failed; 0 ignored; 0 measured; 341 filtered out

[Myx778]

Hi @andygrove, just a gentle ping — I've pushed the changes you requested (threshold as column support + tests). The CI is waiting for approval to run. Whenever you have a moment, could you kick it off? No rush at all, just wanted to make sure it didn't slip through. Thanks for all your guidance on this PR!

[Myx778]

Hi @andygrove, the 2 failing checks are HiveSparkSubmitSuite: SPARK-8020 which is a Hive metastore infrastructure test unrelated to this PR (the executor was killed with exit code 143 — likely OOM in the 512MB local-cluster mode).

Could you re-run the CI or approve when you get a chance? All other 69 checks pass. Thanks!

[Myx778]

Hi @andygrove, I've rebased onto the latest main to resolve the merge conflicts (the branch was 163 commits behind). The conflicts were straightforward — mostly the optExprWithInfo → optExprWithFallbackReason rename from #4508 and an import addition for SubstringIndex. The PR is now mergeable again. Whenever you have a moment, could you kick off CI? Thanks!

[Myx778]

Hi @andygrove and @mbutrovich,

Hope you both are doing well! I just rebased onto the latest main again to resolve fresh merge conflicts (the branch had fallen 82 commits behind since the previous rebase).

A note on one substantive merge decision worth flagging:

• main now contains a CometLevenshtein extends CometCodegenDispatch[Levenshtein] (JVM codegen dispatch path). I replaced it with this PR's native implementation, since the two would otherwise collide and a native path is what this PR was originally designed to deliver. Happy to revisit this if you'd prefer a different approach — for instance, an opt-in config flag that selects between native vs codegen-dispatch, similar to the CometCaseConversionBase pattern.

The PR has been through ~5 review rounds since April and I've tried to address all substantive feedback. The branch is mergeable now and I've kept the commit history clean.

Whenever either of you has a moment, would you be able to take another look, or kick off CI so we can confirm everything is green? No urgency at all — I really appreciate the time you've already put into reviewing this. Thanks so much for your guidance throughout!

Best regards,
Yuxin

[andygrove]

Thanks for the persistence on this one @Myx778, and thanks especially for calling out the codegen-dispatch vs native decision explicitly. A few thoughts after going through the diff and checking the Rust against Spark's UTF8String.levenshteinDistance.

Correctness looks good. I traced the threshold logic against Spark's banded limitedCompare algorithm. Spark returns the true distance D when D <= threshold and -1 otherwise, and the PR's dist > t ? -1 : dist is exactly equivalent for all threshold values including negatives and the empty-string cases. NULL propagation matches Spark too. Nice work on the test coverage here.

The main thing to weigh is whether the native path is worth it. main already routes Levenshtein through CometCodegenDispatch, which runs Spark's own doGenCode inside the kernel and is Spark-exact with no correctness risk. The benchmark numbers show this native version is roughly on par with Spark (~1.0X, and 0.9X for scan-only) rather than faster. Could you spell out the concrete motivation for the native path? For example, that it keeps the pipeline native when COMET_SCALA_UDF_CODEGEN_ENABLED is off. That helps justify replacing a zero-risk path with a hand-written Rust implementation to maintain.

A couple of smaller code points:

• In spark_levenshtein, could we use as_generic_string_array::<i32>(...)? instead of as_string_array(...)? The latter panics on a type mismatch and only handles Utf8, whereas the rest of string_funcs/ (see split.rs, regexp_extract_all.rs) uses the fallible form and also handles LargeUtf8.
• getSupportLevel checks collation on children.headOption only. Since the native code does raw char comparison and ignores collation, checking all string children (children.exists(...)) would be safer for the case where only the right argument is collated.
• The new CometStringExpressionSuite cases overlap heavily with the SQL file tests. Since we prefer CometSqlFileTestSuite for expression coverage, would you be open to trimming the Scala duplicates?

Also the PR is showing merge conflicts again, so it will need another rebase onto main before CI is meaningful.

[Myx778]

Thank you very much for the detailed review and for validating the correctness against Spark's implementation, @andygrove. Your concern about the maintenance trade-off is entirely fair.

I agree that keeping the pipeline native only when COMET_SCALA_UDF_CODEGEN_ENABLED is disabled is not, by itself, a strong enough justification, especially because the dispatcher is enabled by default and the current benchmark is roughly at parity. The intended benefits are to avoid per-task closure serialization/Janino compilation overhead and to retain native execution when the dispatcher is disabled or cannot handle the expression, but these benefits need to be demonstrated with measurements rather than assumed.

I have pushed commit e602d97f, which syncs the branch with current main and addresses the smaller points you raised:

• uses fallible generic string conversion and supports both Utf8 and LargeUtf8;
• checks collation on all string children;
• removes the overlapping Scala tests in favor of the SQL file tests;
• uses a banded implementation for the threshold form instead of computing the full matrix.

For the native-path decision, I plan to optimize and benchmark the implementation more thoroughly, including ASCII and Unicode data, short and long strings, literal and column operands, several threshold values, and dispatcher-enabled/disabled execution. The main optimization candidates are an ASCII fast path, common-prefix/suffix trimming, reusable batch-level work buffers, and a bit-parallel path for short strings.

If those changes do not produce a clear, repeatable advantage over CometCodegenDispatch, I agree that retaining the Spark-exact dispatcher as the default is preferable to adding a hand-written native implementation without sufficient benefit.

The new CI workflows have been triggered, but they are currently waiting for an Apache repository maintainer to approve the fork workflow runs. Thank you again for the thoughtful feedback.