andygrove commented on code in PR #6130: URL: https://github.com/apache/datafusion-comet/pull/6130#discussion_r4219921527
########## 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: The SAFETY note says Arrow validates the schema and buffers, but `from_ffi` in arrow 59 builds the result with `ArrayData::new_unchecked`. That's why `decode_string_arrays` exists in `native/common/src/utf8.rs`, and why the JVM UDF bridge runs it on its result in `native/spark-expr/src/jvm_udf/mod.rs`. A UDF that returns `pc.cast(values, pa.string(), safe=False)` over binary data hands back a `string` array with invalid UTF-8. The type already matches, so the bridge skips the cast and the array reaches native string kernels unchecked, while Spark keeps those bytes raw. Could the result go through `decode_string_arrays` the way the JVM bridge does, with the comment corrected? A test with that unsafe cast feeding a native string expression would cover it. ########## 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 { Review Comment: `CometArrowEvalPythonExec` inherits `output_rows` and `elapsed_compute`, but `ArrowPythonUdfExec` registers no metrics, so the SQL UI shows 0 rows and 0 ms for this node. Spark's node reports `pythonTotalTime`, and time in Python is the number users need when they compare this path against Spark's workers. Every other custom operator in `native/core/src/execution/operators/` implements `metrics()`. Could this one record `BaselineMetrics`, with the timer around the Python call? ########## 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: What happens when Spark kills a task while a UDF is running? Spark's `MonitorThread` destroys the Python worker `spark.python.task.killTimeout` (2s) after the interrupt. Here the call runs inside `block_in_place` on the task thread and nothing interrupts it. A cancelled job, or the losing copy of a speculative task, keeps its core and keeps competing for the GIL until the UDF returns. Could the limitations section say so, with an issue to track interrupting the call, for example with `PyThreadState_SetAsyncExc`? ########## 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]] = Review Comment: `CometExecRule` checks `enabledConfig` before `getSupportLevel`. So on 4.1+ every `ArrowEvalPythonExec`, including `@pandas_udf` and `useArrow=True` UDFs, now explains its fallback as "Set spark.comet.exec.nativeArrowPythonUDF.enabled=true to enable it." That hint can't help those UDFs, or any build without `python-udf`, which includes everything we publish. The release reported "ArrowEvalPython is not supported". Could the feature and eval-type checks run first, for example by checking the config at the end of `getSupportLevel`? ########## 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: Spark 4.2.0 rewrote the scalar Arrow UDF path in `worker.py`. It builds the output with `pa.RecordBatch.from_arrays` and casts through `enforce_schema`, so a UDF that returns a Python list or a NumPy array works there. I replayed those helpers from pyspark 4.2.0 and a list result came back as `[1, 2, 3, 4]`. 4.1.0 through 4.1.3 assert `pa.Array` in `_create_array`, so this check matches 4.1 only, and on 4.2 the native path fails queries that Spark runs. Could the shim pass a flag so that on 4.2 the bridge accepts what `from_arrays` accepts? Taking `pa.RecordBatch.from_arrays([result], ["_0"]).column(0)` would reproduce it exactly, including the `ChunkedArray` rejection. ########## 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: PySpark copies every `spark.executorEnv.*` entry into each UDF's `envVars` (`pyspark/core/context.py`), so this check falls back whenever any executor env var is configured. That includes `spark.executorEnv.PYTHONPATH` and `LD_LIBRARY_PATH`, which are the usual way to do what the setup section in `pyarrow-udfs.md` asks for. On YARN or Kubernetes the native path would quietly never run, even though in cluster mode those variables are already set on the executor process that hosts the interpreter. Could entries that match `spark.executorEnv.*` be accepted? If that isn't safe, could the docs say how to set `PYTHONPATH` without disabling the path? ########## 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: Spark 4.1 added `spark.sql.pyspark.worker.logging.enabled`. When it's on, `ArrowPythonRunner` puts `PYSPARK_SPARK_SESSION_UUID` in the worker environment, and `pyspark/logger/worker_io.py` forwards the UDF's `logging` records to the session. The embedded interpreter has none of that, so those records are silently dropped. Could this fall back when `op.conf.pythonWorkerLoggingEnabled` is set, like the profiler check here, with a planning test next to the profiler one? ########## 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: `python_error` formats the `PyErr` with `Display`, which in pyo3 0.28 is only `TypeName: message`. A failing UDF reports `ZeroDivisionError: division by zero` with no file or line, while Spark's `PythonException` carries the traceback. Could this include the formatted traceback? Is surfacing it as Spark's `PythonException` feasible too, so PySpark handlers that catch that class keep working? ########## docs/source/user-guide/latest/tuning/memory.md: ########## @@ -142,13 +142,15 @@ one line every 10 seconds for the whole executor: Comet native memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB (16 native plans, 8 memory pools); JVM Arrow allocated 310.4 MiB, 96.2 MiB of it imported from native ``` -- `allocated` is the memory that Comet's native code has allocated and not yet freed, whether or not - a pool tracks it. -- `reserved` is the part that Comet's memory pools have reserved from Spark's off-heap memory. It is - charged against `spark.memory.offHeap.size`, so the container already has room for it. A pool - sometimes has to track memory that Spark could not grant, such as a spilled batch read back from - disk while the off-heap memory is full. `reserved` leaves that memory out, since nothing charges it - against `spark.memory.offHeap.size`. +- `allocated` is the memory allocated through Comet's Rust global allocator and not yet freed, + whether or not a pool tracks it. It excludes allocations made by native libraries outside that + allocator, including the embedded Python interpreter and PyArrow when native Arrow UDFs are + enabled. +- `reserved` is the part that Comet's memory pools track. It is charged against Review Comment: This sentence changed from "reserved from Spark's off-heap memory" to "track", which contradicts the end of the same bullet: pools do track memory Spark couldn't grant, and `reserved` leaves it out. Could it go back to the original wording? Two related doc fixes: the "Who allocates what" table in `memory_management.md` should get a row for the embedded interpreter and PyArrow's bundled allocator, which nothing bounds and Spark can't see, and the config doc in `CometConf.scala` says "Spark 4.1" while the path also runs on 4.2. ########## spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala: ########## @@ -0,0 +1,333 @@ +/* + * 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.util.{Base64, Collections} + +import scala.sys.process._ + +import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction} +import org.apache.spark.sql.{CometTestBase, Row} +import org.apache.spark.sql.execution.python.UserDefinedPythonFunction +import org.apache.spark.sql.functions.{array, expr, 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.comet.{CometConf, NativeBase} + +class CometArrowPythonUdfSuite extends CometTestBase { + + test("scalar Arrow UDF falls back when the native feature is unavailable") { + assume(!NativeBase.supportsPythonUdf()) + + val function = SimplePythonFunction( + Array.emptyByteArray, + Collections.emptyMap[String, String](), + Collections.emptyList[String](), + "python3", + "3.13", + Collections.emptyList(), + null) + val udf = UserDefinedPythonFunction( + "arrow_udf", + function, + LongType, + PythonEvalType.SQL_SCALAR_ARROW_UDF, + udfDeterministic = true) + + withSQLConf(CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true") { + val source = spark.range(1, 2) + val plan = source.select(udf(source.col("id"))).queryExecution.executedPlan + assert(plan.collect { case _: CometArrowEvalPythonExec => true }.isEmpty) + } + } + + test("scalar Arrow UDF executes in the native pipeline") { + assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") + + val python = sys.env.getOrElse("PYSPARK_PYTHON", "python3") + val code = + "import base64, pyspark.cloudpickle as cloudpickle, pyarrow.compute as pc; " + + "from pyspark.sql.types import LongType; " + + "print(base64.b64encode(cloudpickle.dumps((pc.negate, LongType()))).decode())" + val command = Base64.getDecoder.decode(Seq(python, "-c", code).!!.trim) + val pythonVersion = + Seq(python, "-c", "import sys; print('%d.%d' % sys.version_info[:2])").!!.trim + val function = SimplePythonFunction( + command, + Collections.emptyMap[String, String](), + Collections.emptyList[String](), + python, + pythonVersion, + Collections.emptyList(), + null) + val udf = UserDefinedPythonFunction( + "negate_arrow", + function, + LongType, + PythonEvalType.SQL_SCALAR_ARROW_UDF, + udfDeterministic = true) + + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val source = spark.range(1, 5) + val df = source.select(udf(source.col("id"))) + assert(df.queryExecution.executedPlan.collect { case _: CometArrowEvalPythonExec => + true + }.nonEmpty) + checkAnswer(df, Seq(Row(-1L), Row(-2L), Row(-3L), Row(-4L))) + + val twoResults = + source.select(udf(source.col("id")).as("first"), udf(source.col("id") + 1L).as("second")) + val nativeUdfs = twoResults.queryExecution.executedPlan.collect { + case op: CometArrowEvalPythonExec => op + } + assert(nativeUdfs.exists(_.nativeOp.getArrowPythonUdf.getFunctionsCount == 2)) + checkAnswer(twoResults, Seq(Row(-1L, -2L), Row(-2L, -3L), Row(-3L, -4L), Row(-4L, -5L))) + + val withSubquery = spark.range(4).select(udf(expr("id + (SELECT max(id) FROM range(8))"))) + val subqueryPlan = withSubquery.queryExecution.executedPlan + assert(subqueryPlan.collect { case _: CometArrowEvalPythonExec => true }.nonEmpty) + checkAnswer(withSubquery, Seq(Row(-7L), Row(-8L), Row(-9L), Row(-10L))) + } + + withSQLConf(CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "false") { + val source = spark.range(1, 2) + val plan = source.select(udf(source.col("id"))).queryExecution.executedPlan + assert(plan.collect { case _: CometArrowEvalPythonExec => true }.isEmpty) + } + + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ARROW_EXECUTION_USE_LARGE_VAR_TYPES.key -> "true") { + val source = spark.range(1, 2) + val plan = source.select(udf(source.col("id"))).queryExecution.executedPlan + assert(plan.collect { case _: CometArrowEvalPythonExec => true }.isEmpty) + } + } + + test("native Arrow UDF plan hides commands and compares UDF identity without plan IDs") { + assume(NativeBase.supportsPythonUdf(), "native library was built without python-udf") + + val secret = "private_arrow_udf_command" + val function = SimplePythonFunction( + secret.getBytes(java.nio.charset.StandardCharsets.UTF_8), + Collections.emptyMap[String, String](), + Collections.emptyList[String](), + "python3", + "3.13", + Collections.emptyList(), + null) + val udf = UserDefinedPythonFunction( + "secret_arrow", + function, + LongType, + PythonEvalType.SQL_SCALAR_ARROW_UDF, + udfDeterministic = true) + + withSQLConf( + CometConf.COMET_NATIVE_ARROW_PYTHON_UDF_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val source = spark.range(1, 2) + val plan = source.select(udf(source.col("id"))).queryExecution.executedPlan + val native = plan.collectFirst { case op: CometArrowEvalPythonExec => op }.get + assert(!plan.treeString.contains(secret)) + assert(!native.toString.contains(secret)) + + val differentPlanId = native.copy(nativeOp = + native.nativeOp.toBuilder.setPlanId(native.nativeOp.getPlanId + 1).build()) + assert(native == differentPlanId) + assert(native.hashCode() == differentPlanId.hashCode()) + + val otherFunction = SimplePythonFunction( + "different_arrow_udf_command".getBytes(java.nio.charset.StandardCharsets.UTF_8), + Collections.emptyMap[String, String](), + Collections.emptyList[String](), + "python3", + "3.13", + Collections.emptyList(), + null) + val otherUdf = UserDefinedPythonFunction( + "secret_arrow", + otherFunction, + LongType, + PythonEvalType.SQL_SCALAR_ARROW_UDF, + udfDeterministic = true) + val otherPlan = source.select(otherUdf(source.col("id"))).queryExecution.executedPlan + val otherNative = otherPlan.collectFirst { case op: CometArrowEvalPythonExec => op }.get + assert(native.copy(udfs = otherNative.udfs) != native) + + val differentBatchSize = native.nativeOp.toBuilder + differentBatchSize.getArrowPythonUdfBuilder.setMaxRecordsPerBatch( + native.nativeOp.getArrowPythonUdf.getMaxRecordsPerBatch + 1) + assert(native.copy(nativeOp = differentBatchSize.build()) != native) + + val sameFunctionPlan = source.select(udf(source.col("id"))).queryExecution.executedPlan + val sameFunctionNative = + sameFunctionPlan.collectFirst { case op: CometArrowEvalPythonExec => op }.get + assert(native.udfs != sameFunctionNative.udfs) + assert(native.sameResult(sameFunctionNative)) + } + } + + test("native Arrow UDF preserves every accepted scalar type and nulls") { Review Comment: This test runs two rows per type, so no argument reaches Python sliced. Could it use more rows with `spark.sql.execution.arrow.maxRecordsPerBatch` set to about 3? Then each type crosses into PyArrow at a non-zero offset (booleans at a non-byte-aligned one), and identity results come back to the JVM sliced. A case that reads a multi-row-group Parquet file would also cover the scan-fed path, since every test here starts from `spark.range`. -- 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]
