apache / apache/datafusion-comet

Make JVM-scalar-UDF dispatch responsive to task cancellation

Open
#4,175 0 comments 0 reactions 0 assignees View on GitHub
area:expressions enhancement priority:medium
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 6h
Merged PRs (30d)
190

Description

## Describe the problem

In `RegExpLikeUDF.evaluate` (`common/src/main/scala/org/apache/comet/udf/RegExpLikeUDF.scala`, on the JVM-scalar-UDF prototype branch) the per-row loop calls `pattern.matcher(s).find()` synchronously and uninterruptibly. A pathological backtracking pattern can run for seconds (or longer) on a single row, during which Spark task cancellation (executor task kill, query cancel) has no effect: the batch must finish before any cancel signal is observed.

The same issue applies to any future stateful CometUDF that does heavy per-row work. Native side `JvmScalarUdfExpr::evaluate` also has no cancellation check between batches.

## Describe the potential solution

1. In `JvmScalarUdfExpr::evaluate`, check DataFusion's cancellation token (via `Context`) before crossing JNI for each batch.
2. In long-running UDFs, periodically check `TaskContext.get().isInterrupted()` inside the row loop (e.g. every N rows) and throw `InterruptedException` so the batch fails fast and Spark's cancel propagates.
3. Document the contract on the `CometUDF` trait: implementations of long-running UDFs should poll `TaskContext.isInterrupted()` periodically.

For `RegExpLikeUDF` specifically, even checking once per batch (before the row loop) would catch most cancel windows.

## Additional context

Identified during code review of the JVM-scalar-UDF prototype. Filed as a follow-up so the prototype PR can ship without a cancellation-protocol design.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.