diff --git a/docs/source/user-guide/latest/pyarrow-udfs.md b/docs/source/user-guide/latest/pyarrow-udfs.md index 48848ecd6e3..0fcc18e27e1 100644 --- a/docs/source/user-guide/latest/pyarrow-udfs.md +++ b/docs/source/user-guide/latest/pyarrow-udfs.md @@ -165,14 +165,20 @@ 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 -([#6811](https://github.com/apache/datafusion-comet/issues/6811)). A failing UDF raises Spark's -`PythonException` with its Python traceback, which PySpark raises as -`pyspark.errors.PythonException`. Errors outside the UDF's own code, such as a result of the wrong -length or type, are still reported as Comet native errors. The operator reports the time spent in +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. 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 raises Spark's `PythonException` +with its Python traceback, which PySpark raises as `pyspark.errors.PythonException`. Errors outside +the UDF's own code, such as a result of the wrong length or type, are still reported as Comet native +errors. 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 494e7acdff9..880965f47fb 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -535,6 +535,9 @@ struct ExecutionContext { pub id: i64, /// The deserialized Spark plan pub spark_plan: Operator, + /// 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` @@ -884,6 +887,7 @@ fn create_execution_context( Ok(Box::new(ExecutionContext { id, spark_plan, + python_udf_task: None, partition_count, root_op: None, scans: vec![], @@ -1538,6 +1542,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_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()) @@ -1780,6 +1785,40 @@ fn get_execution_context<'a>(id: i64) -> CometResult<&'a mut ExecutionContext> { .ok_or_else(|| CometError::Internal("Native execution context handle is null".to_string())) } +/// Whether the plan evaluates a Python UDF in the embedded interpreter, which the JVM then +/// stops 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, + )) + }) +} + +/// 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 7f8d0522fb6..c7daf0e08ab 100644 --- a/native/core/src/execution/operators/arrow_python_udf.rs +++ b/native/core/src/execution/operators/arrow_python_udf.rs @@ -37,7 +37,7 @@ use datafusion_comet_common::unwrap_arrow_external_error; 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 { @@ -59,6 +59,9 @@ pub struct ArrowPythonUdfExec { max_records_per_batch: usize, max_bytes_per_batch: usize, accept_array_like_results: bool, + /// 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, @@ -71,6 +74,7 @@ impl ArrowPythonUdfExec { max_records_per_batch: i32, max_bytes_per_batch: i64, accept_array_like_results: bool, + python_udf_task: Option, ) -> Result { if specs.is_empty() { return exec_err!("ArrowPythonUdfExec requires at least one UDF"); @@ -107,6 +111,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, + python_udf_task, schema, cache, metrics: ExecutionPlanMetricsSet::new(), @@ -300,6 +305,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.python_udf_task, )?)) } @@ -309,6 +315,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 + .python_udf_task + .map(|task| Arc::new(TaskRegistration::new(task))); let workers: Vec<_> = self .specs .iter() @@ -321,6 +331,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); @@ -473,6 +487,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 899d8ff3320..07ed5e9ccbf 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -207,6 +207,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 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 { @@ -229,6 +232,7 @@ impl PhysicalPlanner { output_stream: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, + python_udf_task: None, } } @@ -307,6 +311,11 @@ impl PhysicalPlanner { self } + pub fn with_python_udf_task(mut self, python_udf_task: Option) -> Self { + self.python_udf_task = python_udf_task; + self + } + pub fn with_strict_timestamp_labels(mut self, strict: bool) -> Self { self.strict_timestamp_labels = strict; self @@ -1686,6 +1695,7 @@ impl PhysicalPlanner { udf.max_records_per_batch, udf.max_bytes_per_batch, udf.accept_array_like_results, + 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 f9aca4d2d9e..4cd218e3a06 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,178 @@ 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 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, + killed: bool, + /// The calls running now. + calls: Vec, +} + +struct Call { + /// Python thread id of the thread running the call. + thread: c_ulong, + /// 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, +} + +// 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, 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: i64, +} + +impl TaskRegistration { + pub fn new(task: i64) -> Self { + tasks().entry(task).or_default().plans += 1; + Self { task } + } + + fn is_killed(&self) -> bool { + tasks().get(&self.task).is_some_and(|task| task.killed) + } + + fn killed_error(&self) -> ArrowError { + 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) { + task.plans -= 1; + if task.plans == 0 { + tasks.remove(&self.task); + } + } + } +} + +/// 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, +} + +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).or_default(); + if calls.killed { + return Err(task.killed_error()); + } + calls.calls.push(Call { + thread, + interrupted: false, + }); + 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) { + self.unregister(); + } +} + +/// 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. 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, + }; + 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 mut tasks = tasks(); + let Some(calls) = tasks.get_mut(&task) else { + return 0; + }; + let mut interrupted = 0; + 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()) + }; + if set > 0 { + call.interrupted = true; + interrupted += 1; + } + } + interrupted + }) +} + fn initialize_python() -> Result<()> { use std::ffi::CStr; use std::sync::OnceLock; @@ -181,6 +356,7 @@ pub struct ArrowPythonUdf { allow_cast: bool, safe_cast: bool, accept_array_like: bool, + task: Option>, } impl ArrowPythonUdf { @@ -247,10 +423,17 @@ impl ArrowPythonUdf { allow_cast, safe_cast, accept_array_like, + task: None, }) }) } + /// 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 + } + /// 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. @@ -279,6 +462,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)?; @@ -310,78 +497,24 @@ 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 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 { - 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 positional = PyTuple::new(py, positional).map_err(python_error)?; + let call = self + .task + .as_deref() + .map(|task| RunningCall::enter(py, task)) + .transpose()?; + // 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 (out_array, out_schema) = exported?; + if let Some(task) = self.task.as_deref().filter(|task| task.is_killed()) { + return Err(task.killed_error()); } - - 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. @@ -400,6 +533,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}")) } @@ -620,6 +834,273 @@ mod tests { }); } + fn running_calls(task: i64) -> Option { + tasks().get(&task).map(|task| task.calls.len()) + } + + /// 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, + accept_array_like, + "", + ) + .unwrap() + .with_task(Arc::clone(task)), + ); + let (sender, receiver) = std::sync::mpsc::channel(); + 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 !condition() { + assert!(start.elapsed().as_secs() < 30, "{message}"); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + } + + fn error_of(receiver: std::sync::mpsc::Receiver>) -> String { + receiver + .recv_timeout(std::time::Duration::from_secs(30)) + .expect("the call did not return") + .unwrap_err() + .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!(kill_task(TASK, true), 0); + drop(udf); + drop(task); + assert_eq!(running_calls(TASK), None); + } + + #[test] + fn interrupts_a_call_once_so_its_cleanup_runs() { + const TASK: i64 = -6812; + initialize_python().unwrap(); + 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(); + }); + let task = Arc::new(TaskRegistration::new(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", + ); + // A later poll while the call cleans up must not raise in it again. + assert_eq!(kill_task(TASK, true), 0); + let error = error_of(result); + assert!(error.contains("TaskKilled"), "{error}"); + assert!( + !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 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; + 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(); + // 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" + ); + } + 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 4f468409466..7527d8f060a 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -191,6 +191,24 @@ pub extern "system" fn Java_org_apache_comet_NativeBase_nativeSupportsPythonUdf( cfg!(feature = "python-udf") } +/// 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_killPythonUdfs( + _: EnvUnowned, + _: JClass, + task: jni::sys::jlong, + interrupt_running: jni::sys::jboolean, +) -> jni::sys::jint { + #[cfg(feature = "python-udf")] + return execution::python_udf::kill_task(task, interrupt_running) as jni::sys::jint; + #[cfg(not(feature = "python-udf"))] + { + let _ = (task, interrupt_running); + 0 + } +} + /// JNI: the `major.minor` version of the embedded Python interpreter, or an empty string when /// libcomet was built without the `python-udf` feature. The interpreter is the `libpython` that /// libcomet was linked against, so its version is fixed when libcomet is built, while PySpark's diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index b9349057c99..a4bdb5d70e4 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -329,6 +329,17 @@ public static boolean supportsPythonUdf() { } } + /** + * 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 killPythonUdfs(long task, boolean interruptRunning); + private static native String nativePythonUdfVersion(); /** diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..aa118adb3e9 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -556,6 +556,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 d0db9d9dfcf..306d2adfafa 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -115,7 +115,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 @@ -170,7 +170,20 @@ class CometExecIterator( nativeLib.setShufflePartitionPusher(createdPlan, pusher) } outputStream.foreach(output => nativeLib.setOutputStream(createdPlan, output)) - 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. @@ -374,6 +387,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 f960785647b..5979e8733e3 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -151,6 +151,25 @@ 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 + + /** + * 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 new file mode 100644 index 00000000000..a98d060baee --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala @@ -0,0 +1,111 @@ +/* + * 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, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} + +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging + +/** + * 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 { + + // Repeating the kill reaches calls of plans that the task registers after an earlier poll. + private val PollIntervalMs = 100L + + private val nextTask = new AtomicLong() + + 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() + + /** + * 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() + 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 _ = watches.remove(context, watch) + } + } + } + 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 1e742873a5a..492a34de5d5 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, PythonException, 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, SparkErrorConverter} import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus @@ -583,6 +590,101 @@ 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") + 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()) + } + } + } + + 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() + } + } + } + } + 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")