From 908ae970b018d855759dc08d44b1ee8909ac0cfa Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 23 Oct 2026 11:22:43 -0700 Subject: [PATCH 1/4] feat: interrupt a native scalar Arrow UDF when its task is killed A native scalar Arrow UDF runs its Python callable on the Spark task thread, which a Spark kill does not reach. A cancelled job or the losing copy of a speculative task kept its core and competed for the interpreter lock until the UDF returned on its own. Native code now records, per task attempt, the Python threads running a UDF call. `NativeBase.interruptPythonUdfs` marks the task as killed and raises `comet.TaskKilled`, a `BaseException` subclass, in those threads with `PyThreadState_SetAsyncExc`; later calls of the task fail without entering Python. The registry entry is removed when the task's last native plan is released. When a native plan contains an Arrow UDF, `CometExecIterator` starts one daemon monitor thread per task, like Spark's `MonitorThread`, that polls `TaskContext.isInterrupted` and calls the interrupt until the task completes. Spark reports the resulting non-fatal failure of a killed task as TaskKilled. The exception takes effect when the thread next runs Python bytecode, so a call blocked in native code that released the GIL is interrupted only when it returns. The user guide documents this. Closes #6811. --- docs/source/user-guide/latest/pyarrow-udfs.md | 12 +- native/core/src/execution/jni_api.rs | 23 ++ .../execution/operators/arrow_python_udf.rs | 16 +- native/core/src/execution/planner.rs | 9 + native/core/src/execution/python_udf.rs | 230 +++++++++++++++++- native/core/src/lib.rs | 17 ++ .../java/org/apache/comet/NativeBase.java | 9 + .../org/apache/comet/CometExecIterator.scala | 5 + .../main/scala/org/apache/comet/Native.scala | 8 + .../org/apache/comet/PythonUdfMonitor.scala | 69 ++++++ .../sql/comet/CometArrowPythonUdfSuite.scala | 57 +++++ 11 files changed, 440 insertions(+), 15 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md index 64ead611edc..8a6bd503082 100644 --- a/docs/source/user-guide/latest/pyarrow-udfs.md +++ b/docs/source/user-guide/latest/pyarrow-udfs.md @@ -163,11 +163,13 @@ lock, so multiple partitions may be slower than Spark's separate Python workers; that release the lock can still run concurrently. Python execution stays synchronous on JVM input paths and hands off other async tasks when it runs on a Tokio worker. `pyspark.TaskContext.get()` returns `None` inside a native UDF. A native extension crash or `os._exit` terminates the executor -process. Killing a task does not interrupt a running Python call: Spark stops its Python worker -`spark.python.task.killTimeout` after interrupting the task, but a native UDF runs on the task -thread until it returns, so a cancelled job or the losing copy of a speculative task keeps its core -and competes for the interpreter lock until then. A failing UDF reports its Python traceback in a -Comet native error rather than in Spark's `PythonException`. The operator reports the time spent in +process. Killing a task interrupts its running Python call by raising `comet.TaskKilled`, a +`BaseException` that an `except Exception` block does not catch, in the thread running it, within +about 100 ms and without waiting for `spark.python.task.killTimeout`. Python handles the exception +only when it next runs Python bytecode, so a call blocked in native code, such as a long PyArrow +kernel that released the interpreter lock, keeps the task's core until that code returns. A +failing UDF reports its Python traceback in a Comet native error rather than in Spark's +`PythonException`. The operator reports the time spent in Python as its `python_time` metric. The embedded Python interpreter and PyArrow allocate outside Comet's memory pool. Those allocations are also absent from the executor's `Comet native memory usage: allocated` figure and are not limited by `spark.executor.pyspark.memory`. Budget them in executor memory overhead in diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index c768d81f3f2..f2090f62fae 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -531,6 +531,8 @@ struct ExecutionContext { pub id: i64, /// The deserialized Spark plan pub spark_plan: Operator, + /// The Spark task attempt running the plan + pub task_attempt_id: i64, /// The number of partitions pub partition_count: usize, /// The DataFusion root operator converted from the `spark_plan` @@ -872,6 +874,7 @@ fn create_execution_context( Ok(Box::new(ExecutionContext { id, spark_plan, + task_attempt_id, partition_count, root_op: None, scans: vec![], @@ -1476,6 +1479,7 @@ fn execute_plan( let start = Instant::now(); let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) + .with_task_attempt_id(exec_context.task_attempt_id) .with_strict_timestamp_labels(exec_context.strict_timestamp_labels) .with_sql_text_pool(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) @@ -1705,6 +1709,25 @@ fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext { } } +/// Whether the plan evaluates a Python UDF in the embedded interpreter, which the JVM then +/// interrupts when the task is killed. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_hasPythonUdf( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, +) -> jboolean { + fn contains_python_udf(op: &Operator) -> bool { + matches!(op.op_struct, Some(OpStruct::ArrowPythonUdf(_))) + || op.children.iter().any(contains_python_udf) + } + try_unwrap_or_throw(&e, |_| { + Ok(contains_python_udf( + &get_execution_context(exec_context).spark_plan, + )) + }) +} + /// Returns the partition offsets published by a finished native shuffle write. /// /// The returned array holds `num_output_partitions + 1` offsets, the last being the total data diff --git a/native/core/src/execution/operators/arrow_python_udf.rs b/native/core/src/execution/operators/arrow_python_udf.rs index 8652c1f283f..d4df3da95b9 100644 --- a/native/core/src/execution/operators/arrow_python_udf.rs +++ b/native/core/src/execution/operators/arrow_python_udf.rs @@ -36,7 +36,7 @@ use datafusion::physical_plan::{ use futures::stream; use futures::StreamExt; -use crate::execution::python_udf::ArrowPythonUdf; +use crate::execution::python_udf::{ArrowPythonUdf, TaskRegistration}; #[derive(Debug, Clone)] pub struct ArrowPythonUdfSpec { @@ -58,6 +58,8 @@ pub struct ArrowPythonUdfExec { max_records_per_batch: usize, max_bytes_per_batch: usize, accept_array_like_results: bool, + /// The Spark task attempt whose kill interrupts the UDF calls; `None` without a task. + task_attempt_id: Option, schema: SchemaRef, cache: Arc, metrics: ExecutionPlanMetricsSet, @@ -70,6 +72,7 @@ impl ArrowPythonUdfExec { max_records_per_batch: i32, max_bytes_per_batch: i64, accept_array_like_results: bool, + task_attempt_id: Option, ) -> Result { if specs.is_empty() { return exec_err!("ArrowPythonUdfExec requires at least one UDF"); @@ -106,6 +109,7 @@ impl ArrowPythonUdfExec { max_records_per_batch: max_records_per_batch.max(0) as usize, max_bytes_per_batch: max_bytes_per_batch.max(0) as usize, accept_array_like_results, + task_attempt_id, schema, cache, metrics: ExecutionPlanMetricsSet::new(), @@ -293,6 +297,7 @@ impl ExecutionPlan for ArrowPythonUdfExec { self.max_records_per_batch as i32, self.max_bytes_per_batch as i64, self.accept_array_like_results, + self.task_attempt_id, )?)) } @@ -302,6 +307,10 @@ impl ExecutionPlan for ArrowPythonUdfExec { context: Arc, ) -> Result { let input = self.child.execute(partition, context)?; + // Dropped with the stream, when the native plan is released. + let task = self + .task_attempt_id + .map(|id| Arc::new(TaskRegistration::new(id))); let workers: Vec<_> = self .specs .iter() @@ -314,6 +323,10 @@ impl ExecutionPlan for ArrowPythonUdfExec { self.accept_array_like_results, &spec.python_version, ) + .map(|worker| match &task { + Some(task) => worker.with_task(Arc::clone(task)), + None => worker, + }) }) .collect::>()?; let workers = Arc::new(workers); @@ -465,6 +478,7 @@ mod tests { 1, 0, false, + None, ) .unwrap(); let results: Vec<_> = exec diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a47c19d1a23..dfc130e6e2a 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -203,6 +203,8 @@ pub struct PhysicalPlanner { /// Whether a boundary that would cast a timestamp's timezone label to the declared one fails /// instead. Test-only; see `spark.comet.testing.strictTimestampLabels.enabled`. strict_timestamp_labels: bool, + /// The Spark task attempt running the plan, if any. Killing it interrupts Arrow UDFs. + task_attempt_id: Option, } impl Default for PhysicalPlanner { @@ -224,6 +226,7 @@ impl PhysicalPlanner { shuffle_partition_pusher: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, + task_attempt_id: None, } } @@ -302,6 +305,11 @@ impl PhysicalPlanner { self } + pub fn with_task_attempt_id(mut self, task_attempt_id: i64) -> Self { + self.task_attempt_id = Some(task_attempt_id); + self + } + pub fn with_strict_timestamp_labels(mut self, strict: bool) -> Self { self.strict_timestamp_labels = strict; self @@ -1665,6 +1673,7 @@ impl PhysicalPlanner { udf.max_records_per_batch, udf.max_bytes_per_batch, udf.accept_array_like_results, + self.task_attempt_id, )?); Ok(( scans, diff --git a/native/core/src/execution/python_udf.rs b/native/core/src/execution/python_udf.rs index eb110ff954b..b69de44f512 100644 --- a/native/core/src/execution/python_udf.rs +++ b/native/core/src/execution/python_udf.rs @@ -29,7 +29,10 @@ use pyo3::ffi::Py_uintptr_t; use pyo3::prelude::*; use pyo3::sync::MutexExt; use pyo3::types::{PyBytes, PyList, PyTuple}; -use std::sync::{Mutex, PoisonError}; +use pyo3::PyTypeInfo; +use std::collections::HashMap; +use std::os::raw::{c_long, c_ulong}; +use std::sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError}; // Tasks create their callables concurrently. Concurrent first imports of a package // can hand a thread a partially initialized module (observed with PyArrow on Python @@ -37,6 +40,153 @@ use std::sync::{Mutex, PoisonError}; // which waits for the lock without holding the GIL. static IMPORT_LOCK: Mutex<()> = Mutex::new(()); +pyo3::create_exception!( + comet, + TaskKilled, + pyo3::exceptions::PyBaseException, + "Raised in a running Arrow UDF when Spark kills its task." +); + +unsafe extern "C" { + // Part of the stable Python C API, but not declared by pyo3-ffi. + fn PyThread_get_thread_ident() -> c_ulong; +} + +/// The native plans of one Spark task attempt that hold Arrow UDFs. +#[derive(Default)] +struct TaskCalls { + plans: usize, + killed: bool, + /// Python thread ids of the calls running now, one entry per call. + threads: Vec, +} + +// Holders of this lock never wait for the GIL, so a thread attached to Python may take it. +static TASKS: LazyLock>> = LazyLock::new(Default::default); + +fn tasks() -> MutexGuard<'static, HashMap> { + TASKS.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// Registers a native plan's Arrow UDFs under their Spark task attempt, so that +/// [`interrupt_task`] reaches their calls. The task's entry is removed when its last +/// registration is dropped, which happens when the plan is released. +pub struct TaskRegistration { + task_attempt_id: i64, +} + +impl TaskRegistration { + pub fn new(task_attempt_id: i64) -> Self { + tasks().entry(task_attempt_id).or_default().plans += 1; + Self { task_attempt_id } + } + + fn is_killed(&self) -> bool { + tasks() + .get(&self.task_attempt_id) + .is_some_and(|task| task.killed) + } + + fn killed_error(&self) -> ArrowError { + ArrowError::ComputeError(format!( + "Arrow UDF stopped because task {} was killed", + self.task_attempt_id + )) + } +} + +impl Drop for TaskRegistration { + fn drop(&mut self) { + let mut tasks = tasks(); + if let Some(task) = tasks.get_mut(&self.task_attempt_id) { + task.plans -= 1; + if task.plans == 0 { + tasks.remove(&self.task_attempt_id); + } + } + } +} + +/// A call to the user's callable, recorded so that [`interrupt_task`] can raise in its thread. +struct RunningCall<'a> { + task: &'a TaskRegistration, + thread: c_ulong, +} + +impl<'a> RunningCall<'a> { + /// Fails without recording the call when the task has already been killed. + fn enter(_py: Python<'_>, task: &'a TaskRegistration) -> Result { + // SAFETY: the caller is attached to the interpreter. + let thread = unsafe { PyThread_get_thread_ident() }; + let mut tasks = tasks(); + let calls = tasks.entry(task.task_attempt_id).or_default(); + if calls.killed { + return Err(task.killed_error()); + } + calls.threads.push(thread); + Ok(Self { task, thread }) + } +} + +impl Drop for RunningCall<'_> { + fn drop(&mut self) { + let mut tasks = tasks(); + let Some(calls) = tasks.get_mut(&self.task.task_attempt_id) else { + return; + }; + if let Some(index) = calls.threads.iter().position(|t| *t == self.thread) { + calls.threads.swap_remove(index); + } + if calls.killed { + // An interrupt that arrived as the callable returned is still pending on this + // thread. Nothing can set another one now that the call is unrecorded, so clear + // it rather than let it fire in later Python code on this thread. + // SAFETY: the call is made while attached, on the thread it names. + unsafe { + pyo3::ffi::PyThreadState_SetAsyncExc(self.thread as c_long, std::ptr::null_mut()) + }; + } + } +} + +/// Marks a Spark task attempt as killed and raises [`TaskKilled`] in its running Arrow UDF +/// calls. Later calls of the task fail without entering Python. Returns the number of calls +/// interrupted. A task with no registered plan is not recorded, so the caller repeats this +/// until the task completes. +/// +/// The exception takes effect when the thread next runs Python bytecode: a call blocked in +/// native code that released the GIL is interrupted only when that code returns. +pub fn interrupt_task(task_attempt_id: i64) -> usize { + let running = match tasks().get_mut(&task_attempt_id) { + Some(task) => { + task.killed = true; + !task.threads.is_empty() + } + None => false, + }; + if !running { + return 0; + } + // Calls cannot start after the flag is set, so only calls recorded before it remain. + Python::attach(|py| { + let exception = TaskKilled::type_object(py); + let tasks = tasks(); + let Some(task) = tasks.get(&task_attempt_id) else { + return 0; + }; + task.threads + .iter() + .map(|thread| { + // SAFETY: attached to the interpreter, and `exception` is a live type object. + let set = unsafe { + pyo3::ffi::PyThreadState_SetAsyncExc(*thread as c_long, exception.as_ptr()) + }; + set.max(0) as usize + }) + .sum() + }) +} + fn initialize_python() -> Result<()> { use std::ffi::CStr; use std::sync::OnceLock; @@ -168,6 +318,7 @@ pub struct ArrowPythonUdf { allow_cast: bool, safe_cast: bool, accept_array_like: bool, + task: Option>, } impl ArrowPythonUdf { @@ -234,10 +385,17 @@ impl ArrowPythonUdf { allow_cast, safe_cast, accept_array_like, + task: None, }) }) } + /// Lets [`interrupt_task`] interrupt this UDF's calls for the registered task. + pub fn with_task(mut self, task: Arc) -> Self { + self.task = Some(task); + self + } + /// Evaluate one Arrow batch, with the same row count for every argument. /// Python receives and returns `pyarrow.Array` objects via the Arrow C Data /// interface; no row conversion or Arrow IPC serialization occurs here. @@ -297,14 +455,18 @@ impl ArrowPythonUdf { kwargs.set_item(name, arg).map_err(python_error)?; } } - let result = self - .callable - .bind(py) - .call( - PyTuple::new(py, positional).map_err(python_error)?, - Some(&kwargs), - ) - .map_err(|error| python_traceback_error(py, error))?; + let positional = PyTuple::new(py, positional).map_err(python_error)?; + let call = self + .task + .as_deref() + .map(|task| RunningCall::enter(py, task)) + .transpose()?; + let result = self.callable.bind(py).call(positional, Some(&kwargs)); + drop(call); + let result = result.map_err(|error| python_traceback_error(py, error))?; + if let Some(task) = self.task.as_deref().filter(|task| task.is_killed()) { + return Err(task.killed_error()); + } let result = if self.accept_array_like { // Spark 4.2 builds its output batch with `RecordBatch.from_arrays`, // which converts lists and NumPy arrays and rejects chunked arrays. @@ -600,6 +762,56 @@ mod tests { }); } + fn running_calls(task_attempt_id: i64) -> Option { + tasks().get(&task_attempt_id).map(|task| task.threads.len()) + } + + #[test] + fn interrupts_running_call_of_killed_task() { + const TASK: i64 = -6811; + initialize_python().unwrap(); + // `except Exception` must not swallow the interrupt. + let command = Python::attach(|py| { + cloudpickled_lambda( + py, + c"(lambda ns: (exec('def f(a):\\n n = 0\\n while True:\\n n += 1\\n try:\\n n %= 7\\n except Exception:\\n pass\\n', ns), ns['f'])[1])({})", + ) + }); + let task = Arc::new(TaskRegistration::new(TASK)); + let udf = Arc::new( + ArrowPythonUdf::from_command(&command, DataType::Int64, true, true, false, "") + .unwrap() + .with_task(Arc::clone(&task)), + ); + let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let (sender, receiver) = std::sync::mpsc::channel(); + { + let (udf, input) = (Arc::clone(&udf), Arc::clone(&input)); + // Not joined: a call that ignores the interrupt would never return. + std::thread::spawn(move || sender.send(udf.evaluate(&[input], 1))); + } + let start = std::time::Instant::now(); + while running_calls(TASK) != Some(1) { + assert!(start.elapsed().as_secs() < 30, "the UDF call did not start"); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert_eq!(interrupt_task(TASK), 1); + let error = receiver + .recv_timeout(std::time::Duration::from_secs(30)) + .expect("the interrupted call did not return") + .unwrap_err() + .to_string(); + assert!(error.contains("TaskKilled"), "{error}"); + assert_eq!(running_calls(TASK), Some(0)); + // A killed task does not enter Python again. + let error = udf.evaluate(&[input], 1).unwrap_err().to_string(); + assert!(error.contains("was killed"), "{error}"); + assert_eq!(interrupt_task(TASK), 0); + drop(udf); + drop(task); + assert_eq!(running_calls(TASK), None); + } + const UNBUFFERED_CHILD: &str = "COMET_ARROW_UDF_UNBUFFERED_CHILD"; const UNBUFFERED_MARKER: &str = "comet-arrow-udf-unbuffered-marker"; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index bc382540e32..edffd09389b 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -190,6 +190,23 @@ pub extern "system" fn Java_org_apache_comet_NativeBase_nativeSupportsPythonUdf( cfg!(feature = "python-udf") } +/// Interrupts the Python UDF calls of a killed Spark task; see +/// `execution::python_udf::interrupt_task`. Returns the number of calls interrupted. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeBase_interruptPythonUdfs( + _: EnvUnowned, + _: JClass, + task_attempt_id: jni::sys::jlong, +) -> jni::sys::jint { + #[cfg(feature = "python-udf")] + return execution::python_udf::interrupt_task(task_attempt_id) as jni::sys::jint; + #[cfg(not(feature = "python-udf"))] + { + let _ = task_attempt_id; + 0 + } +} + #[no_mangle] /// Releases the global Tokio runtime used by Comet native execution. pub extern "system" fn Java_org_apache_comet_NativeBase_release(_e: EnvUnowned, _class: JClass) { diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index bce5da59d8b..315e5b03e5f 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -329,6 +329,15 @@ public static boolean supportsPythonUdf() { } } + /** + * Raises an exception in the running Python UDF calls of a killed task, and makes its later calls + * fail without entering Python. A no-op for a task with no native plan holding a Python UDF, so + * callers repeat it until the task completes. + * + * @return the number of calls interrupted + */ + public static native int interruptPythonUdfs(long taskAttemptId); + /** Release native resources through JNI */ static native void release(); diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 5790554a6bd..f51bf5aedb3 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -235,6 +235,11 @@ class CometExecIterator( CometExecIterator.startMemoryUsageLog() + // A native Python UDF ignores the task's kill, so watch the task to interrupt its calls. + if (PythonUdfMonitor.supported && nativeLib.hasPythonUdf(plan)) { + PythonUdfMonitor.watch(TaskContext.get()) + } + /** * Matches any throwable once a JVM input of this plan has failed, and extracts what that input * threw. Native saw only its text. See `CometArrowStream.inputFailure`. diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 4e46b664e8b..e58f150664f 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -130,6 +130,14 @@ class Native extends NativeBase { */ @native def getShufflePartitionOffsets(plan: Long): Array[Long] + /** + * Whether the native query plan runs a Python UDF in the embedded interpreter. + * + * @param plan + * the address to native query plan. + */ + @native def hasPythonUdf(plan: Long): Boolean + /** * Release and drop the native query plan object and context object. * diff --git a/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala new file mode 100644 index 00000000000..6559f838f4f --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import java.util.concurrent.ConcurrentHashMap + +import org.apache.spark.TaskContext +import org.apache.spark.internal.Logging + +/** + * Interrupts the native Python UDF calls of killed tasks. Spark's Python runner has a + * `MonitorThread` that destroys the worker of a killed task. A native UDF runs in the embedded + * interpreter on the task thread instead, where neither a Java interrupt nor a stopped input + * reaches it. So one daemon thread per task polls its kill state and, while the task is killed + * but still running, raises an exception in the task's running Python calls. Calls that start + * after the first interrupt fail without entering Python. + */ +object PythonUdfMonitor extends Logging { + + // Raising the exception is cheap and lets the call unwind on its own, so unlike Spark's + // monitor this one does not wait `spark.python.task.killTimeout` before acting. Repeating it + // reaches calls that start after an earlier attempt and calls that caught the exception. + private val PollIntervalMs = 100L + + private val watchedTasks = ConcurrentHashMap.newKeySet[java.lang.Long]() + + /** Whether the native library can run Python UDFs, and so may need its calls interrupted. */ + lazy val supported: Boolean = NativeBase.supportsPythonUdf() + + /** Starts watching the task unless it is watched already. The watch ends with the task. */ + def watch(context: TaskContext): Unit = { + val taskAttemptId = context.taskAttemptId() + if (watchedTasks.add(taskAttemptId)) { + val thread = new Thread(s"Comet Python UDF monitor for task $taskAttemptId") { + override def run(): Unit = { + try { + while (!context.isCompleted()) { + if (context.isInterrupted() && NativeBase.interruptPythonUdfs(taskAttemptId) > 0) { + logInfo(s"Interrupted Python UDF calls of killed task $taskAttemptId") + } + Thread.sleep(PollIntervalMs) + } + } finally { + watchedTasks.remove(taskAttemptId) + } + } + } + thread.setDaemon(true) + thread.start() + } + } +} diff --git a/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala index 41ddba9f8cc..dbe6a8f039d 100644 --- a/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala +++ b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala @@ -19,19 +19,26 @@ package org.apache.spark.sql.comet +import java.io.File import java.util.{Base64, Collections} +import java.util.concurrent.{CountDownLatch, TimeUnit} +import scala.concurrent.{ExecutionContext, Future} +import scala.concurrent.duration._ import scala.sys.process._ import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.spark.{SparkException, TaskKilled} import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction} +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.execution.python.{ArrowEvalPythonExec, UserDefinedPythonFunction} import org.apache.spark.sql.functions.{array, col, expr, length, lit, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, TimeType, VariantType, YearMonthIntervalType} +import org.apache.spark.util.ThreadUtils import org.apache.comet.{CometConf, CometExplainInfo, NativeBase} import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus @@ -528,6 +535,56 @@ class CometArrowPythonUdfSuite extends CometTestBase { } } + test("killing a task interrupts its running native Arrow UDF") { + assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") + + withTempDir { dir => + val started = new File(dir, "started") + // Loops in Python bytecode, where the interrupt takes effect, and survives a user + // `except Exception`. + val loop = "(lambda ns: (exec('def f(a):\\n" + + " open(\"" + started.getCanonicalPath + "\", \"w\").close()\\n" + + " n = 0\\n while True:\\n n += 1\\n try:\\n" + + " n %= 7\\n except Exception:\\n pass\\n', ns), " + + "ns['f'])[1])({})" + val udf = arrowUdf("loop", pickledCommand(loop, "LongType()"), LongType) + val killed = new CountDownLatch(1) + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEnd.reason match { + case _: TaskKilled => killed.countDown() + case _ => + } + } + spark.sparkContext.addSparkListener(listener) + try { + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = spark.range(0, 1, 1, 1).select(udf(col("id"))) + assert(isNative(df)) + val group = "comet-arrow-udf-kill" + val job = Future { + spark.sparkContext.setJobGroup( + group, + "native Arrow UDF loop", + interruptOnCancel = true) + df.collect() + }(ExecutionContext.global) + val deadline = 60.seconds.fromNow + while (!started.exists()) { + assert(deadline.hasTimeLeft(), "the UDF did not start") + Thread.sleep(50) + } + spark.sparkContext.cancelJobGroup(group) + assert(killed.await(30, TimeUnit.SECONDS), "the task running the UDF was not killed") + intercept[SparkException](ThreadUtils.awaitResult(job, 30.seconds)) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + } + test("native Arrow UDF accepts array-like results on Spark 4.2") { assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") assume(isSpark42Plus, "Spark 4.1 requires a pyarrow.Array result") From 89da483c1e1a1dccb212fae9cb62fd21c492f200 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 23 Oct 2026 18:10:08 -0700 Subject: [PATCH 2/4] fix: raise TaskKilled once per Python UDF call and fix strict Scala build The monitor polls every 100 ms and raised `comet.TaskKilled` in a killed task's running calls on every poll. A call stays recorded until it returns, so a later poll raised again inside the `finally` blocks or context manager exits that the first exception started, which could leave a shared lock held. Each recorded call now remembers whether it was interrupted, and `interrupt_task` raises only in calls it has not interrupted yet. Later calls of a killed task are still rejected. Also discard the result of `watchedTasks.remove` explicitly, which the strict-warnings build reported as a discarded non-Unit value. --- docs/source/user-guide/latest/pyarrow-udfs.md | 12 +- native/core/src/execution/python_udf.rs | 127 ++++++++++++++---- .../org/apache/comet/PythonUdfMonitor.scala | 9 +- 3 files changed, 115 insertions(+), 33 deletions(-) diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md index 8a6bd503082..6ba06b01031 100644 --- a/docs/source/user-guide/latest/pyarrow-udfs.md +++ b/docs/source/user-guide/latest/pyarrow-udfs.md @@ -165,11 +165,13 @@ paths and hands off other async tasks when it runs on a Tokio worker. `pyspark.T returns `None` inside a native UDF. A native extension crash or `os._exit` terminates the executor process. Killing a task interrupts its running Python call by raising `comet.TaskKilled`, a `BaseException` that an `except Exception` block does not catch, in the thread running it, within -about 100 ms and without waiting for `spark.python.task.killTimeout`. Python handles the exception -only when it next runs Python bytecode, so a call blocked in native code, such as a long PyArrow -kernel that released the interpreter lock, keeps the task's core until that code returns. A -failing UDF reports its Python traceback in a Comet native error rather than in Spark's -`PythonException`. The operator reports the time spent in +about 100 ms and without waiting for `spark.python.task.killTimeout`. It is raised once per call, +so `finally` blocks and context managers can clean up, and a call that catches it and keeps +running is not interrupted again. Python handles the exception only when it next runs Python +bytecode, so a call blocked in native code, such as a long PyArrow kernel that released the +interpreter lock, keeps the task's core until that code returns. A failing UDF reports its Python +traceback in a Comet native error rather than in Spark's `PythonException`. The operator reports +the time spent in Python as its `python_time` metric. The embedded Python interpreter and PyArrow allocate outside Comet's memory pool. Those allocations are also absent from the executor's `Comet native memory usage: allocated` figure and are not limited by `spark.executor.pyspark.memory`. Budget them in executor memory overhead in diff --git a/native/core/src/execution/python_udf.rs b/native/core/src/execution/python_udf.rs index b69de44f512..bc48eb2e7bc 100644 --- a/native/core/src/execution/python_udf.rs +++ b/native/core/src/execution/python_udf.rs @@ -57,8 +57,16 @@ unsafe extern "C" { struct TaskCalls { plans: usize, killed: bool, - /// Python thread ids of the calls running now, one entry per call. - threads: Vec, + /// The calls running now. + calls: Vec, +} + +struct Call { + /// Python thread id of the thread running the call. + thread: c_ulong, + /// Whether [`interrupt_task`] has raised in the call. It raises once, so that a second + /// exception does not abort the cleanup the first one started. + interrupted: bool, } // Holders of this lock never wait for the GIL, so a thread attached to Python may take it. @@ -123,7 +131,10 @@ impl<'a> RunningCall<'a> { if calls.killed { return Err(task.killed_error()); } - calls.threads.push(thread); + calls.calls.push(Call { + thread, + interrupted: false, + }); Ok(Self { task, thread }) } } @@ -134,10 +145,10 @@ impl Drop for RunningCall<'_> { let Some(calls) = tasks.get_mut(&self.task.task_attempt_id) else { return; }; - if let Some(index) = calls.threads.iter().position(|t| *t == self.thread) { - calls.threads.swap_remove(index); - } - if calls.killed { + let Some(index) = calls.calls.iter().position(|c| c.thread == self.thread) else { + return; + }; + if calls.calls.swap_remove(index).interrupted { // An interrupt that arrived as the callable returned is still pending on this // thread. Nothing can set another one now that the call is unrecorded, so clear // it rather than let it fire in later Python code on this thread. @@ -150,9 +161,9 @@ impl Drop for RunningCall<'_> { } /// Marks a Spark task attempt as killed and raises [`TaskKilled`] in its running Arrow UDF -/// calls. Later calls of the task fail without entering Python. Returns the number of calls -/// interrupted. A task with no registered plan is not recorded, so the caller repeats this -/// until the task completes. +/// calls that it has not interrupted yet. Later calls of the task fail without entering +/// Python. Returns the number of calls interrupted. A task with no registered plan is not +/// recorded, so the caller repeats this until the task completes. /// /// The exception takes effect when the thread next runs Python bytecode: a call blocked in /// native code that released the GIL is interrupted only when that code returns. @@ -160,7 +171,7 @@ pub fn interrupt_task(task_attempt_id: i64) -> usize { let running = match tasks().get_mut(&task_attempt_id) { Some(task) => { task.killed = true; - !task.threads.is_empty() + task.calls.iter().any(|call| !call.interrupted) } None => false, }; @@ -170,20 +181,22 @@ pub fn interrupt_task(task_attempt_id: i64) -> usize { // Calls cannot start after the flag is set, so only calls recorded before it remain. Python::attach(|py| { let exception = TaskKilled::type_object(py); - let tasks = tasks(); - let Some(task) = tasks.get(&task_attempt_id) else { + let mut tasks = tasks(); + let Some(task) = tasks.get_mut(&task_attempt_id) else { return 0; }; - task.threads - .iter() - .map(|thread| { - // SAFETY: attached to the interpreter, and `exception` is a live type object. - let set = unsafe { - pyo3::ffi::PyThreadState_SetAsyncExc(*thread as c_long, exception.as_ptr()) - }; - set.max(0) as usize - }) - .sum() + let mut interrupted = 0; + for call in task.calls.iter_mut().filter(|call| !call.interrupted) { + // SAFETY: attached to the interpreter, and `exception` is a live type object. + let set = unsafe { + pyo3::ffi::PyThreadState_SetAsyncExc(call.thread as c_long, exception.as_ptr()) + }; + if set > 0 { + call.interrupted = true; + interrupted += 1; + } + } + interrupted }) } @@ -763,7 +776,7 @@ mod tests { } fn running_calls(task_attempt_id: i64) -> Option { - tasks().get(&task_attempt_id).map(|task| task.threads.len()) + tasks().get(&task_attempt_id).map(|task| task.calls.len()) } #[test] @@ -812,6 +825,72 @@ mod tests { assert_eq!(running_calls(TASK), None); } + #[test] + fn interrupts_a_call_once_so_its_cleanup_runs() { + const TASK: i64 = -6812; + initialize_python().unwrap(); + let command = Python::attach(|py| { + let sys = py.import("sys").unwrap(); + let lock = py + .import("threading") + .unwrap() + .call_method0("Lock") + .unwrap(); + sys.setattr("comet_cleanup_lock", lock).unwrap(); + sys.setattr("comet_cleaning", false).unwrap(); + // The cleanup outlasts the monitor's 100 ms polling interval. + cloudpickled_lambda( + py, + c"(lambda ns: (exec('def f(a):\\n import sys, time\\n sys.comet_cleanup_lock.acquire()\\n try:\\n n = 0\\n while True:\\n n += 1\\n finally:\\n sys.comet_cleaning = True\\n time.sleep(0.4)\\n sys.comet_cleanup_lock.release()\\n', ns), ns['f'])[1])({})", + ) + }); + let sys_attr = |expression: &std::ffi::CStr| -> bool { + Python::attach(|py| { + let sys = py.import("sys").unwrap(); + let locals = pyo3::types::PyDict::new(py); + locals.set_item("sys", sys).unwrap(); + py.eval(expression, None, Some(&locals)) + .unwrap() + .extract() + .unwrap() + }) + }; + let task = Arc::new(TaskRegistration::new(TASK)); + let udf = Arc::new( + ArrowPythonUdf::from_command(&command, DataType::Int64, true, true, false, "") + .unwrap() + .with_task(Arc::clone(&task)), + ); + let (sender, receiver) = std::sync::mpsc::channel(); + { + let udf = Arc::clone(&udf); + let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); + std::thread::spawn(move || sender.send(udf.evaluate(&[input], 1))); + } + let start = std::time::Instant::now(); + while running_calls(TASK) != Some(1) { + assert!(start.elapsed().as_secs() < 30, "the UDF call did not start"); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + assert_eq!(interrupt_task(TASK), 1); + while !sys_attr(c"sys.comet_cleaning") { + assert!(start.elapsed().as_secs() < 30, "the cleanup did not start"); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + // A later poll while the call cleans up must not raise in it again. + assert_eq!(interrupt_task(TASK), 0); + let error = receiver + .recv_timeout(std::time::Duration::from_secs(30)) + .expect("the interrupted call did not return") + .unwrap_err() + .to_string(); + assert!(error.contains("TaskKilled"), "{error}"); + assert!( + !sys_attr(c"sys.comet_cleanup_lock.locked()"), + "the cleanup did not release its lock" + ); + } + const UNBUFFERED_CHILD: &str = "COMET_ARROW_UDF_UNBUFFERED_CHILD"; const UNBUFFERED_MARKER: &str = "comet-arrow-udf-unbuffered-marker"; diff --git a/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala index 6559f838f4f..5bd32f80f4b 100644 --- a/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala +++ b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala @@ -29,14 +29,15 @@ import org.apache.spark.internal.Logging * `MonitorThread` that destroys the worker of a killed task. A native UDF runs in the embedded * interpreter on the task thread instead, where neither a Java interrupt nor a stopped input * reaches it. So one daemon thread per task polls its kill state and, while the task is killed - * but still running, raises an exception in the task's running Python calls. Calls that start - * after the first interrupt fail without entering Python. + * but still running, raises an exception in the task's running Python calls, once per call so + * that its cleanup can run. Calls that start after the first interrupt fail without entering + * Python. */ object PythonUdfMonitor extends Logging { // Raising the exception is cheap and lets the call unwind on its own, so unlike Spark's // monitor this one does not wait `spark.python.task.killTimeout` before acting. Repeating it - // reaches calls that start after an earlier attempt and calls that caught the exception. + // reaches calls of plans that the task registers after an earlier attempt. private val PollIntervalMs = 100L private val watchedTasks = ConcurrentHashMap.newKeySet[java.lang.Long]() @@ -58,7 +59,7 @@ object PythonUdfMonitor extends Logging { Thread.sleep(PollIntervalMs) } } finally { - watchedTasks.remove(taskAttemptId) + val _ = watchedTasks.remove(taskAttemptId) } } } From 3d9fe9ec923c39a178d53b0ca8a676f29b10e31f Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 24 Oct 2026 18:17:10 -0700 Subject: [PATCH 3/4] fix: address review of native Arrow UDF task kill - Read whether the plan has a Python UDF before the completion listener that can release the plan at once, instead of after it. - Keep a task's monitor running until its native UDF plans are released. Releasing a plan waits for the call in flight, which a kill must still stop after the task has completed. - Key the native registry by a token the JVM assigns to each TaskContext, bound to the plan with `Native.setPythonUdfTask`. A task attempt id is unique only within one SparkContext. - Keep a call recorded while Spark 4.2 converts its array-like result, which can run a user-defined `__arrow_array__`, and check the kill after the conversion. - Consume an interrupt still pending when a call ends by running Python bytecode, instead of clearing it with `PyThreadState_SetAsyncExc(id, NULL)`, which leaves the interpreter's pending-event flag set on CPython 3.11 and 3.12. - Fail a killed task's batch before waiting for the GIL. - Make later calls of a killed task fail at once, but raise in a running call only after `spark.python.task.killTimeout`, like Spark's monitor. The exception can strand a lock acquired just before a `try` block, so `spark.comet.exec.nativeArrowPythonUDF.interruptOnKill.enabled=false` leaves running calls alone. The user guide documents the risk. - Release the test worker's registration before it reports its result. --- docs/source/user-guide/latest/pyarrow-udfs.md | 17 +- native/core/src/execution/jni_api.rs | 26 +- .../execution/operators/arrow_python_udf.rs | 15 +- native/core/src/execution/planner.rs | 13 +- native/core/src/execution/python_udf.rs | 427 ++++++++++++------ native/core/src/lib.rs | 13 +- .../java/org/apache/comet/NativeBase.java | 10 +- .../scala/org/apache/comet/CometConf.scala | 13 + .../org/apache/comet/CometExecIterator.scala | 25 +- .../main/scala/org/apache/comet/Native.scala | 11 + .../org/apache/comet/PythonUdfMonitor.scala | 95 ++-- .../sql/comet/CometArrowPythonUdfSuite.scala | 123 +++-- 12 files changed, 539 insertions(+), 249 deletions(-) diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md index 6ba06b01031..bc21a74cd67 100644 --- a/docs/source/user-guide/latest/pyarrow-udfs.md +++ b/docs/source/user-guide/latest/pyarrow-udfs.md @@ -163,15 +163,18 @@ lock, so multiple partitions may be slower than Spark's separate Python workers; that release the lock can still run concurrently. Python execution stays synchronous on JVM input paths and hands off other async tasks when it runs on a Tokio worker. `pyspark.TaskContext.get()` returns `None` inside a native UDF. A native extension crash or `os._exit` terminates the executor -process. Killing a task interrupts its running Python call by raising `comet.TaskKilled`, a -`BaseException` that an `except Exception` block does not catch, in the thread running it, within -about 100 ms and without waiting for `spark.python.task.killTimeout`. It is raised once per call, -so `finally` blocks and context managers can clean up, and a call that catches it and keeps +process. Once a task is killed, its later Python calls do not run. A call still running +`spark.python.task.killTimeout` (2 seconds by default) after the kill gets `comet.TaskKilled` raised +in it, a `BaseException` that an `except Exception` block does not catch. It is raised once per +call, so `finally` blocks and context managers can clean up, and a call that catches it and keeps running is not interrupted again. Python handles the exception only when it next runs Python bytecode, so a call blocked in native code, such as a long PyArrow kernel that released the -interpreter lock, keeps the task's core until that code returns. A failing UDF reports its Python -traceback in a Comet native error rather than in Spark's `PythonException`. The operator reports -the time spent in +interpreter lock, keeps the task's core until that code returns. The exception can also arrive +between a lock's `acquire()` and the `try` block that releases it, as in `logging` before Python +3.13, and then the lock stays held for every task in the executor's shared interpreter. Set +`spark.comet.exec.nativeArrowPythonUDF.interruptOnKill.enabled=false` to let a killed task's +running call return on its own instead. A failing UDF reports its Python traceback in a Comet +native error rather than in Spark's `PythonException`. The operator reports the time spent in Python as its `python_time` metric. The embedded Python interpreter and PyArrow allocate outside Comet's memory pool. Those allocations are also absent from the executor's `Comet native memory usage: allocated` figure and are not limited by `spark.executor.pyspark.memory`. Budget them in executor memory overhead in diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index f2090f62fae..32538644edc 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -531,8 +531,9 @@ struct ExecutionContext { pub id: i64, /// The deserialized Spark plan pub spark_plan: Operator, - /// The Spark task attempt running the plan - pub task_attempt_id: i64, + /// The JVM's token for the Spark task running the plan, set when the JVM stops the plan's + /// Python UDFs on a kill + pub python_udf_task: Option, /// The number of partitions pub partition_count: usize, /// The DataFusion root operator converted from the `spark_plan` @@ -874,7 +875,7 @@ fn create_execution_context( Ok(Box::new(ExecutionContext { id, spark_plan, - task_attempt_id, + python_udf_task: None, partition_count, root_op: None, scans: vec![], @@ -1479,7 +1480,7 @@ fn execute_plan( let start = Instant::now(); let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) - .with_task_attempt_id(exec_context.task_attempt_id) + .with_python_udf_task(exec_context.python_udf_task) .with_strict_timestamp_labels(exec_context.strict_timestamp_labels) .with_sql_text_pool(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) @@ -1710,7 +1711,7 @@ fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext { } /// Whether the plan evaluates a Python UDF in the embedded interpreter, which the JVM then -/// interrupts when the task is killed. +/// stops when the task is killed. #[no_mangle] pub extern "system" fn Java_org_apache_comet_Native_hasPythonUdf( e: EnvUnowned, @@ -1728,6 +1729,21 @@ pub extern "system" fn Java_org_apache_comet_Native_hasPythonUdf( }) } +/// Sets the JVM's token for the plan's Spark task, which `NativeBase.killPythonUdfs` takes to +/// stop the plan's Python UDFs. Must precede the first `executePlan`, which creates them. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_setPythonUdfTask( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, + task: jlong, +) { + try_unwrap_or_throw(&e, |_| { + get_execution_context(exec_context).python_udf_task = Some(task); + Ok(()) + }) +} + /// Returns the partition offsets published by a finished native shuffle write. /// /// The returned array holds `num_output_partitions + 1` offsets, the last being the total data diff --git a/native/core/src/execution/operators/arrow_python_udf.rs b/native/core/src/execution/operators/arrow_python_udf.rs index d4df3da95b9..72577ffcfb3 100644 --- a/native/core/src/execution/operators/arrow_python_udf.rs +++ b/native/core/src/execution/operators/arrow_python_udf.rs @@ -58,8 +58,9 @@ pub struct ArrowPythonUdfExec { max_records_per_batch: usize, max_bytes_per_batch: usize, accept_array_like_results: bool, - /// The Spark task attempt whose kill interrupts the UDF calls; `None` without a task. - task_attempt_id: Option, + /// The JVM's token for the Spark task whose kill stops the UDF calls; `None` when no JVM + /// monitors the task. + python_udf_task: Option, schema: SchemaRef, cache: Arc, metrics: ExecutionPlanMetricsSet, @@ -72,7 +73,7 @@ impl ArrowPythonUdfExec { max_records_per_batch: i32, max_bytes_per_batch: i64, accept_array_like_results: bool, - task_attempt_id: Option, + python_udf_task: Option, ) -> Result { if specs.is_empty() { return exec_err!("ArrowPythonUdfExec requires at least one UDF"); @@ -109,7 +110,7 @@ impl ArrowPythonUdfExec { max_records_per_batch: max_records_per_batch.max(0) as usize, max_bytes_per_batch: max_bytes_per_batch.max(0) as usize, accept_array_like_results, - task_attempt_id, + python_udf_task, schema, cache, metrics: ExecutionPlanMetricsSet::new(), @@ -297,7 +298,7 @@ impl ExecutionPlan for ArrowPythonUdfExec { self.max_records_per_batch as i32, self.max_bytes_per_batch as i64, self.accept_array_like_results, - self.task_attempt_id, + self.python_udf_task, )?)) } @@ -309,8 +310,8 @@ impl ExecutionPlan for ArrowPythonUdfExec { let input = self.child.execute(partition, context)?; // Dropped with the stream, when the native plan is released. let task = self - .task_attempt_id - .map(|id| Arc::new(TaskRegistration::new(id))); + .python_udf_task + .map(|task| Arc::new(TaskRegistration::new(task))); let workers: Vec<_> = self .specs .iter() diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index dfc130e6e2a..85e4e9d63e0 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -203,8 +203,9 @@ pub struct PhysicalPlanner { /// Whether a boundary that would cast a timestamp's timezone label to the declared one fails /// instead. Test-only; see `spark.comet.testing.strictTimestampLabels.enabled`. strict_timestamp_labels: bool, - /// The Spark task attempt running the plan, if any. Killing it interrupts Arrow UDFs. - task_attempt_id: Option, + /// The JVM's token for the Spark task running the plan, if the JVM stops its Arrow UDFs + /// when the task is killed. + python_udf_task: Option, } impl Default for PhysicalPlanner { @@ -226,7 +227,7 @@ impl PhysicalPlanner { shuffle_partition_pusher: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, - task_attempt_id: None, + python_udf_task: None, } } @@ -305,8 +306,8 @@ impl PhysicalPlanner { self } - pub fn with_task_attempt_id(mut self, task_attempt_id: i64) -> Self { - self.task_attempt_id = Some(task_attempt_id); + pub fn with_python_udf_task(mut self, python_udf_task: Option) -> Self { + self.python_udf_task = python_udf_task; self } @@ -1673,7 +1674,7 @@ impl PhysicalPlanner { udf.max_records_per_batch, udf.max_bytes_per_batch, udf.accept_array_like_results, - self.task_attempt_id, + self.python_udf_task, )?); Ok(( scans, diff --git a/native/core/src/execution/python_udf.rs b/native/core/src/execution/python_udf.rs index bc48eb2e7bc..b18395af95c 100644 --- a/native/core/src/execution/python_udf.rs +++ b/native/core/src/execution/python_udf.rs @@ -52,7 +52,9 @@ unsafe extern "C" { fn PyThread_get_thread_ident() -> c_ulong; } -/// The native plans of one Spark task attempt that hold Arrow UDFs. +/// The native plans of one Spark task that hold Arrow UDFs. Tasks are keyed by a token the JVM +/// assigns to each `TaskContext`, because a task attempt id is unique only within one +/// `SparkContext`. #[derive(Default)] struct TaskCalls { plans: usize, @@ -64,8 +66,8 @@ struct TaskCalls { struct Call { /// Python thread id of the thread running the call. thread: c_ulong, - /// Whether [`interrupt_task`] has raised in the call. It raises once, so that a second - /// exception does not abort the cleanup the first one started. + /// Whether [`kill_task`] has raised in the call. It raises once, so that a second exception + /// does not abort the cleanup the first one started. interrupted: bool, } @@ -76,46 +78,41 @@ fn tasks() -> MutexGuard<'static, HashMap> { TASKS.lock().unwrap_or_else(PoisonError::into_inner) } -/// Registers a native plan's Arrow UDFs under their Spark task attempt, so that -/// [`interrupt_task`] reaches their calls. The task's entry is removed when its last -/// registration is dropped, which happens when the plan is released. +/// Registers a native plan's Arrow UDFs under their Spark task, so that [`kill_task`] reaches +/// their calls. The task's entry is removed when its last registration is dropped, which +/// happens when the plan is released. pub struct TaskRegistration { - task_attempt_id: i64, + task: i64, } impl TaskRegistration { - pub fn new(task_attempt_id: i64) -> Self { - tasks().entry(task_attempt_id).or_default().plans += 1; - Self { task_attempt_id } + pub fn new(task: i64) -> Self { + tasks().entry(task).or_default().plans += 1; + Self { task } } fn is_killed(&self) -> bool { - tasks() - .get(&self.task_attempt_id) - .is_some_and(|task| task.killed) + tasks().get(&self.task).is_some_and(|task| task.killed) } fn killed_error(&self) -> ArrowError { - ArrowError::ComputeError(format!( - "Arrow UDF stopped because task {} was killed", - self.task_attempt_id - )) + ArrowError::ComputeError("Arrow UDF stopped because its Spark task was killed".to_string()) } } impl Drop for TaskRegistration { fn drop(&mut self) { let mut tasks = tasks(); - if let Some(task) = tasks.get_mut(&self.task_attempt_id) { + if let Some(task) = tasks.get_mut(&self.task) { task.plans -= 1; if task.plans == 0 { - tasks.remove(&self.task_attempt_id); + tasks.remove(&self.task); } } } } -/// A call to the user's callable, recorded so that [`interrupt_task`] can raise in its thread. +/// Python code run by a UDF call, recorded so that [`kill_task`] can raise in its thread. struct RunningCall<'a> { task: &'a TaskRegistration, thread: c_ulong, @@ -127,7 +124,7 @@ impl<'a> RunningCall<'a> { // SAFETY: the caller is attached to the interpreter. let thread = unsafe { PyThread_get_thread_ident() }; let mut tasks = tasks(); - let calls = tasks.entry(task.task_attempt_id).or_default(); + let calls = tasks.entry(task.task).or_default(); if calls.killed { return Err(task.killed_error()); } @@ -137,41 +134,56 @@ impl<'a> RunningCall<'a> { }); Ok(Self { task, thread }) } + + /// Unrecords the call. Returns whether [`kill_task`] raised in it. + fn unregister(&self) -> bool { + let mut tasks = tasks(); + tasks + .get_mut(&self.task.task) + .and_then(|calls| { + let index = calls.calls.iter().position(|c| c.thread == self.thread)?; + Some(calls.calls.swap_remove(index).interrupted) + }) + .unwrap_or(false) + } + + /// Unrecords the call. If it was interrupted, the exception may still be pending, for + /// example when the interrupt arrived as native code returned to Rust. Nothing can raise + /// another one once the call is unrecorded, so let this one fire now rather than in later + /// Python code on this thread. Clearing it with `PyThreadState_SetAsyncExc(id, NULL)` + /// instead would leave the interpreter's pending-event flag set on CPython 3.11 and 3.12, + /// slowing down every other thread. + fn exit(self, py: Python<'_>) { + let interrupted = self.unregister(); + std::mem::forget(self); + if interrupted { + // The backward jump checks for pending events. The registry lock is released, so + // the GIL may switch here. + let _ = py.run(c"for _ in (0, 0):\n pass\n", None, None); + } + } } impl Drop for RunningCall<'_> { fn drop(&mut self) { - let mut tasks = tasks(); - let Some(calls) = tasks.get_mut(&self.task.task_attempt_id) else { - return; - }; - let Some(index) = calls.calls.iter().position(|c| c.thread == self.thread) else { - return; - }; - if calls.calls.swap_remove(index).interrupted { - // An interrupt that arrived as the callable returned is still pending on this - // thread. Nothing can set another one now that the call is unrecorded, so clear - // it rather than let it fire in later Python code on this thread. - // SAFETY: the call is made while attached, on the thread it names. - unsafe { - pyo3::ffi::PyThreadState_SetAsyncExc(self.thread as c_long, std::ptr::null_mut()) - }; - } + self.unregister(); } } -/// Marks a Spark task attempt as killed and raises [`TaskKilled`] in its running Arrow UDF -/// calls that it has not interrupted yet. Later calls of the task fail without entering -/// Python. Returns the number of calls interrupted. A task with no registered plan is not -/// recorded, so the caller repeats this until the task completes. +/// Marks a Spark task as killed, so that its later Arrow UDF calls fail without entering +/// Python. With `interrupt_running`, also raises [`TaskKilled`] in its running calls that it +/// has not interrupted yet. Returns the number of calls interrupted. A task with no +/// registered plan is not recorded, so the caller repeats this until the task completes. /// /// The exception takes effect when the thread next runs Python bytecode: a call blocked in -/// native code that released the GIL is interrupted only when that code returns. -pub fn interrupt_task(task_attempt_id: i64) -> usize { - let running = match tasks().get_mut(&task_attempt_id) { - Some(task) => { - task.killed = true; - task.calls.iter().any(|call| !call.interrupted) +/// native code that released the GIL is interrupted only when that code returns. It can also +/// take effect between a lock's acquisition and the `try` that releases it, leaving the lock +/// held, which is why the JVM can turn `interrupt_running` off. +pub fn kill_task(task: i64, interrupt_running: bool) -> usize { + let running = match tasks().get_mut(&task) { + Some(calls) => { + calls.killed = true; + interrupt_running && calls.calls.iter().any(|call| !call.interrupted) } None => false, }; @@ -182,11 +194,11 @@ pub fn interrupt_task(task_attempt_id: i64) -> usize { Python::attach(|py| { let exception = TaskKilled::type_object(py); let mut tasks = tasks(); - let Some(task) = tasks.get_mut(&task_attempt_id) else { + let Some(calls) = tasks.get_mut(&task) else { return 0; }; let mut interrupted = 0; - for call in task.calls.iter_mut().filter(|call| !call.interrupted) { + for call in calls.calls.iter_mut().filter(|call| !call.interrupted) { // SAFETY: attached to the interpreter, and `exception` is a live type object. let set = unsafe { pyo3::ffi::PyThreadState_SetAsyncExc(call.thread as c_long, exception.as_ptr()) @@ -403,7 +415,7 @@ impl ArrowPythonUdf { }) } - /// Lets [`interrupt_task`] interrupt this UDF's calls for the registered task. + /// Lets [`kill_task`] stop this UDF's calls for the registered task. pub fn with_task(mut self, task: Arc) -> Self { self.task = Some(task); self @@ -437,6 +449,10 @@ impl ArrowPythonUdf { } } + // Fail before waiting for the GIL, which another UDF may hold for a long time. + if let Some(task) = self.task.as_deref().filter(|task| task.is_killed()) { + return Err(task.killed_error()); + } Python::attach(|py| { let pa = py.import("pyarrow").map_err(python_error)?; let array_class = pa.getattr("Array").map_err(python_error)?; @@ -474,30 +490,32 @@ impl ArrowPythonUdf { .as_deref() .map(|task| RunningCall::enter(py, task)) .transpose()?; - let result = self.callable.bind(py).call(positional, Some(&kwargs)); - drop(call); + let mut result = self.callable.bind(py).call(positional, Some(&kwargs)); + if self.accept_array_like { + // Spark 4.2 builds its output batch with `RecordBatch.from_arrays`, which + // converts lists and NumPy arrays and rejects chunked arrays. It can run a + // user-defined `__arrow_array__`, so the call stays recorded meanwhile. + result = result.and_then(|result| { + let arrays = PyList::new(py, [result])?; + let names = PyList::new(py, ["_0"])?; + pa.getattr("RecordBatch")? + .call_method1("from_arrays", (arrays, names))? + .call_method1("column", (0,)) + }); + } + if let Some(call) = call { + call.exit(py); + } let result = result.map_err(|error| python_traceback_error(py, error))?; if let Some(task) = self.task.as_deref().filter(|task| task.is_killed()) { return Err(task.killed_error()); } - let result = if self.accept_array_like { - // Spark 4.2 builds its output batch with `RecordBatch.from_arrays`, - // which converts lists and NumPy arrays and rejects chunked arrays. - let arrays = PyList::new(py, [result]).map_err(python_error)?; - let names = PyList::new(py, ["_0"]).map_err(python_error)?; - pa.getattr("RecordBatch") - .map_err(python_error)? - .call_method1("from_arrays", (arrays, names)) - .map_err(|error| python_traceback_error(py, error))? - .call_method1("column", (0,)) - .map_err(python_error)? - } else if result.is_instance(&array_class).map_err(python_error)? { - result - } else { + // The conversion above already produced an array for Spark 4.2. + if !self.accept_array_like && !result.is_instance(&array_class).map_err(python_error)? { return Err(ArrowError::ComputeError( "Arrow UDF must return a pyarrow.Array".to_string(), )); - }; + } let result_len = result.len().map_err(python_error)?; if result_len != num_rows { return Err(ArrowError::ComputeError(format!( @@ -775,51 +793,98 @@ mod tests { }); } - fn running_calls(task_attempt_id: i64) -> Option { - tasks().get(&task_attempt_id).map(|task| task.calls.len()) + fn running_calls(task: i64) -> Option { + tasks().get(&task).map(|task| task.calls.len()) } - #[test] - fn interrupts_running_call_of_killed_task() { - const TASK: i64 = -6811; - initialize_python().unwrap(); - // `except Exception` must not swallow the interrupt. - let command = Python::attach(|py| { - cloudpickled_lambda( - py, - c"(lambda ns: (exec('def f(a):\\n n = 0\\n while True:\\n n += 1\\n try:\\n n %= 7\\n except Exception:\\n pass\\n', ns), ns['f'])[1])({})", - ) - }); - let task = Arc::new(TaskRegistration::new(TASK)); + /// Evaluates `python` (a function of one argument) as a UDF of `task` on another thread. + /// Returns the UDF, and a receiver for its result once the call has started. + fn start_call( + task: &Arc, + python: &std::ffi::CStr, + accept_array_like: bool, + ) -> ( + Arc, + std::sync::mpsc::Receiver>, + ) { + let command = Python::attach(|py| cloudpickled_lambda(py, python)); let udf = Arc::new( - ArrowPythonUdf::from_command(&command, DataType::Int64, true, true, false, "") - .unwrap() - .with_task(Arc::clone(&task)), + ArrowPythonUdf::from_command( + &command, + DataType::Int64, + true, + true, + accept_array_like, + "", + ) + .unwrap() + .with_task(Arc::clone(task)), ); - let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); let (sender, receiver) = std::sync::mpsc::channel(); - { - let (udf, input) = (Arc::clone(&udf), Arc::clone(&input)); - // Not joined: a call that ignores the interrupt would never return. - std::thread::spawn(move || sender.send(udf.evaluate(&[input], 1))); - } + let worker = Arc::clone(&udf); + // Not joined: a call that ignores the interrupt would never return. + std::thread::spawn(move || { + let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let result = worker.evaluate(&[input], 1); + // Release the registration before the test can check that it is gone. + drop(worker); + sender.send(result) + }); + wait_for( + || running_calls(task.task) == Some(1), + "the UDF call did not start", + ); + (udf, receiver) + } + + fn wait_for(condition: impl Fn() -> bool, message: &str) { let start = std::time::Instant::now(); - while running_calls(TASK) != Some(1) { - assert!(start.elapsed().as_secs() < 30, "the UDF call did not start"); + while !condition() { + assert!(start.elapsed().as_secs() < 30, "{message}"); std::thread::sleep(std::time::Duration::from_millis(10)); } - assert_eq!(interrupt_task(TASK), 1); - let error = receiver + } + + fn error_of(receiver: std::sync::mpsc::Receiver>) -> String { + receiver .recv_timeout(std::time::Duration::from_secs(30)) - .expect("the interrupted call did not return") + .expect("the call did not return") .unwrap_err() - .to_string(); + .to_string() + } + + /// Whether the Python expression `expression`, which may use `sys`, is true. + fn python_is(expression: &std::ffi::CStr) -> bool { + Python::attach(|py| { + let locals = pyo3::types::PyDict::new(py); + locals.set_item("sys", py.import("sys").unwrap()).unwrap(); + py.eval(expression, None, Some(&locals)) + .unwrap() + .extract() + .unwrap() + }) + } + + #[test] + fn interrupts_running_call_of_killed_task() { + const TASK: i64 = -6811; + initialize_python().unwrap(); + let task = Arc::new(TaskRegistration::new(TASK)); + // `except Exception` must not swallow the interrupt. + let (udf, result) = start_call( + &task, + c"(lambda ns: (exec('def f(a):\\n n = 0\\n while True:\\n n += 1\\n try:\\n n %= 7\\n except Exception:\\n pass\\n', ns), ns['f'])[1])({})", + false, + ); + assert_eq!(kill_task(TASK, true), 1); + let error = error_of(result); assert!(error.contains("TaskKilled"), "{error}"); assert_eq!(running_calls(TASK), Some(0)); // A killed task does not enter Python again. + let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); let error = udf.evaluate(&[input], 1).unwrap_err().to_string(); assert!(error.contains("was killed"), "{error}"); - assert_eq!(interrupt_task(TASK), 0); + assert_eq!(kill_task(TASK, true), 0); drop(udf); drop(task); assert_eq!(running_calls(TASK), None); @@ -829,7 +894,7 @@ mod tests { fn interrupts_a_call_once_so_its_cleanup_runs() { const TASK: i64 = -6812; initialize_python().unwrap(); - let command = Python::attach(|py| { + Python::attach(|py| { let sys = py.import("sys").unwrap(); let lock = py .import("threading") @@ -838,59 +903,139 @@ mod tests { .unwrap(); sys.setattr("comet_cleanup_lock", lock).unwrap(); sys.setattr("comet_cleaning", false).unwrap(); - // The cleanup outlasts the monitor's 100 ms polling interval. - cloudpickled_lambda( - py, - c"(lambda ns: (exec('def f(a):\\n import sys, time\\n sys.comet_cleanup_lock.acquire()\\n try:\\n n = 0\\n while True:\\n n += 1\\n finally:\\n sys.comet_cleaning = True\\n time.sleep(0.4)\\n sys.comet_cleanup_lock.release()\\n', ns), ns['f'])[1])({})", - ) }); - let sys_attr = |expression: &std::ffi::CStr| -> bool { - Python::attach(|py| { - let sys = py.import("sys").unwrap(); - let locals = pyo3::types::PyDict::new(py); - locals.set_item("sys", sys).unwrap(); - py.eval(expression, None, Some(&locals)) - .unwrap() - .extract() - .unwrap() - }) - }; let task = Arc::new(TaskRegistration::new(TASK)); - let udf = Arc::new( - ArrowPythonUdf::from_command(&command, DataType::Int64, true, true, false, "") - .unwrap() - .with_task(Arc::clone(&task)), + // The cleanup outlasts the monitor's 100 ms polling interval. + let (_udf, result) = start_call( + &task, + c"(lambda ns: (exec('def f(a):\\n import sys, time\\n sys.comet_cleanup_lock.acquire()\\n try:\\n n = 0\\n while True:\\n n += 1\\n finally:\\n sys.comet_cleaning = True\\n time.sleep(0.4)\\n sys.comet_cleanup_lock.release()\\n', ns), ns['f'])[1])({})", + false, + ); + assert_eq!(kill_task(TASK, true), 1); + wait_for( + || python_is(c"sys.comet_cleaning"), + "the cleanup did not start", ); - let (sender, receiver) = std::sync::mpsc::channel(); - { - let udf = Arc::clone(&udf); - let input: ArrayRef = Arc::new(Int64Array::from(vec![1])); - std::thread::spawn(move || sender.send(udf.evaluate(&[input], 1))); - } - let start = std::time::Instant::now(); - while running_calls(TASK) != Some(1) { - assert!(start.elapsed().as_secs() < 30, "the UDF call did not start"); - std::thread::sleep(std::time::Duration::from_millis(10)); - } - assert_eq!(interrupt_task(TASK), 1); - while !sys_attr(c"sys.comet_cleaning") { - assert!(start.elapsed().as_secs() < 30, "the cleanup did not start"); - std::thread::sleep(std::time::Duration::from_millis(10)); - } // A later poll while the call cleans up must not raise in it again. - assert_eq!(interrupt_task(TASK), 0); - let error = receiver - .recv_timeout(std::time::Duration::from_secs(30)) - .expect("the interrupted call did not return") - .unwrap_err() - .to_string(); + assert_eq!(kill_task(TASK, true), 0); + let error = error_of(result); assert!(error.contains("TaskKilled"), "{error}"); assert!( - !sys_attr(c"sys.comet_cleanup_lock.locked()"), + !python_is(c"sys.comet_cleanup_lock.locked()"), "the cleanup did not release its lock" ); } + #[test] + fn kill_without_interrupt_lets_the_running_call_finish() { + const TASK: i64 = -6813; + initialize_python().unwrap(); + Python::attach(|py| { + py.import("sys") + .unwrap() + .setattr("comet_finish_call", false) + .unwrap() + }); + let task = Arc::new(TaskRegistration::new(TASK)); + let (_udf, result) = start_call( + &task, + c"(lambda ns: (exec('def f(a):\\n import sys, time\\n while not sys.comet_finish_call:\\n time.sleep(0.01)\\n return a\\n', ns), ns['f'])[1])({})", + false, + ); + assert_eq!(kill_task(TASK, false), 0); + Python::attach(|py| { + py.import("sys") + .unwrap() + .setattr("comet_finish_call", true) + .unwrap() + }); + // The call ran to its end, and its result is dropped because the task was killed. + let error = error_of(result); + assert!(error.contains("was killed"), "{error}"); + assert!(!error.contains("TaskKilled"), "{error}"); + } + + #[test] + fn interrupts_a_user_defined_array_conversion() { + const TASK: i64 = -6814; + initialize_python().unwrap(); + Python::attach(|py| { + py.import("sys") + .unwrap() + .setattr("comet_converting", false) + .unwrap() + }); + let task = Arc::new(TaskRegistration::new(TASK)); + // Spark 4.2 converts the result with `RecordBatch.from_arrays`, which calls + // `__arrow_array__`. + let (_udf, result) = start_call( + &task, + c"(lambda ns: (exec('import sys\\nclass Slow:\\n def __len__(self):\\n return 1\\n def __arrow_array__(self, type=None):\\n sys.comet_converting = True\\n n = 0\\n while True:\\n n += 1\\ndef f(a):\\n return Slow()\\n', ns), ns['f'])[1])({})", + true, + ); + wait_for( + || python_is(c"sys.comet_converting"), + "the conversion did not start", + ); + assert_eq!(kill_task(TASK, true), 1); + let error = error_of(result); + assert!(error.contains("TaskKilled"), "{error}"); + } + + #[test] + fn consumes_an_interrupt_pending_when_the_call_returns() { + const TASK: i64 = -6815; + initialize_python().unwrap(); + // A callable implemented in C returns straight to Rust, so an interrupt that arrives + // while it sleeps is still pending when the call ends. + let command = Python::attach(|py| { + let sleep = py + .import("functools") + .unwrap() + .getattr("partial") + .unwrap() + .call1((py.import("time").unwrap().getattr("sleep").unwrap(), 0.3)) + .unwrap(); + let _imports = IMPORT_LOCK + .lock_py_attached(py) + .unwrap_or_else(PoisonError::into_inner); + py.import("cloudpickle") + .unwrap() + .call_method1("dumps", ((sleep, py.None()),)) + .unwrap() + .extract::>() + .unwrap() + }); + let task = Arc::new(TaskRegistration::new(TASK)); + let udf = ArrowPythonUdf::from_command(&command, DataType::Int64, true, true, false, "") + .unwrap() + .with_task(Arc::clone(&task)); + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::scope(|scope| { + scope.spawn(|| { + // Keep this thread's Python state across the call, as a thread that runs + // more Python afterwards would. + Python::attach(|py| { + let result = udf.evaluate(&[], 0).map(drop); + let after = py.run(c"for _ in (0, 0):\n pass\n", None, None); + sender.send((result, after.is_ok())).unwrap(); + }) + }); + wait_for( + || running_calls(TASK) == Some(1), + "the UDF call did not start", + ); + assert_eq!(kill_task(TASK, true), 1); + }); + let (result, later_python_ok) = receiver.recv().unwrap(); + let error = result.unwrap_err().to_string(); + assert!(error.contains("was killed"), "{error}"); + assert!( + later_python_ok, + "the interrupt fired in Python code after the call" + ); + } + const UNBUFFERED_CHILD: &str = "COMET_ARROW_UDF_UNBUFFERED_CHILD"; const UNBUFFERED_MARKER: &str = "comet-arrow-udf-unbuffered-marker"; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index edffd09389b..1ad391757a3 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -190,19 +190,20 @@ pub extern "system" fn Java_org_apache_comet_NativeBase_nativeSupportsPythonUdf( cfg!(feature = "python-udf") } -/// Interrupts the Python UDF calls of a killed Spark task; see -/// `execution::python_udf::interrupt_task`. Returns the number of calls interrupted. +/// Stops the Python UDF calls of a killed Spark task; see `execution::python_udf::kill_task`. +/// Returns the number of calls interrupted. #[no_mangle] -pub extern "system" fn Java_org_apache_comet_NativeBase_interruptPythonUdfs( +pub extern "system" fn Java_org_apache_comet_NativeBase_killPythonUdfs( _: EnvUnowned, _: JClass, - task_attempt_id: jni::sys::jlong, + task: jni::sys::jlong, + interrupt_running: jni::sys::jboolean, ) -> jni::sys::jint { #[cfg(feature = "python-udf")] - return execution::python_udf::interrupt_task(task_attempt_id) as jni::sys::jint; + return execution::python_udf::kill_task(task, interrupt_running) as jni::sys::jint; #[cfg(not(feature = "python-udf"))] { - let _ = task_attempt_id; + let _ = (task, interrupt_running); 0 } } diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index 315e5b03e5f..ccb91384a11 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -330,13 +330,15 @@ public static boolean supportsPythonUdf() { } /** - * Raises an exception in the running Python UDF calls of a killed task, and makes its later calls - * fail without entering Python. A no-op for a task with no native plan holding a Python UDF, so - * callers repeat it until the task completes. + * Makes the later Python UDF calls of a killed task fail without entering Python and, with + * `interruptRunning`, raises an exception in its running calls that it has not interrupted yet. A + * no-op for a task with no native plan holding a Python UDF, so callers repeat it until the task + * completes. * + * @param task the token that `Native.setPythonUdfTask` bound the task's plans to * @return the number of calls interrupted */ - public static native int interruptPythonUdfs(long taskAttemptId); + public static native int killPythonUdfs(long task, boolean interruptRunning); /** Release native resources through JNI */ static native void release(); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 24f38826463..f10611cf26c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -541,6 +541,19 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_NATIVE_ARROW_PYTHON_UDF_INTERRUPT_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.nativeArrowPythonUDF.interruptOnKill.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether killing a task raises an exception in its running native Arrow UDF call, " + + "once the call has run for spark.python.task.killTimeout after the kill. The " + + "exception can arrive while Python code holds a lock it would release in a later " + + "`finally` block, leaving the lock held for every task in the executor's shared " + + "interpreter. When disabled, a killed task waits for its running call to return. " + + "Either way, the task's later calls do not run.") + .booleanConf + .createWithDefault(true) + val COMET_TRACING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.tracing.enabled") .category(CATEGORY_TUNING) .doc(s"Enable fine-grained tracing of events and memory usage. $TRACING_GUIDE.") diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index f51bf5aedb3..8b5a50f75e4 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -111,7 +111,7 @@ class CometExecIterator( private val taskMemory = CometExecIterator.taskMemory(TaskContext.get(), id) private val cometTaskMemoryManager = taskMemory.manager - private val plan = { + private val (plan, pythonUdfWatch) = { val conf = SparkEnv.get.conf val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs @@ -164,7 +164,20 @@ class CometExecIterator( shufflePartitionPusher.foreach { pusher => nativeLib.setShufflePartitionPusher(createdPlan, pusher) } - createdPlan + // A native Python UDF ignores the task's kill, so watch the task to stop its calls. This + // reads the plan before the completion listener below, which can release it at once. + val pythonUdfWatch = + if (PythonUdfMonitor.supported && nativeLib.hasPythonUdf(createdPlan)) { + val watch = PythonUdfMonitor.watch( + TaskContext.get(), + COMET_NATIVE_ARROW_PYTHON_UDF_INTERRUPT_ENABLED.get()) + nativeLib.setPythonUdfTask(createdPlan, watch.task) + watch.planOpened() + Some(watch) + } else { + None + } + (createdPlan, pythonUdfWatch) } catch { case failure: Throwable => // The task-completion listener is not installed until iterator construction succeeds. @@ -235,11 +248,6 @@ class CometExecIterator( CometExecIterator.startMemoryUsageLog() - // A native Python UDF ignores the task's kill, so watch the task to interrupt its calls. - if (PythonUdfMonitor.supported && nativeLib.hasPythonUdf(plan)) { - PythonUdfMonitor.watch(TaskContext.get()) - } - /** * Matches any throwable once a JVM input of this plan has failed, and extracts what that input * threw. Native saw only its text. See `CometArrowStream.inputFailure`. @@ -373,6 +381,9 @@ class CometExecIterator( // execution context frees this plan's task-shared memory pool reference and several JNI // global refs. attempt(nativeLib.releasePlan(plan)) + // Only now: releasing the plan waits for a Python UDF call in flight, which a kill + // meanwhile must still stop. + pythonUdfWatch.foreach(_.planClosed()) // Drops what the task's JVM UDFs keep for this plan, such as its dispatched kernels. `id` is // the plan id the native plan passed them. diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index e58f150664f..7ef2cbe8a56 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -138,6 +138,17 @@ class Native extends NativeBase { */ @native def hasPythonUdf(plan: Long): Boolean + /** + * Binds the native query plan's Python UDFs to the token of its Spark task, which + * `NativeBase.killPythonUdfs` takes. Must precede the first `executePlan`. + * + * @param plan + * the address to native query plan. + * @param task + * the token of the plan's task, unique in this process. + */ + @native def setPythonUdfTask(plan: Long, task: Long): Unit + /** * Release and drop the native query plan object and context object. * diff --git a/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala index 5bd32f80f4b..a98d060baee 100644 --- a/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala +++ b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala @@ -19,52 +19,93 @@ package org.apache.comet -import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.{ConcurrentHashMap, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} -import org.apache.spark.TaskContext +import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.internal.Logging /** - * Interrupts the native Python UDF calls of killed tasks. Spark's Python runner has a - * `MonitorThread` that destroys the worker of a killed task. A native UDF runs in the embedded - * interpreter on the task thread instead, where neither a Java interrupt nor a stopped input - * reaches it. So one daemon thread per task polls its kill state and, while the task is killed - * but still running, raises an exception in the task's running Python calls, once per call so - * that its cleanup can run. Calls that start after the first interrupt fail without entering - * Python. + * Stops the native Python UDF calls of killed tasks. Spark's Python runner has a `MonitorThread` + * that destroys the worker of a killed task. A native UDF runs in the embedded interpreter on the + * task thread instead, where neither a Java interrupt nor a stopped input reaches it. So one + * daemon thread per task polls its kill state. Once the task is killed, its later calls fail + * without entering Python. A call still running `spark.python.task.killTimeout` after the kill + * gets an exception raised in it, once so that its cleanup can run, unless that is disabled. */ object PythonUdfMonitor extends Logging { - // Raising the exception is cheap and lets the call unwind on its own, so unlike Spark's - // monitor this one does not wait `spark.python.task.killTimeout` before acting. Repeating it - // reaches calls of plans that the task registers after an earlier attempt. + // Repeating the kill reaches calls of plans that the task registers after an earlier poll. private val PollIntervalMs = 100L - private val watchedTasks = ConcurrentHashMap.newKeySet[java.lang.Long]() + private val nextTask = new AtomicLong() - /** Whether the native library can run Python UDFs, and so may need its calls interrupted. */ + private val watches = new ConcurrentHashMap[TaskContext, Watch]() + + /** Whether the native library can run Python UDFs, and so may need its calls stopped. */ lazy val supported: Boolean = NativeBase.supportsPythonUdf() - /** Starts watching the task unless it is watched already. The watch ends with the task. */ - def watch(context: TaskContext): Unit = { + /** + * The watch over one task. `task` is unique in this process, while a task attempt id is unique + * only within one SparkContext. + */ + final class Watch private[PythonUdfMonitor] (val task: Long) { + private val openPlans = new AtomicInteger() + + /** Counts a native plan bound to `task`, which the watch outlives. */ + def planOpened(): Unit = { + val _ = openPlans.incrementAndGet() + } + + /** Ends what [[planOpened]] started, once the plan has been released. */ + def planClosed(): Unit = { + val _ = openPlans.decrementAndGet() + } + + private[PythonUdfMonitor] def hasOpenPlans: Boolean = openPlans.get() > 0 + } + + /** + * Returns the task's watch, starting it unless the task is watched already. The watch ends once + * the task has completed and its counted plans are released, since releasing a plan waits for + * its running call. + */ + def watch(context: TaskContext, interruptRunning: Boolean): Watch = + watches.computeIfAbsent( + context, + _ => { + val watch = new Watch(nextTask.incrementAndGet()) + start(context, watch, interruptRunning) + watch + }) + + private def start(context: TaskContext, watch: Watch, interruptRunning: Boolean): Unit = { + val killTimeoutNs = TimeUnit.MILLISECONDS.toNanos( + SparkEnv.get.conf.getTimeAsMs("spark.python.task.killTimeout", "2s")) val taskAttemptId = context.taskAttemptId() - if (watchedTasks.add(taskAttemptId)) { - val thread = new Thread(s"Comet Python UDF monitor for task $taskAttemptId") { - override def run(): Unit = { - try { - while (!context.isCompleted()) { - if (context.isInterrupted() && NativeBase.interruptPythonUdfs(taskAttemptId) > 0) { + val thread = new Thread(s"Comet Python UDF monitor for task $taskAttemptId") { + override def run(): Unit = { + try { + var killedAt = Option.empty[Long] + while (!context.isCompleted() || watch.hasOpenPlans) { + if (context.isInterrupted()) { + val now = System.nanoTime() + val since = killedAt.getOrElse(now) + killedAt = Some(since) + // Like Spark's monitor, give a running call the kill timeout to return. + val interrupt = interruptRunning && now - since >= killTimeoutNs + if (NativeBase.killPythonUdfs(watch.task, interrupt) > 0) { logInfo(s"Interrupted Python UDF calls of killed task $taskAttemptId") } - Thread.sleep(PollIntervalMs) } - } finally { - val _ = watchedTasks.remove(taskAttemptId) + Thread.sleep(PollIntervalMs) } + } finally { + val _ = watches.remove(context, watch) } } - thread.setDaemon(true) - thread.start() } + thread.setDaemon(true) + thread.start() } } diff --git a/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala index dbe6a8f039d..ddd247e5605 100644 --- a/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala +++ b/spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala @@ -535,52 +535,97 @@ class CometArrowPythonUdfSuite extends CometTestBase { } } + /** + * Runs `job` in a job group, waits until `started` exists, cancels the group, and checks that + * the task ends as killed. + */ + private def assertCancelKillsTask(started: File)(job: => Any): Unit = { + val killed = new CountDownLatch(1) + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEnd.reason match { + case _: TaskKilled => killed.countDown() + case _ => + } + } + spark.sparkContext.addSparkListener(listener) + try { + val group = "comet-arrow-udf-kill" + val running = Future { + spark.sparkContext.setJobGroup(group, "native Arrow UDF loop", interruptOnCancel = true) + job + }(ExecutionContext.global) + val deadline = 60.seconds.fromNow + while (!started.exists()) { + assert(deadline.hasTimeLeft(), "the UDF did not start") + Thread.sleep(50) + } + spark.sparkContext.cancelJobGroup(group) + // The monitor waits spark.python.task.killTimeout, 2 seconds, before interrupting. + assert(killed.await(30, TimeUnit.SECONDS), "the task running the UDF was not killed") + intercept[SparkException](ThreadUtils.awaitResult(running, 30.seconds)) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + /** + * A UDF that loops in Python bytecode, where the interrupt takes effect, and survives a user + * `except Exception`. It creates `started` first, for its batches that pass `condition`. + */ + private def loopingUdf(started: File, condition: String = "True") = { + val loop = "(lambda ns: (exec('def f(a):\\n" + + s" if not ($condition):\\n return a\\n" + + " open(\"" + started.getCanonicalPath + "\", \"w\").close()\\n" + + " n = 0\\n while True:\\n n += 1\\n try:\\n" + + " n %= 7\\n except Exception:\\n pass\\n', ns), " + + "ns['f'])[1])({})" + arrowUdf("loop", pickledCommand(loop, "LongType()"), LongType) + } + test("killing a task interrupts its running native Arrow UDF") { assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") withTempDir { dir => val started = new File(dir, "started") - // Loops in Python bytecode, where the interrupt takes effect, and survives a user - // `except Exception`. - val loop = "(lambda ns: (exec('def f(a):\\n" + - " open(\"" + started.getCanonicalPath + "\", \"w\").close()\\n" + - " n = 0\\n while True:\\n n += 1\\n try:\\n" + - " n %= 7\\n except Exception:\\n pass\\n', ns), " + - "ns['f'])[1])({})" - val udf = arrowUdf("loop", pickledCommand(loop, "LongType()"), LongType) - val killed = new CountDownLatch(1) - val listener = new SparkListener { - override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEnd.reason match { - case _: TaskKilled => killed.countDown() - case _ => - } + val udf = loopingUdf(started) + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = spark.range(0, 1, 1, 1).select(udf(col("id"))) + assert(isNative(df)) + assertCancelKillsTask(started)(df.collect()) } - spark.sparkContext.addSparkListener(listener) - try { - withSQLConf( - CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - val df = spark.range(0, 1, 1, 1).select(udf(col("id"))) - assert(isNative(df)) - val group = "comet-arrow-udf-kill" - val job = Future { - spark.sparkContext.setJobGroup( - group, - "native Arrow UDF loop", - interruptOnCancel = true) - df.collect() - }(ExecutionContext.global) - val deadline = 60.seconds.fromNow - while (!started.exists()) { - assert(deadline.hasTimeLeft(), "the UDF did not start") - Thread.sleep(50) - } - spark.sparkContext.cancelJobGroup(group) - assert(killed.await(30, TimeUnit.SECONDS), "the task running the UDF was not killed") - intercept[SparkException](ThreadUtils.awaitResult(job, 30.seconds)) + } + } + + test("killing a completed task interrupts the UDF call its plan release waits for") { + assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") + + withTempDir { dir => + val path = new File(dir, "data").getCanonicalPath + spark.range(0, 2, 1, 1).write.parquet(path) + val started = new File(dir, "started") + // Loops only on the second row, which a native scan's plan evaluates ahead of the task. + val udf = loopingUdf(started, "a[0].as_py() == 1") + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "1") { + val df = spark.read.parquet(path).select(udf(col("id"))) + assert(isNative(df)) + val marker = started.getCanonicalPath + assertCancelKillsTask(started) { + // Takes the first row, then completes the task once the second is in the UDF. Releasing + // the plan then waits for that call. + df.rdd + .mapPartitions { rows => + assert(rows.hasNext) + val first = rows.next().getLong(0) + while (!new File(marker).exists()) Thread.sleep(10) + Iterator(first) + } + .collect() } - } finally { - spark.sparkContext.removeSparkListener(listener) } } } From 69697372fcbfb1c53a0a69485a561c729d7d85e8 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 27 Oct 2026 10:06:13 -0700 Subject: [PATCH 4/4] fix: keep a Python UDF call recorded until its result is exported After the callable returns, checking, casting and exporting its result can run more user code, such as the `cast` of an extension array subclass. The call was unrecorded before that, so a kill could not interrupt it and a successful batch came back. The call now stays recorded until its result is exported, and the kill check follows. --- native/core/src/execution/python_udf.rs | 181 +++++++++++++++--------- 1 file changed, 113 insertions(+), 68 deletions(-) diff --git a/native/core/src/execution/python_udf.rs b/native/core/src/execution/python_udf.rs index b18395af95c..35e89cee382 100644 --- a/native/core/src/execution/python_udf.rs +++ b/native/core/src/execution/python_udf.rs @@ -490,78 +490,18 @@ impl ArrowPythonUdf { .as_deref() .map(|task| RunningCall::enter(py, task)) .transpose()?; - let mut result = self.callable.bind(py).call(positional, Some(&kwargs)); - if self.accept_array_like { - // Spark 4.2 builds its output batch with `RecordBatch.from_arrays`, which - // converts lists and NumPy arrays and rejects chunked arrays. It can run a - // user-defined `__arrow_array__`, so the call stays recorded meanwhile. - result = result.and_then(|result| { - let arrays = PyList::new(py, [result])?; - let names = PyList::new(py, ["_0"])?; - pa.getattr("RecordBatch")? - .call_method1("from_arrays", (arrays, names))? - .call_method1("column", (0,)) - }); - } + // Everything up to the export can run user code, such as a result's + // `__arrow_array__` or an extension array's `cast`, so the call stays recorded + // until it ends. + let exported = + self.call_and_export(py, &pa, &array_class, &positional, &kwargs, num_rows); if let Some(call) = call { call.exit(py); } - let result = result.map_err(|error| python_traceback_error(py, error))?; + let (out_array, out_schema) = exported?; if let Some(task) = self.task.as_deref().filter(|task| task.is_killed()) { return Err(task.killed_error()); } - // The conversion above already produced an array for Spark 4.2. - if !self.accept_array_like && !result.is_instance(&array_class).map_err(python_error)? { - return Err(ArrowError::ComputeError( - "Arrow UDF must return a pyarrow.Array".to_string(), - )); - } - let result_len = result.len().map_err(python_error)?; - if result_len != num_rows { - return Err(ArrowError::ComputeError(format!( - "Arrow UDF returned {result_len} rows, expected {num_rows}" - ))); - } - - let mut ffi_return_type = FFI_ArrowSchema::try_from(&self.return_type)?; - let expected_type = pa - .getattr("DataType") - .map_err(python_error)? - .call_method1( - "_import_from_c", - (&raw mut ffi_return_type as Py_uintptr_t,), - ) - .map_err(python_error)?; - let actual_type = result.getattr("type").map_err(python_error)?; - let typed_result = if actual_type.eq(&expected_type).map_err(python_error)? { - result - } else if self.allow_cast { - let kwargs = pyo3::types::PyDict::new(py); - kwargs - .set_item("safe", self.safe_cast) - .map_err(python_error)?; - result - .call_method("cast", (expected_type,), Some(&kwargs)) - .map_err(python_error)? - } else { - return Err(ArrowError::ComputeError(format!( - "Arrow UDF returned type {}, expected {}", - actual_type.str().map_err(python_error)?, - expected_type.str().map_err(python_error)? - ))); - }; - - let mut out_array = FFI_ArrowArray::empty(); - let mut out_schema = FFI_ArrowSchema::empty(); - typed_result - .call_method1( - "_export_to_c", - ( - &raw mut out_array as Py_uintptr_t, - &raw mut out_schema as Py_uintptr_t, - ), - ) - .map_err(python_error)?; // SAFETY: PyArrow filled both C Data structs and transferred ownership // of the array to `out_array`. `from_ffi` checks the schema but builds the // array without validating its buffers, so string data is checked below. @@ -580,6 +520,87 @@ impl ArrowPythonUdf { } } +impl ArrowPythonUdf { + /// Calls the callable, checks and casts its result, and exports it through the Arrow C Data + /// interface. + fn call_and_export( + &self, + py: Python<'_>, + pa: &Bound<'_, PyModule>, + array_class: &Bound<'_, PyAny>, + positional: &Bound<'_, PyTuple>, + kwargs: &Bound<'_, pyo3::types::PyDict>, + num_rows: usize, + ) -> Result<(FFI_ArrowArray, FFI_ArrowSchema)> { + let mut result = self.callable.bind(py).call(positional, Some(kwargs)); + if self.accept_array_like { + // Spark 4.2 builds its output batch with `RecordBatch.from_arrays`, which converts + // lists and NumPy arrays and rejects chunked arrays. + result = result.and_then(|result| { + let arrays = PyList::new(py, [result])?; + let names = PyList::new(py, ["_0"])?; + pa.getattr("RecordBatch")? + .call_method1("from_arrays", (arrays, names))? + .call_method1("column", (0,)) + }); + } + let result = result.map_err(|error| python_traceback_error(py, error))?; + // The conversion above already produced an array for Spark 4.2. + if !self.accept_array_like && !result.is_instance(array_class).map_err(python_error)? { + return Err(ArrowError::ComputeError( + "Arrow UDF must return a pyarrow.Array".to_string(), + )); + } + let result_len = result.len().map_err(python_error)?; + if result_len != num_rows { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF returned {result_len} rows, expected {num_rows}" + ))); + } + + let mut ffi_return_type = FFI_ArrowSchema::try_from(&self.return_type)?; + let expected_type = pa + .getattr("DataType") + .map_err(python_error)? + .call_method1( + "_import_from_c", + (&raw mut ffi_return_type as Py_uintptr_t,), + ) + .map_err(python_error)?; + let actual_type = result.getattr("type").map_err(python_error)?; + let typed_result = if actual_type.eq(&expected_type).map_err(python_error)? { + result + } else if self.allow_cast { + let kwargs = pyo3::types::PyDict::new(py); + kwargs + .set_item("safe", self.safe_cast) + .map_err(python_error)?; + result + .call_method("cast", (expected_type,), Some(&kwargs)) + .map_err(python_error)? + } else { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF returned type {}, expected {}", + actual_type.str().map_err(python_error)?, + expected_type.str().map_err(python_error)? + ))); + }; + + let mut out_array = FFI_ArrowArray::empty(); + let mut out_schema = FFI_ArrowSchema::empty(); + typed_result + .call_method1( + "_export_to_c", + ( + &raw mut out_array as Py_uintptr_t, + &raw mut out_schema as Py_uintptr_t, + ), + ) + .map_err(python_error)?; + Ok((out_array, out_schema)) + } +} + fn python_error(error: impl std::fmt::Display) -> ArrowError { ArrowError::ComputeError(format!("Arrow UDF Python error: {error}")) } @@ -982,6 +1003,30 @@ mod tests { assert!(error.contains("TaskKilled"), "{error}"); } + #[test] + fn interrupts_a_user_defined_result_cast() { + const TASK: i64 = -6816; + initialize_python().unwrap(); + Python::attach(|py| { + py.import("sys") + .unwrap() + .setattr("comet_casting", false) + .unwrap() + }); + let task = Arc::new(TaskRegistration::new(TASK)); + // The result's extension type differs from the declared Int64, so it is cast with the + // array's own `cast`. + let (_udf, result) = start_call( + &task, + c"(lambda ns: (exec('import sys\\nimport pyarrow as pa\\nclass LoopArray(pa.ExtensionArray):\\n def cast(self, *args, **kwargs):\\n sys.comet_casting = True\\n n = 0\\n while True:\\n n += 1\\nclass LoopType(pa.ExtensionType):\\n def __init__(self):\\n super().__init__(pa.int32(), \"comet.loop\")\\n def __arrow_ext_serialize__(self):\\n return b\"\"\\n @classmethod\\n def __arrow_ext_deserialize__(cls, storage_type, serialized):\\n return cls()\\n def __arrow_ext_class__(self):\\n return LoopArray\\ndef f(a):\\n return pa.ExtensionArray.from_storage(LoopType(), a.cast(pa.int32()))\\n', ns), ns['f'])[1])({})", + false, + ); + wait_for(|| python_is(c"sys.comet_casting"), "the cast did not start"); + assert_eq!(kill_task(TASK, true), 1); + let error = error_of(result); + assert!(error.contains("TaskKilled"), "{error}"); + } + #[test] fn consumes_an_interrupt_pending_when_the_call_returns() { const TASK: i64 = -6815; @@ -1028,8 +1073,8 @@ mod tests { assert_eq!(kill_task(TASK, true), 1); }); let (result, later_python_ok) = receiver.recv().unwrap(); - let error = result.unwrap_err().to_string(); - assert!(error.contains("was killed"), "{error}"); + // The sleep returns None, so the call fails on its result. + assert!(result.is_err()); assert!( later_python_ok, "the interrupt fired in Python code after the call"