[viirya]

Which issue does this PR close?

Closes #6811.

Rationale for this change

The native scalar Arrow UDF path from #6130 runs the Python callable on the Spark task thread, inside tokio::task::block_in_place. A Spark kill does not reach it: the Java interrupt is ignored while the thread is in Python, and Spark's PythonRunner MonitorThread, which destroys the Python worker of a killed task, does not apply. A cancelled job or the losing copy of a speculative task therefore kept its executor core, and kept competing for the embedded interpreter's GIL, until the UDF returned on its own. A stuck UDF never released them.

What changes are included in this PR?

Native side (python_udf.rs, arrow_python_udf.rs):

• A global registry maps a task attempt id to the Python thread ids currently running one of its UDF calls. ArrowPythonUdfExec::execute registers the plan under its task attempt (now passed from createPlan through ExecutionContext and the planner). Each call records its thread right before entering the callable and removes it right after, on success or error. The entry is removed when the task's last plan drops its registration, i.e. when the plan is released.
• interrupt_task marks the task as killed and raises comet.TaskKilled in its recorded threads with PyThreadState_SetAsyncExc. TaskKilled subclasses BaseException, so a user except Exception block does not swallow it. Once a task is killed, later calls fail without entering Python, and a call that returned just as the interrupt arrived clears the pending exception on its thread and fails with the same error.
• Lock ordering: holders of the registry lock never wait for the GIL. interrupt_task sets the flag without the GIL and attaches only when a call is running, so it does not touch Python for a task that is not in a UDF call.
• New JNI functions: NativeBase.interruptPythonUdfs(taskAttemptId) (a no-op that returns 0 in builds without python-udf) and Native.hasPythonUdf(plan).

JVM side:

• CometExecIterator asks hasPythonUdf only when NativeBase.supportsPythonUdf(). For a plan that contains an Arrow UDF, it starts PythonUdfMonitor, one daemon thread per task like Spark's MonitorThread, which polls TaskContext.isInterrupted every 100 ms until the task completes and calls the interrupt while the task is killed. The interrupt is repeated, so it also reaches calls that start after an earlier attempt and calls that caught the exception. Unlike Spark's monitor, it does not wait spark.python.task.killTimeout, because the interrupt is cooperative and lets the call unwind instead of killing a process. Plans without a native Arrow UDF, and builds without the feature, start no thread.
• The UDF then fails with a non-fatal error. Spark's Executor reports a non-fatal failure of a task whose reasonIfKilled is set as TaskKilled; this is the same code in Spark 4.1.3 and 4.2.0 (Executor.scala, case _: InterruptedException | NonFatal(_) if task != null && task.reasonIfKilled.isDefined).

Known limitation, now documented in the PyArrow UDF guide in place of the old "killing a task does not interrupt" note: an asynchronous exception takes effect only when the thread next runs Python bytecode. A call blocked in native code, such as a long PyArrow kernel that released the GIL, keeps the task's core until that code returns.

How are these changes tested?

• Rust: interrupts_running_call_of_killed_task runs a while True UDF that wraps its body in try/except Exception, interrupts it from another thread, and checks that the call fails with TaskKilled, that its thread is unregistered, that a later call of the killed task fails without entering Python, and that the registry entry is removed when the registration is dropped. With the PyThreadState_SetAsyncExc call stubbed out, the test fails ("the interrupted call did not return").
• Scala: CometArrowPythonUdfSuite "killing a task interrupts its running native Arrow UDF" cancels the job group of a query whose native UDF loops forever, and checks that the task ends as TaskKilled and the job fails within 30 seconds. With the monitor disabled, the task is never killed and the test fails.
• I ran the Rust tests and the full CometArrowPythonUdfSuite with -Pspark-4.1 on Python 3.10 and 3.13.

This touches the planner and a native operator, so it would be good to run it with the run-pyarrow-udf-tests and run-spark-4.1-tests labels.

[sunchao]

Reviewed the full 11-file diff against 785fdbcb4cbf9057dfcfae3bd9dd45a77bb505bb. The PR remains open and non-draft at the reviewed SHA. The snapshot and live discussion contained no reviews or comments. Two introduced P2 issues were identified.

Routed skills: review-comet-pr and review-comet-ffi-pr. Read AGENTS.md and the relevant contributor documentation.

Summary

• Prior state and problem: Native scalar Arrow UDFs continued running after Spark killed their tasks, retaining executor resources until the callable returned.
• Design approach: A per-task monitor polls cancellation state. Native registrations associate running Python threads with task attempts, and cancellation injects comet.TaskKilled through PyThreadState_SetAsyncExc.
• Correctness: Registration lifetime, lock ordering, pending-exception clearing, and rejection of subsequent calls are coherent. However, repeated injection can interrupt finally cleanup and leave shared Python state locked. A bounded reproduction demonstrated this. Spark's Executor sources confirm that the resulting nonfatal native error is classified as TaskKilled when the task already has a kill reason.
• Compatibility analysis: Checked Spark's task-failure classification across the supported 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0 sources. Native scalar Arrow execution remains restricted to its existing supported path. The shared Scala addition introduces a strict-compilation failure on Spark 3.5. The touched existing files on branch-1.2 at 889856dc5af821de1f7c6ca2e9538bd663431d04 match the PR base, so the release-relative behavior change is the cancellation mechanism reviewed here.
• Key design decisions: Using BaseException avoids ordinary except Exception handlers. Task-attempt IDs keep speculative attempts separate. The monitor avoids accessing plan pointers concurrently, and feature-disabled builds skip monitoring.
• Implementation sketch: ExecutionContext passes the task ID through the planner into ArrowPythonUdfExec. Stream-owned registrations track active calls. JNI exposes plan detection and interruption, while CometExecIterator starts one monitor per task.
• Performance: The feature adds a daemon thread and approximately ten polls per second per watched task, plus registry locking around calls. Feature-enabled plans also incur a plan traversal to detect Python UDFs. No quantified throughput regression was established. The intended benefit is reclaiming resources from cancelled Python loops.
• Design: The registry and RAII guards provide understandable ownership. Repeated exception delivery undermines the stated cooperative-unwinding behavior. Delivering cancellation once per active call would simplify that policy while retaining rejection of later calls.
• Abstraction & complexity: TaskRegistration and RunningCall separate stream lifetime from call lifetime appropriately. The JNI additions are narrow and do not change Arrow buffer ownership or conversion semantics. No additional abstraction-related P1/P2 issue was identified.
• Behavioral changes worth calling out: Cancellation now raises comet.TaskKilled without waiting for spark.python.task.killTimeout. Delivery remains deferred while native code runs. These intended differences are documented. Aborting cleanup with subsequent interrupts is an unintended additional effect.
• Suggested improvements: Prevent repeated delivery into an already-cancelling call and add a bounded cleanup regression. Explicitly discard the Boolean returned by watchedTasks.remove so the strict Scala build succeeds.

Exact-head CI: PyArrow jobs passed on Spark 4.0, 4.1, and 4.2. The Spark 4.1 log confirms the new cancellation test passed. The strict Scala job failed with the new PythonUdfMonitor.scala:61 error and four additional diagnostics in unchanged shuffle files. At the final check, 29 checks had succeeded, one failed, five were running, nine Spark SQL shards were queued, and 52 were skipped. Relevant labels were already applied.

Validation limits: No full local Comet integration build was run. A disposable Rust harness used this head's registry and call-guard code unchanged with PyO3 0.28.3 and CPython 3.12.13. Its two bounded tests demonstrated successful cleanup after one interrupt and a retained lock after repeated interruption. This was not a local Spark end-to-end reproduction. git diff --check passed, and the project checkout remained unchanged.

[sunchao]

[P2] [P2] Avoid reinjecting cancellation while a call is unwinding

A call remains in task.threads until its entire Python invocation returns, including its finally blocks. Because the monitor calls this method every 100 ms, subsequent TaskKilled exceptions also interrupt cleanup that is already handling the first one. For example, a UDF holding a module-level lock with finally: flush(); lock.release() leaves that lock permanently held when the cleanup takes longer than a polling interval. Other tasks share that module and can then block indefinitely. A single interrupt lets the same cleanup finish. Could cancellation delivery be recorded per active call so subsequent polls do not reinject into it, while retaining the killed flag that rejects later calls?

Evidence: A bounded Rust harness copied this head's TaskRegistration, RunningCall, and interrupt_task implementation unchanged. With PyO3 0.28.3 and CPython 3.12.13, the callable acquired a shared threading.Lock, entered a loop, and used finally: cleaning.set(); time.sleep(0.4); lock.release(). One interrupt produced TaskKilled and left the lock available. A second interrupt 100 ms later, after confirming cleanup had begun, also produced TaskKilled but left the lock unavailable after the call ended. Both calls were correctly unregistered, so the retained lock outlived cancellation. Command: cargo test --offline -- --nocapture --test-threads=1 in /tmp/comet7181-908ae-root-probe-xk1ssxaq. Output: repeat=false, shared lock available after cancellation=true and repeat=true, shared lock available after cancellation=false. The production monitor repeats delivery at this same interval.

[viirya]

Good catch, thanks. I reproduced it on Python 3.10 and 3.13 with a lock released in a finally after a 0.4 s sleep: a second injection during the cleanup leaves the lock held. Fixed in 89da483. Each recorded call now carries an interrupted flag, and interrupt_task raises only in calls it has not interrupted yet, so a running call gets TaskKilled exactly once. The killed flag is unchanged, so later calls of the task still fail without entering Python. The monitor keeps polling, because a task killed before its plan registers is not recorded, and a plan the task registers later still needs the flag set.

The trade-off is that a call that catches TaskKilled (e.g. except BaseException) and keeps running is no longer interrupted again. The user guide and the PythonUdfMonitor comment now say so. The new interrupts_a_call_once_so_its_cleanup_runs test interrupts a call, waits until its finally cleanup has started, checks that a second interrupt_task raises in nothing, and checks that the cleanup released its lock. It fails with the previous repeated delivery.

[sunchao]

[P2] [P2] Explicitly discard the removal result for strict Scala compilation

watchedTasks.remove(taskAttemptId) returns a Boolean, which is discarded by this finally block. The supported Spark 3.5 strict-warnings profile treats that discard as a compilation error, so the new shared source fails the PR's strict Scala job even though native Arrow UDF execution is disabled on that Spark version. Explicitly discard the result, for example with val _ = watchedTasks.remove(taskAttemptId).

Evidence: Exact-head run 40003235556, job 120591854928, executes ./mvnw -B test-compile -Pspark-3.5,delta -Pstrict-warnings -DskipTests. Its compiler output reports PythonUdfMonitor.scala:61: discarded non-Unit value, then exits with compilation failure. This diagnostic points to the newly added file. Log: https://github.com/apache/datafusion-comet/actions/runs/40003235556/job/120591854928.

[viirya]

Fixed in 89da483 with val _ = watchedTasks.remove(taskAttemptId), matching CometExecIterator. test-compile -Pspark-3.5,delta -Pstrict-warnings now passes. The four @nowarn errors in the shuffle managers came from the same failed compile and are gone with it.

[sunchao]

Reviewed commit 89da483c1e1a1dccb212fae9cb62fd21c492f200. New P1/P2 findings are detailed inline.

[sunchao]

[P1] Compute hasPythonUdf before registering the listener that calls close(). Spark can invoke a completion listener immediately when an input feeder constructs its parent iterator after the consuming task completes. close() then frees plan before this line reads it. That can terminate the executor for any plan in a Python-enabled build, even with native UDF execution disabled in the session. Cache the Boolean in the cleanup-protected createdPlan block.

Evidence: A bounded Spark 4.1.3 PipedRDD fixture recorded completed=true closeRanBeforeProbe=true using an inert handle. At this head, line 232 publishes the listener, line 375 calls releasePlan, and hasPythonUdf dereferences the released ExecutionContext to read spark_plan. The base constructor made no native plan access after publishing that listener.

[viirya]

Right: addTaskCompletionListener runs the listener at once on a completed task, and close() releases the plan before the probe read it. Fixed in 3d9fe9e. The probe now runs inside the cleanup-protected block that creates the plan, before the listener is registered. That block also binds the task's token (see the token thread) and counts the plan for the monitor, so a failure there releases the plan like any other setup failure.

[sunchao]

[P2] A first TaskKilled can strand a shared Python lock. If a UDF is waiting in logging.Handler.acquire(), delivery can occur after it acquires the RLock and before Handler.handle() enters its try/finally. The call exits owning the lock, and healthy tasks using the same handler block. Use explicit cooperative cancellation points in the shared interpreter, or isolate calls that require forced interruption.

Evidence: On CPython 3.11.15 and 3.12.13, one interrupt made the cancelled UDF return TaskKilled while still owning the handler lock. A peer’s logger.warning() timed out after 300 ms and completed when the victim thread released the lock. With zero interrupts, all three UDFs completed. This happens during acquisition, before the cleanup covered by the earlier review. CPython handler

[viirya]

This one holds and the code cannot rule it out: the exception can fire right after a C acquire() returns and before the try whose finally releases it, and a kill that lands while the call waits for the lock makes that the likely outcome rather than a narrow race. A Python __enter__ that acquires a lock has the same gap. Moving acquire() into the try only trades it for releasing a lock that was never taken. Spark avoids the problem by destroying the worker process, which a shared interpreter cannot do.

In 3d9fe9e the monitor reduces and exposes the risk instead:

• A killed task's later calls fail at once without entering Python, which is always safe.
• A running call gets TaskKilled only after it has run for spark.python.task.killTimeout since the kill, like Spark's monitor, so calls that finish within the timeout are never interrupted.
• spark.comet.exec.nativeArrowPythonUDF.interruptOnKill.enabled=false disables raising in running calls; a killed task then waits for its call to return. kill_without_interrupt_lets_the_running_call_finish covers that mode.
• The user guide describes the lock hazard, including logging before Python 3.13, next to the setting.

Explicit cooperative cancellation would not stop an arbitrary loop, which is the case #6811 is about, so I kept the forced interrupt as the default with the opt-out.

[sunchao]

[P2] Keep the watch alive until the last native UDF plan finishes closing. Spark sets isCompleted before invoking completion listeners. When a consumer takes the first batch and the producer is evaluating the next one, releasePlan can wait in BatchProducer.stop() after the monitor exits. A kill during that wait cannot interrupt even a pure Python loop, which can retain the executor slot indefinitely.

Evidence: A Spark 4.1.3 take(1) fixture used the head’s monitor, UDF module, and byte-identical BatchProducer. After cancelJobGroup, it observed completed=true, interrupted=true, zero native interrupts, and both plan close and task end pending. A direct interrupt returned 1, released close, and Spark reported TaskKilled. Spark completion order

[viirya]

Agreed. BatchProducer.stop takes the stream's mutex, which the producer holds for the whole poll, including the Python call, so releasePlan waits for that call after isCompleted is already true. Fixed in 3d9fe9e. Each watch counts the native UDF plans bound to it, CometExecIterator.close uncounts its plan after releasePlan returns, and the monitor runs until the task has completed and no counted plan remains. The new "killing a completed task interrupts the UDF call its plan release waits for" test reads a two-row Parquet file one row per batch, takes the first row while the second loops in the UDF, completes the task, and cancels the job; it fails with the previous lifetime.

[sunchao]

[P2] Key cancellation by a process-unique token shared by the plans of one TaskContext. A cancelled local task can remain in native code after its SparkContext stops. If the replacement context reuses its attempt ID, this registration inherits the old killed flag and rejects a healthy UDF. Use the same token for watchedTasks and the native registry.

Evidence: Two sequential Spark 4.1.3 local[1] contexts both assigned attempt ID 0. While the first context’s cancelled call was still waiting, the replacement’s identity UDF received input 73 and returned Arrow UDF stopped because task 0 was killed with isInterrupted=false. Attempt 1 returned 73 for the same UDF and input. Spark attempt-ID contract

[viirya]

Agreed, taskAttemptId is only unique within one SparkContext. Fixed in 3d9fe9e. PythonUdfMonitor assigns each watched TaskContext a token from a process-wide counter. CometExecIterator binds it to the plan with a new Native.setPythonUdfTask before the first executePlan, and the native registry, ArrowPythonUdfExec, and NativeBase.killPythonUdfs all use the token instead of the attempt id.

[sunchao]

[P2] Keep the call registered through user-defined result conversion and check cancellation before returning the converted batch. Spark 4.2 accepts array-like results whose __arrow_array__ executes Python inside RecordBatch.from_arrays. Dropping RunningCall here makes that code invisible to the monitor. A cancelled conversion can keep running and ultimately return a successful batch.

Evidence: On CPython 3.11.15 and 3.12.13, two interrupts during a bounded __arrow_array__ loop both returned 0. After the loop was released, native evaluation returned [73]. Stock Spark 4.2 accepted the same result shape and stopped the worker during conversion on cancellation. Spark result materialization

[viirya]

Fixed in 3d9fe9e. The call now stays recorded through the Spark 4.2 RecordBatch.from_arrays conversion, and the killed check runs after it, so a cancelled conversion fails instead of returning a batch. interrupts_a_user_defined_array_conversion interrupts a result whose __arrow_array__ loops, and fails if the call is unrecorded before the conversion.

[rich7420]

Keep registration through cast; ExtensionArray.cast still escapes cancellation and returns a successful batch after kill_task.

[viirya]

Good catch, thanks. An extension type's __arrow_ext_class__ can return an ExtensionArray subclass that overrides cast, and the same goes for __len__, the type comparison, and the type's __str__ in the mismatch error. All of them ran after the call was unrecorded and after the kill check. Fixed in 6969737. Everything from the callable to _export_to_c now runs in call_and_export while the call is still recorded, so a kill interrupts any user code in that span. The call is unrecorded only after the export, on success and on error alike, and the kill check follows it. interrupts_a_user_defined_result_cast returns an extension array whose cast loops, interrupts it, and checks that the call fails with TaskKilled. It fails when the call is unrecorded right after the callable.

[sunchao]

[P2] Avoid unconditionally clearing an interrupt that Python already delivered. On CPython 3.11 and 3.12, PyThreadState_SetAsyncExc(id, NULL) re-arms the interpreter-wide pending flag even when no thread has an exception to consume. Other UDFs then repeatedly enter the pending-event path. After unregistering the call and releasing the mutex, consume any genuinely late interrupt on its originating thread before reusing it.

Evidence: For a surviving UDF summing three million integers, median thread CPU over 13 samples rose from 100.22 to 111.71 ms on 3.11.15 and from 134.66 to 154.77 ms on 3.12.13. Both reported a pending flag of 1 with zero target threads. A disposable Python checkpoint kept the flag at 0 and also passed the case where a C callable returns directly to Rust with an interrupt pending. CPython setter · Pending-event handler

[viirya]

Agreed. PyThreadState_SetAsyncExc signals the interpreter even for a NULL exception, and only a thread that holds an exception clears the flag. Fixed in 3d9fe9e. After unregistering an interrupted call and releasing the registry lock, the thread runs a short loop of Python bytecode, so a still-pending TaskKilled fires and is discarded there, and nothing is ever cleared with NULL. consumes_an_interrupt_pending_when_the_call_returns interrupts a C callable (functools.partial(time.sleep, 0.3)) that returns straight to Rust, then runs more Python on the same thread state and checks that the interrupt does not fire there.

[sunchao]

[P2] Check the killed flag before Python::attach, and retain this check in RunningCall::enter for races. A batch submitted after the registration is already marked killed currently waits for the GIL and imports its arguments before it can fail. An unrelated UDF holding the GIL can therefore retain this cancelled task’s executor slot even though its callable will never run.

Evidence: A CPython 3.12.13 probe marked task 101 killed with no active calls, then evaluated its next batch while another UDF held the GIL in math.factorial(420000). Returning the kill error took 1,481 ms. With a pre-attach check, the same evaluation returned in under 0.1 ms.

[viirya]

Fixed in 3d9fe9e. evaluate_named checks the killed flag before Python::attach, and RunningCall::enter keeps its check for a kill that lands while the batch's arguments are imported.

[sunchao]

[P2] Release the worker’s cloned Arc<ArrowPythonUdf> before sending the result, or join the worker after the timed receive succeeds. The receiver can wake while the detached worker still owns the registration. The final assertion then intermittently fails even though the Python call was interrupted and unregistered correctly.

Evidence: The unmodified interrupts_running_call_of_killed_task failed 5 of 24 isolated CPython 3.12.13 runs at line 825 with actual Some(0) and expected None. A control changing only the worker to drop its Arc before send passed all 24 runs.

[viirya]

Fixed in 3d9fe9e. The worker drops its Arc<ArrowPythonUdf> before it sends the result, in a shared start_call helper that the interrupt tests now use.