viirya commented on code in PR #6130: URL: https://github.com/apache/datafusion-comet/pull/6130#discussion_r4236703413
########## native/core/src/execution/python_udf.rs: ########## @@ -0,0 +1,453 @@ +// 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. + +//! In-process bridge for Spark 4.1+ scalar Arrow UDFs. Each instance owns one +//! unpickled Python callable and must be created for one Spark task/partition. +//! The public API deliberately deals in Arrow arrays; the physical operator is +//! responsible for evaluating Catalyst arguments and preserving input columns. + +use arrow::array::{make_array, Array, ArrayRef}; +use arrow::datatypes::DataType; +use arrow::error::{ArrowError, Result}; +use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; +use pyo3::ffi::Py_uintptr_t; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyTuple}; + +fn initialize_python() -> Result<()> { + use std::ffi::CStr; + use std::sync::OnceLock; + + static RESULT: OnceLock<std::result::Result<(), String>> = OnceLock::new(); + RESULT + .get_or_init(|| { + // Spark sets PYTHONHASHSEED=0 on its Python workers by default. + // Match that seed before any Python object is created in the embedded + // interpreter, without changing the JVM process environment. + // SAFETY: OnceLock serializes initialization by Comet. No other Comet + // code accesses the Python C API before this function returns. + unsafe { + if pyo3::ffi::Py_IsInitialized() != 0 { + return Ok(()); + } + let mut config = std::mem::MaybeUninit::<pyo3::ffi::PyConfig>::uninit(); + pyo3::ffi::PyConfig_InitPythonConfig(config.as_mut_ptr()); + let mut config = config.assume_init(); + config.install_signal_handlers = 0; + config.use_hash_seed = 1; + config.hash_seed = 0; + let status = pyo3::ffi::Py_InitializeFromConfig(&config); + let error = if pyo3::ffi::PyStatus_Exception(status) != 0 { + if status.err_msg.is_null() { + "Python interpreter initialization failed".to_string() + } else { + CStr::from_ptr(status.err_msg) + .to_string_lossy() + .into_owned() + } + } else { + String::new() + }; + pyo3::ffi::PyConfig_Clear(&mut config); + if !error.is_empty() { + return Err(error); + } + pyo3::ffi::PyEval_SaveThread(); + Ok(()) + } + }) + .clone() + .map_err(ArrowError::ComputeError) +} + +#[cfg(target_os = "linux")] +fn make_python_symbols_global() -> Result<()> { + use std::ffi::CStr; + use std::sync::OnceLock; + + static RESULT: OnceLock<std::result::Result<(), String>> = OnceLock::new(); + RESULT + .get_or_init(|| { + // The JVM loads libcomet with RTLD_LOCAL. Its libpython dependency is + // local too, but CPython extension modules resolve Python C API + // symbols from the global namespace when they are imported. + let mut info = std::mem::MaybeUninit::<libc::Dl_info>::uninit(); + // SAFETY: Py_Initialize is a linked function address and info is + // writable storage for dladdr's result. + if unsafe { + libc::dladdr( + pyo3::ffi::Py_Initialize as *const () as *const libc::c_void, + info.as_mut_ptr(), + ) + } == 0 + { + return Err("cannot locate the linked Python library".to_string()); + } + // SAFETY: dladdr initialized info on success and dli_fname is a + // null-terminated path valid for the duration of this call. + let info = unsafe { info.assume_init() }; + if info.dli_fname.is_null() { + return Err("linked Python library has no path".to_string()); + } + let path = unsafe { CStr::from_ptr(info.dli_fname) }; + // RTLD_NOLOAD promotes the already-loaded libpython rather than + // loading a second copy with separate interpreter state. Keep the + // handle for the executor lifetime so its symbols remain global. + // SAFETY: path points to a valid C string returned by dladdr. + if unsafe { + libc::dlopen( + path.as_ptr(), + libc::RTLD_NOW | libc::RTLD_GLOBAL | libc::RTLD_NOLOAD, + ) + } + .is_null() + { + // SAFETY: dlerror returns a null-terminated message, if any. + let error = unsafe { libc::dlerror() }; + let detail = if error.is_null() { + "unknown dynamic loader error".to_string() + } else { + unsafe { CStr::from_ptr(error) } + .to_string_lossy() + .into_owned() + }; + return Err(format!("cannot expose Python C API symbols: {detail}")); + } + Ok(()) + }) + .clone() + .map_err(ArrowError::ComputeError) +} + +#[cfg(not(target_os = "linux"))] +fn make_python_symbols_global() -> Result<()> { + Ok(()) +} + +/// A scalar Arrow UDF loaded from Spark's pickled `(function, returnType)` command. +/// Spark serializes the return type for its worker; Comet uses the separately +/// serialized Arrow type from the physical plan instead. +pub struct ArrowPythonUdf { + callable: Py<PyAny>, + return_type: DataType, + allow_cast: bool, + safe_cast: bool, +} + +impl ArrowPythonUdf { + pub fn from_command( + command: &[u8], + return_type: DataType, + allow_cast: bool, + safe_cast: bool, + python_version: &str, + ) -> Result<Self> { + make_python_symbols_global()?; + initialize_python()?; + Python::attach(|py| { + if !python_version.is_empty() { + let info = py + .import("sys") + .map_err(python_error)? + .getattr("version_info") + .map_err(python_error)?; + let major: u8 = info + .get_item(0) + .map_err(python_error)? + .extract() + .map_err(python_error)?; + let minor: u8 = info + .get_item(1) + .map_err(python_error)? + .extract() + .map_err(python_error)?; + let actual = format!("{major}.{minor}"); + if actual != python_version { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF requires Python {python_version}, embedded interpreter is {actual}" + ))); + } + } + let pickle = py.import("pickle").map_err(python_error)?; + let loaded = pickle + .call_method1("loads", (PyBytes::new(py, command),)) + .map_err(python_error)?; + let tuple = loaded.cast::<PyTuple>().map_err(python_error)?; + if tuple.len() != 2 { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF command must contain (function, returnType), got {} items", + tuple.len() + ))); + } + let callable = tuple.get_item(0).map_err(python_error)?; + if !callable.is_callable() { + return Err(ArrowError::ComputeError( + "Arrow UDF command does not contain a callable".to_string(), + )); + } + Ok(Self { + callable: callable.unbind(), + return_type, + allow_cast, + safe_cast, + }) + }) + } + + /// 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. + pub fn evaluate(&self, args: &[ArrayRef], num_rows: usize) -> Result<ArrayRef> { + let names = vec![String::new(); args.len()]; + self.evaluate_named(args, &names, num_rows) + } + + pub fn evaluate_named( + &self, + args: &[ArrayRef], + names: &[String], + num_rows: usize, + ) -> Result<ArrayRef> { + if args.len() != names.len() { + return Err(ArrowError::ComputeError( + "Arrow UDF argument names are not aligned with arguments".to_string(), + )); + } + for (index, arg) in args.iter().enumerate() { + if arg.len() != num_rows { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF argument {index} has {} rows, expected {num_rows}", + arg.len() + ))); + } + } + + Python::attach(|py| { + let pa = py.import("pyarrow").map_err(python_error)?; + let array_class = pa.getattr("Array").map_err(python_error)?; + let mut py_args = Vec::with_capacity(args.len()); + for arg in args { + let data = arg.to_data(); + // PyArrow takes ownership of these C Data structs and clears their + // release callbacks, so the pointed-to storage must be writable. + let mut ffi_array = FFI_ArrowArray::new(&data); + let mut ffi_schema = FFI_ArrowSchema::try_from(data.data_type())?; + let py_arg = array_class + .call_method1( + "_import_from_c", + ( + &raw mut ffi_array as Py_uintptr_t, + &raw mut ffi_schema as Py_uintptr_t, + ), + ) + .map_err(python_error)?; + py_args.push(py_arg); + } + + let kwargs = pyo3::types::PyDict::new(py); + let mut positional = Vec::new(); + for (arg, name) in py_args.into_iter().zip(names) { + if name.is_empty() { + positional.push(arg); + } else { + 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(python_error)?; + if !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`; Arrow validates the schema and buffers. + let data = unsafe { from_ffi(out_array, &out_schema) }?; + if data.data_type() != &self.return_type { + return Err(ArrowError::ComputeError(format!( + "Arrow UDF returned type {}, expected {}", + data.data_type(), + self.return_type + ))); + } + Ok(make_array(data)) + }) + } +} + +fn python_error(error: impl std::fmt::Display) -> ArrowError { Review Comment: @sunchao Thanks, confirmed. Before Python 3.12, PyO3 keeps the traceback apart from the exception value, so passing only the value to `traceback.format_exception` gave it nothing to format. My local runs used 3.13, so I missed it. Fixed in `939144228`: the call now passes the type, the value, and `PyErr::traceback()`. Reproducing it on Python 3.10 also exposed a separate race. When several task threads imported PyArrow for the first time at once, one of them could get a partially initialized `pyarrow.lib`. PyArrow is now imported while the interpreter is initialized, before any other thread can run Python. Callables are unpickled one at a time under a lock that waits without holding the GIL. The feature tests passed 30 consecutive runs on Python 3.10, and the Spark 4.2 Scala suite passes with a native library built against Python 3.10. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
