viirya commented on code in PR #6130: URL: https://github.com/apache/datafusion-comet/pull/6130#discussion_r4101574739
########## spark/src/main/spark-4.1+/org/apache/spark/sql/comet/CometArrowEvalPythonExec.scala: ########## @@ -0,0 +1,159 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.PythonEvalType +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.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] { + + 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 (op.udfs.exists(_.children.exists(expr => !hasCompatibleArrowSchema(expr.dataType))) || + op.resultAttrs.exists(attr => !hasCompatibleArrowSchema(attr.dataType))) { + return Unsupported(Some("Arrow UDF type is outside the verified native Arrow schema set")) + } + op.udfs.collectFirst { + case udf if udf.func.broadcastVars != null && !udf.func.broadcastVars.isEmpty => + "Arrow UDF broadcast variables are not supported in-process" + case udf if udf.func.pythonIncludes != null && !udf.func.pythonIncludes.isEmpty => + "Arrow UDF Python includes are not supported in-process" + case udf if udf.func.envVars != null && !udf.func.envVars.isEmpty => + "Arrow UDF Python environment overrides are not supported in-process" Review Comment: Thanks for catching this. Fixed in `4830cee65`: native execution now accepts Spark’s default `PYTHONHASHSEED=0`, while custom values still fall back to Spark. The embedded interpreter uses hash seed 0 as well. I added a test using a real PySpark `@arrow_udf` that asserts native execution is selected and compares its results, including `hash()`, with Spark. ########## native/core/src/execution/python_udf.rs: ########## @@ -0,0 +1,400 @@ +// 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}; + +#[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()?; + 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(); + let ffi_array = FFI_ArrowArray::new(&data); + let ffi_schema = FFI_ArrowSchema::try_from(data.data_type())?; + let py_arg = array_class + .call_method1( + "_import_from_c", + ( + &raw const ffi_array as Py_uintptr_t, + &raw const ffi_schema as Py_uintptr_t, Review Comment: Fixed in `4830cee65`. The input `FFI_ArrowArray` and `FFI_ArrowSchema`, as well as `ffi_return_type`, are now mutable and passed with `&raw mut` so PyArrow can update their release fields. The feature-enabled Rust tests pass. ########## native/core/src/execution/operators/arrow_python_udf.rs: ########## @@ -0,0 +1,198 @@ +// 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, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +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::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>, + schema: SchemaRef, + cache: Arc<PlanProperties>, +} + +impl ArrowPythonUdfExec { + pub fn try_new(child: Arc<dyn ExecutionPlan>, specs: Vec<ArrowPythonUdfSpec>) -> 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, + schema, + cache, + }) + } + + fn evaluate_batch( + specs: &[ArrowPythonUdfSpec], + workers: &[ArrowPythonUdf], + schema: SchemaRef, + batch: RecordBatch, + ) -> Result<RecordBatch> { + let mut columns = batch.columns().to_vec(); + for (spec, worker) in specs.iter().zip(workers) { + let args: Vec<ArrayRef> = spec + .args + .iter() + .map(|arg| arg.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::<Result<_>>()?; + columns.push(worker.evaluate_named(&args, &spec.arg_names, batch.num_rows())?); + } + 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(), + )?)) + } + + 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 specs = self.specs.clone(); + let schema = Arc::clone(&self.schema); + let stream = input.map(move |batch| { + // 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. + tokio::task::block_in_place(|| { + Self::evaluate_batch(&specs, &workers, Arc::clone(&schema), batch?) Review Comment: Fixed in `4830cee65`. We now pass `arrowMaxRecordsPerBatch` to the Rust operator and slice each input batch at that limit before invoking Python. I added a real PySpark test with the limit set to 2; it asserts that native execution is selected and that Python sees the same batch sizes as Spark. -- 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]
