viirya commented on code in PR #6130: URL: https://github.com/apache/datafusion-comet/pull/6130#discussion_r4227464515
########## spark/src/main/spark-4.1+/org/apache/spark/sql/comet/CometArrowEvalPythonExec.scala: ########## @@ -0,0 +1,215 @@ +/* + * 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.spark.sql.comet + +import java.nio.charset.StandardCharsets + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expression, NamedArgumentExpression, NamedExpression, PythonUDF} +import org.apache.spark.sql.execution.{PartitioningPreservingUnaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.python.ArrowEvalPythonExec +import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType} + +import com.google.common.base.Objects +import com.google.protobuf.ByteString + +import org.apache.comet.{CometConf, ConfigEntry, NativeBase} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** Native execution for Spark 4.1+ scalar `@arrow_udf` functions. */ +object CometArrowEvalPythonExec extends CometOperatorSerde[ArrowEvalPythonExec] { + + // SparkContext adds this entry even when the user has not configured a Python + // environment. Keep other overrides on Spark's worker path. + private def hasUnsupportedEnvironment(env: java.util.Map[String, String]): Boolean = + env != null && env.asScala.exists { case (key, value) => + key != "PYTHONHASHSEED" || value != "0" + } + + // PySpark's Accumulator.__reduce__ serializes a reference to + // pyspark.accumulators._deserialize_accumulator. Spark's worker forwards its + // task-local updates to the JVM when the task finishes; embedded Python does + // not have that worker protocol. A match may also come from a harmless string + // in the pickle, in which case Spark's worker path is the safe choice. + private def hasSerializedAccumulator(command: Seq[Byte]): Boolean = + new String(command.toArray, StandardCharsets.ISO_8859_1).contains("pyspark.accumulators") + + private def hasCompatibleArrowSchema(dataType: DataType): Boolean = dataType match { + case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType | + _: FloatType | _: DoubleType | _: BinaryType | _: DateType | _: DecimalType | + _: TimestampNTZType => + true + // Spark's Arrow conversion accepts plain strings. Collated and constrained strings + // may carry semantics that are not represented by Comet's Utf8 Arrow type. + case s: StringType if s == StringType => true + case _ => false + } + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED) + + override def getSupportLevel(op: ArrowEvalPythonExec): SupportLevel = { + if (!NativeBase.supportsPythonUdf()) { + return Unsupported(Some("Native library lacks the python-udf feature")) + } + if (op.evalType != PythonEvalType.SQL_SCALAR_ARROW_UDF) { + return Unsupported(Some("Only scalar @arrow_udf is supported")) + } + if (op.udfs.isEmpty || op.udfs.length != op.resultAttrs.length) { + return Unsupported(Some("Arrow UDF functions and result attributes do not match")) + } + if (op.conf.arrowUseLargeVarTypes) { + return Unsupported(Some("Arrow UDF large variable types are not supported in-process")) + } + if (op.conf.pythonUDFProfiler.nonEmpty) { + return Unsupported(Some("Arrow UDF profiling is not supported in-process")) + } + if (SparkSession.active.sparkContext.listFiles().nonEmpty) { Review Comment: Fixed in `794b2767f`: the support check now also falls back when `SparkContext.listArchives()` is non-empty. `test_native_arrow_udf_files.py` has a new `addArchive` case that reads `offset.txt` from an extracted ZIP through `SparkFiles.get()` and asserts that the plan stays on `ArrowEvalPythonExec` with results `[7, 8]`. ########## 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 Review Comment: Fixed in `794b2767f`. The result now goes through `decode_string_arrays`, as the JVM UDF bridge does, and the SAFETY comment now says that `from_ffi` checks the schema but does not validate buffers. A Rust test returns `pa.array([b'ok', b'\xff'], pa.binary()).cast(pa.string(), safe=False)` and checks that the imported array is valid UTF-8. A Scala test feeds the same unsafe cast into a native `length` and compares the result with Spark. ########## 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)? { Review Comment: Fixed in `794b2767f`. On Spark 4.2+ the serde sets a new `accept_array_like_results` flag. When it is set, the bridge takes `pa.RecordBatch.from_arrays([result], ["_0"]).column(0)` before the length check and the cast, so lists and NumPy arrays work and a `ChunkedArray` is still rejected. Spark 4.1 still requires a `pa.Array`. A Rust test covers both modes and the chunked rejection. A Spark 4.2 test compares `to_pylist()` and NumPy results with Spark. ########## spark/src/main/spark-4.1+/org/apache/spark/sql/comet/CometArrowEvalPythonExec.scala: ########## @@ -0,0 +1,215 @@ +/* + * 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.spark.sql.comet + +import java.nio.charset.StandardCharsets + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expression, NamedArgumentExpression, NamedExpression, PythonUDF} +import org.apache.spark.sql.execution.{PartitioningPreservingUnaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.python.ArrowEvalPythonExec +import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType} + +import com.google.common.base.Objects +import com.google.protobuf.ByteString + +import org.apache.comet.{CometConf, ConfigEntry, NativeBase} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** Native execution for Spark 4.1+ scalar `@arrow_udf` functions. */ +object CometArrowEvalPythonExec extends CometOperatorSerde[ArrowEvalPythonExec] { + + // SparkContext adds this entry even when the user has not configured a Python + // environment. Keep other overrides on Spark's worker path. + private def hasUnsupportedEnvironment(env: java.util.Map[String, String]): Boolean = + env != null && env.asScala.exists { case (key, value) => + key != "PYTHONHASHSEED" || value != "0" + } + + // PySpark's Accumulator.__reduce__ serializes a reference to + // pyspark.accumulators._deserialize_accumulator. Spark's worker forwards its + // task-local updates to the JVM when the task finishes; embedded Python does + // not have that worker protocol. A match may also come from a harmless string + // in the pickle, in which case Spark's worker path is the safe choice. + private def hasSerializedAccumulator(command: Seq[Byte]): Boolean = + new String(command.toArray, StandardCharsets.ISO_8859_1).contains("pyspark.accumulators") + + private def hasCompatibleArrowSchema(dataType: DataType): Boolean = dataType match { + case _: BooleanType | _: ByteType | _: ShortType | _: IntegerType | _: LongType | + _: FloatType | _: DoubleType | _: BinaryType | _: DateType | _: DecimalType | + _: TimestampNTZType => + true + // Spark's Arrow conversion accepts plain strings. Collated and constrained strings + // may carry semantics that are not represented by Comet's Utf8 Arrow type. + case s: StringType if s == StringType => true + case _ => false + } + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED) + + override def getSupportLevel(op: ArrowEvalPythonExec): SupportLevel = { + if (!NativeBase.supportsPythonUdf()) { + return Unsupported(Some("Native library lacks the python-udf feature")) + } + if (op.evalType != PythonEvalType.SQL_SCALAR_ARROW_UDF) { + return Unsupported(Some("Only scalar @arrow_udf is supported")) + } + if (op.udfs.isEmpty || op.udfs.length != op.resultAttrs.length) { + return Unsupported(Some("Arrow UDF functions and result attributes do not match")) + } + if (op.conf.arrowUseLargeVarTypes) { + return Unsupported(Some("Arrow UDF large variable types are not supported in-process")) + } + if (op.conf.pythonUDFProfiler.nonEmpty) { Review Comment: Fixed in `794b2767f`: native planning falls back when `op.conf.pythonWorkerLoggingEnabled` is set. The planning test next to the profiler check enables `spark.sql.pyspark.worker.logging.enabled` and asserts both the fallback and its reason. ########## spark/src/main/spark-4.1+/org/apache/spark/sql/comet/CometArrowEvalPythonExec.scala: ########## @@ -0,0 +1,215 @@ +/* + * 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.spark.sql.comet + +import java.nio.charset.StandardCharsets + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expression, NamedArgumentExpression, NamedExpression, PythonUDF} +import org.apache.spark.sql.execution.{PartitioningPreservingUnaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.python.ArrowEvalPythonExec +import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType} + +import com.google.common.base.Objects +import com.google.protobuf.ByteString + +import org.apache.comet.{CometConf, ConfigEntry, NativeBase} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** Native execution for Spark 4.1+ scalar `@arrow_udf` functions. */ +object CometArrowEvalPythonExec extends CometOperatorSerde[ArrowEvalPythonExec] { + + // SparkContext adds this entry even when the user has not configured a Python + // environment. Keep other overrides on Spark's worker path. + private def hasUnsupportedEnvironment(env: java.util.Map[String, String]): Boolean = Review Comment: Changed in `794b2767f`: an entry is now accepted when `spark.executorEnv.<key>` has the same value, since cluster managers set those variables on the executor process that hosts the interpreter. A local-mode executor runs in the driver JVM, which Spark does not launch with `spark.executorEnv.*`, so in local mode the driver's own environment must also have that value; otherwise the query falls back. `PYTHONHASHSEED` still has to be `0`, because the embedded interpreter always starts with that seed. The setup section now points to `spark.executorEnv.PYTHONPATH` and describes the local-mode caveat. ########## native/core/src/execution/operators/arrow_python_udf.rs: ########## @@ -0,0 +1,441 @@ +// 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. + +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BinaryArray, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{exec_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::EmissionType; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, SendableRecordBatchStream, +}; +use futures::stream; +use futures::StreamExt; + +use crate::execution::python_udf::ArrowPythonUdf; + +#[derive(Debug, Clone)] +pub struct ArrowPythonUdfSpec { + pub command: Vec<u8>, + pub args: Vec<Arc<dyn PhysicalExpr>>, + pub arg_names: Vec<String>, + pub return_type: DataType, + pub return_name: String, + pub python_version: String, +} + +/// Evaluates scalar PyArrow UDFs inside the native pipeline. Workers are +/// instantiated in `execute`, once per partition. Python module state remains +/// shared by every task in the executor's embedded interpreter. +#[derive(Debug)] +pub struct ArrowPythonUdfExec { + child: Arc<dyn ExecutionPlan>, + specs: Vec<ArrowPythonUdfSpec>, + max_records_per_batch: usize, + max_bytes_per_batch: usize, + schema: SchemaRef, + cache: Arc<PlanProperties>, +} + +impl ArrowPythonUdfExec { + pub fn try_new( + child: Arc<dyn ExecutionPlan>, + specs: Vec<ArrowPythonUdfSpec>, + max_records_per_batch: i32, + max_bytes_per_batch: i64, + ) -> Result<Self> { + if specs.is_empty() { + return exec_err!("ArrowPythonUdfExec requires at least one UDF"); + } + let mut fields: Vec<Field> = child + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + for spec in &specs { + if spec.args.len() != spec.arg_names.len() { + return exec_err!("ArrowPythonUdf argument names are not aligned with arguments"); + } + for arg in &spec.args { + arg.data_type(&child.schema())?; + } + fields.push(Field::new( + &spec.return_name, + spec.return_type.clone(), + true, + )); + } + let schema = Arc::new(Schema::new(fields)); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + child.output_partitioning().clone(), + EmissionType::Incremental, + child.boundedness(), + )); + Ok(Self { + child, + specs, + max_records_per_batch: max_records_per_batch.max(0) as usize, + max_bytes_per_batch: max_bytes_per_batch.max(0) as usize, + schema, + cache, + }) + } + + fn evaluate_args( + specs: &[ArrowPythonUdfSpec], + batch: &RecordBatch, + ) -> Result<Vec<Vec<ArrayRef>>> { + specs + .iter() + .map(|spec| { + spec.args + .iter() + .map(|arg| arg.evaluate(batch)?.into_array(batch.num_rows())) + .collect::<Result<Vec<_>>>() + }) + .collect() + } + + // Spark's row-based Arrow writer checks its buffer size after each row, so + // the row that reaches the byte limit remains in that batch. The native + // path uses logical Arrow buffer sizes for its verified scalar types. + fn input_bytes(args: &[Vec<ArrayRef>], offset: usize, length: usize) -> Result<usize> { + if length == 0 { + return Ok(0); + } + let mut bytes = 0usize; + for array in args.iter().flatten() { + let value_bytes = match array.data_type() { + DataType::Boolean => length.div_ceil(8), + DataType::Int8 | DataType::UInt8 => length, + DataType::Int16 | DataType::UInt16 => length.saturating_mul(2), + DataType::Int32 | DataType::UInt32 | DataType::Float32 | DataType::Date32 => { + length.saturating_mul(4) + } + DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::Timestamp(TimeUnit::Microsecond, None) => length.saturating_mul(8), + DataType::Decimal128(_, _) => length.saturating_mul(16), + DataType::Utf8 => { + let values = array + .as_any() + .downcast_ref::<StringArray>() + .ok_or_else(|| { + datafusion::error::DataFusionError::Execution( + "Arrow UDF string argument has an unexpected array type" + .to_string(), + ) + })?; + let offsets = values.value_offsets(); + (offsets[offset + length] - offsets[offset]) as usize + + (length + 1).saturating_mul(4) + } + DataType::Binary => { + let values = array + .as_any() + .downcast_ref::<BinaryArray>() + .ok_or_else(|| { + datafusion::error::DataFusionError::Execution( + "Arrow UDF binary argument has an unexpected array type" + .to_string(), + ) + })?; + let offsets = values.value_offsets(); + (offsets[offset + length] - offsets[offset]) as usize + + (length + 1).saturating_mul(4) + } + other => return exec_err!("Unsupported Arrow UDF argument type: {other}"), + }; + bytes = bytes.saturating_add(value_bytes); + // Arrow Java's getBufferSizeFor counts the validity bitmap even + // when every value is non-null. + bytes = bytes.saturating_add(length.div_ceil(8)); + } + Ok(bytes) + } + + fn next_batch_length( + args: &[Vec<ArrayRef>], + offset: usize, + remaining: usize, + max_records: usize, + max_bytes: usize, + ) -> Result<usize> { + let limit = if max_records == 0 { + remaining + } else { + remaining.min(max_records) + }; + if limit == 0 || max_bytes == 0 || args.iter().all(Vec::is_empty) { + return Ok(limit); + } + if Self::input_bytes(args, offset, limit)? < max_bytes { + return Ok(limit); + } + let (mut low, mut high) = (1, limit); + while low < high { + let middle = low + (high - low) / 2; + if Self::input_bytes(args, offset, middle)? >= max_bytes { + high = middle; + } else { + low = middle + 1; + } + } + Ok(low) + } + + fn evaluate_batch( + specs: &[ArrowPythonUdfSpec], + workers: &[ArrowPythonUdf], + schema: SchemaRef, + batch: &RecordBatch, + args: &[Vec<ArrayRef>], + offset: usize, + length: usize, + ) -> Result<RecordBatch> { + let mut columns = batch.slice(offset, length).columns().to_vec(); + for ((spec, worker), function_args) in specs.iter().zip(workers).zip(args) { + let sliced_args: Vec<_> = function_args + .iter() + .map(|array| array.slice(offset, length)) + .collect(); + columns.push(worker.evaluate_named(&sliced_args, &spec.arg_names, length)?); + } + Ok(RecordBatch::try_new(schema, columns)?) + } +} + +impl DisplayAs for ArrowPythonUdfExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => { + write!(f, "CometArrowPythonUdfExec: {} UDF(s)", self.specs.len()) + } + } + } +} + +impl ExecutionPlan for ArrowPythonUdfExec { + fn name(&self) -> &str { + "CometArrowPythonUdfExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.child] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + apply_expression_roots(self.specs.iter().flat_map(|spec| spec.args.iter()), f) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + if children.len() != 1 { + return exec_err!("ArrowPythonUdfExec requires exactly one child"); + } + Ok(Arc::new(Self::try_new( + Arc::clone(&children[0]), + self.specs.clone(), + self.max_records_per_batch as i32, + self.max_bytes_per_batch as i64, + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.child.execute(partition, context)?; + let workers: Vec<_> = self + .specs + .iter() + .map(|spec| { + ArrowPythonUdf::from_command( + &spec.command, + spec.return_type.clone(), + true, + true, + &spec.python_version, + ) + }) + .collect::<std::result::Result<_, _>>()?; + let workers = Arc::new(workers); + let specs = Arc::new(self.specs.clone()); + let schema = Arc::clone(&self.schema); + let max_records_per_batch = self.max_records_per_batch; + let max_bytes_per_batch = self.max_bytes_per_batch; + let stream = input.flat_map(move |batch| { + let workers = Arc::clone(&workers); + let specs = Arc::clone(&specs); + let schema = Arc::clone(&schema); + let (batch, args, mut error) = match batch { + // Spark's Arrow writer does not invoke a scalar UDF for an empty input + // batch. Native scans may still emit one, so skip it before evaluating + // arguments or entering Python. + Ok(batch) if batch.num_rows() == 0 => (None, None, None), + Ok(batch) => { + match tokio::task::block_in_place(|| Self::evaluate_args(&specs, &batch)) { + Ok(args) => (Some(batch), Some(args), None), + Err(error) => (None, None, Some(error)), + } + } + Err(error) => (None, None, Some(error)), + }; + let mut offset = 0; + let mut failed = false; + // RecordBatch::slice shares Arrow buffers. Produce one result per poll + // so the remaining slices do not pin a second set of output batches. + stream::iter(std::iter::from_fn(move || { + if let Some(error) = error.take() { + return Some(Err(error)); + } + if failed { + return None; + } + let batch = batch.as_ref()?; + let args = args.as_ref()?; + if offset == batch.num_rows() { + return None; + } + let length = match Self::next_batch_length( + args, + offset, + batch.num_rows() - offset, + max_records_per_batch, + max_bytes_per_batch, + ) { + Ok(length) => length, + Err(error) => { + failed = true; + return Some(Err(error)); + } + }; + // Keep the JVM scan path synchronous so its Pending loop does not spin while + // Python runs. On a tokio worker, this hands its other tasks to another worker. + let result = tokio::task::block_in_place(|| { Review Comment: Nothing interrupts it today. The limitations section in `794b2767f` now says that killing a task does not stop a running Python call: a cancelled job or the losing copy of a speculative task keeps its core and competes for the GIL until the UDF returns. Interrupting the call, for example with `PyThreadState_SetAsyncExc`, is tracked in #6811, and we'll start on it once this PR is merged. -- 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]
