viirya commented on code in PR #6130: URL: https://github.com/apache/datafusion-comet/pull/6130#discussion_r4227466265
########## 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: Added in `794b2767f`. `ArrowPythonUdfExec` records `BaselineMetrics` (`output_rows` and `elapsed_compute`) and a `python_time` subset timer around each Python call. `CometArrowEvalPythonExec` declares the matching SQL metrics. The multi-row-group Parquet test checks `output_rows`, checks that `python_time` is non-zero, and checks that it stays within `elapsed_compute`. ########## 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: The traceback is in `794b2767f`: an exception from the callable is formatted with `traceback.format_exception`, so the error has the file, line, and exception type. A Rust test and a Scala end-to-end test check for the traceback. I didn't change the exception class. PySpark's `convert_exception` turns a `org.apache.spark.api.python.PythonException` cause into its Python `PythonException` only when one of the cause's stack frames is in `org.apache.spark.sql.execution.python`. This error is thrown from `CometExecIterator`, so a handler catching `pyspark.errors.PythonException` would still get an `UnknownException`. Matching it would mean constructing the exception from a helper in that package just to pass the frame check. If you think that's worth doing, I'd keep it as a follow-up. The docs now say that the traceback comes in a Comet native error. ########## 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: Fixed in `794b2767f`. `enabledConfig` is now `None`, and `getSupportLevel` checks the eval type first, then the `python-udf` feature, then the remaining conditions. It reports the config only when the UDF would otherwise run natively. A test checks that a scalar pandas UDF reports "Only scalar @arrow_udf is supported" without the config hint, and that a supported Arrow UDF with the config off still gets the hint. ########## 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: Fixed in `794b2767f`: the original "reserved from Spark's off-heap memory" wording is back. The "Who allocates what" table has a row for the embedded interpreter and PyArrow, which nothing bounds and Spark cannot see. The config doc now says Spark 4.1+. ########## 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: Updated in `794b2767f`. The type test now runs eight rows with `maxRecordsPerBatch=3` and nulls at rows 1, 4, and 7. Each type therefore reaches PyArrow at offsets 3 and 6 (non-byte-aligned for booleans), and identity results return to the JVM sliced. A new test writes a Parquet file with `parquet.block.size=512`, asserts from the footer that it has more than one row group, and reads it through the native scan with a 777-row cap, comparing the results with 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]
