viirya commented on code in PR #58978:
URL: https://github.com/apache/spark/pull/58978#discussion_r4088322208
##########
dev/spark-test-image/python-312/Dockerfile:
##########
@@ -44,6 +44,7 @@ RUN apt-get update && apt-get install -y \
libssl-dev \
openjdk-17-jdk-headless \
python3.12 \
+ python3.12-dev \
Review Comment:
The headers are needed when building JEP's JNI library from source, as this
image does. They are not required to run an already-built compatible JEP
installation; the matching Python shared library is still needed at runtime.
I've clarified this distinction in the dependency documentation.
##########
dev/spark-test-image/python-312/Dockerfile:
##########
@@ -66,3 +67,10 @@ RUN python3.12 -m pip install -U pip
RUN python3.12 -m pip install --group ml_torch --index-url
https://download.pytorch.org/whl/cpu && \
python3.12 -m pip install --group ci_classic_standard --group
ci_connect_standard && \
python3.12 -m pip cache purge
+
+# Exercise in-process UDFs in the existing pyspark-sql job. JEP needs Python
+# headers and a JDK when building its JNI library; no separate test job is
needed.
+RUN JAVA_HOME="$(dirname "$(dirname "$(readlink -f "$(command -v javac)")")")"
\
+ python3.12 -m pip install --no-cache-dir "jep==4.3.1" cffi
Review Comment:
Updated to 4.3.2 in the POM, test image, and documentation/benchmark setup
instructions. The local integration suite also passed with JEP 4.3.2.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractInProcessPythonUDFs.scala:
##########
@@ -0,0 +1,101 @@
+/*
+ * 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.execution.python
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, Expression, ExternalUserDefinedFunction, NamedExpression,
PythonUDF, WindowExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project}
+import org.apache.spark.sql.catalyst.rules.Rule
+
+/**
+ * Extracts [[InProcessPythonUDF]] expressions from logical plan nodes,
rewriting the plan
+ * so each batch of UDFs is evaluated in a dedicated [[InProcessEvalPython]]
node.
+ *
+ * Simpler than [[ExtractPythonUDFs]]: in-process UDFs are scalar and don't
support
+ * iterator mode, aggregate mode, or nested chaining in Phase 1.
+ *
+ * Example rewrite:
+ * Project [double(a), triple(b)]
+ * Scan
+ * becomes:
+ * Project [inprocessUDF0, inprocessUDF1]
+ * InProcessEvalPython [double(a), triple(b)] to [inprocessUDF0,
inprocessUDF1]
+ * Scan
+ */
+object ExtractInProcessPythonUDFs extends Rule[LogicalPlan] {
+
+ private def hasInProcessUDF(e: Expression): Boolean =
+ e.exists(_.isInstanceOf[InProcessPythonUDF])
+
+ override def apply(plan: LogicalPlan): LogicalPlan = plan.transformUp {
+ // Already extracted - skip to avoid double-wrapping
+ case p: InProcessEvalPython => p
+
+ case node: LogicalPlan if node.expressions.exists(hasInProcessUDF) =>
+ extract(node)
+ }
+
+ private def extract(plan: LogicalPlan): LogicalPlan = {
+ // Collect all distinct InProcessPythonUDFs from this plan's expressions
+ val udfs = plan.expressions
+ .flatMap(_.collect { case u: InProcessPythonUDF => u })
+ .distinct
+
+ if (udfs.isEmpty) return plan
+
+ udfs.foreach { udf =>
+ require(!udf.children.exists(_.exists {
+ case _: InProcessPythonUDF | _: PythonUDF | _:
ExternalUserDefinedFunction |
+ _: AggregateExpression | _: WindowExpression => true
+ case _ => false
+ }), "In-process Python UDFs do not support nested UDF, aggregate or
window arguments")
+ }
+
+ // Map each UDF to a fresh AttributeReference that will hold its result
+ val attributeMap = mutable.LinkedHashMap[InProcessPythonUDF,
NamedExpression]()
+
+ // For each child plan, find UDFs whose inputs are fully satisfied by that
child
+ val newChildren = plan.children.map { child =>
+ val validUdfs = udfs.filter { udf =>
Review Comment:
Removed the custom extraction rule and switched to a regular `PythonUDF`
with the new in-process eval type. `ExtractPythonUDFFromAggregate` now handles
these expressions and places their evaluation above the aggregate. Added
regression tests for UDFs over grouping keys, aggregate results, and constants.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractInProcessPythonUDFs.scala:
##########
@@ -0,0 +1,101 @@
+/*
+ * 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.execution.python
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, Expression, ExternalUserDefinedFunction, NamedExpression,
PythonUDF, WindowExpression}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project}
+import org.apache.spark.sql.catalyst.rules.Rule
+
+/**
+ * Extracts [[InProcessPythonUDF]] expressions from logical plan nodes,
rewriting the plan
+ * so each batch of UDFs is evaluated in a dedicated [[InProcessEvalPython]]
node.
+ *
+ * Simpler than [[ExtractPythonUDFs]]: in-process UDFs are scalar and don't
support
+ * iterator mode, aggregate mode, or nested chaining in Phase 1.
+ *
+ * Example rewrite:
+ * Project [double(a), triple(b)]
+ * Scan
+ * becomes:
+ * Project [inprocessUDF0, inprocessUDF1]
+ * InProcessEvalPython [double(a), triple(b)] to [inprocessUDF0,
inprocessUDF1]
+ * Scan
+ */
+object ExtractInProcessPythonUDFs extends Rule[LogicalPlan] {
+
+ private def hasInProcessUDF(e: Expression): Boolean =
+ e.exists(_.isInstanceOf[InProcessPythonUDF])
+
+ override def apply(plan: LogicalPlan): LogicalPlan = plan.transformUp {
+ // Already extracted - skip to avoid double-wrapping
+ case p: InProcessEvalPython => p
+
+ case node: LogicalPlan if node.expressions.exists(hasInProcessUDF) =>
+ extract(node)
+ }
+
+ private def extract(plan: LogicalPlan): LogicalPlan = {
+ // Collect all distinct InProcessPythonUDFs from this plan's expressions
+ val udfs = plan.expressions
+ .flatMap(_.collect { case u: InProcessPythonUDF => u })
+ .distinct
+
+ if (udfs.isEmpty) return plan
+
+ udfs.foreach { udf =>
+ require(!udf.children.exists(_.exists {
Review Comment:
These now go through the existing Python UDF extraction rules. Aggregate
arguments are handled by `ExtractPythonUDFFromAggregate`, and dependent
in-process calls are extracted into separate evaluation nodes. Added coverage
for collapsed projects, UDFs over aggregate results, nested calls, and mixed
worker/in-process UDFs. Removed the separate-stage workaround from the docs.
##########
python/pyspark/inprocess/runtime.py:
##########
@@ -0,0 +1,96 @@
+#
+# 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 Python UDF runtime entry point.
+
+``_inprocess_invoke`` is imported into the jep SharedInterpreter's global
namespace
+during executor initialization (see ``InProcessPythonRuntime.initialize()``),
then called
+directly from the JVM via ``interp.invoke("_inprocess_invoke", ...)``.
+
+Both input and output use the Arrow C Data Interface (CDI). The JVM
pre-allocates
+ArrowArray/ArrowSchema C structs for every input column and for the output,
passing
+their native addresses as Python ints. Input arrays are reconstructed via
+``pa.Array._import_from_c`` (zero-copy). The output is written via
``arr._export_to_c``
+into the JVM-owned structs (zero-copy).
+
+jep type conversions (Java -> Python):
+ byte[] -> bytes (or sequence of signed ints; masked to
unsigned below)
+ List<Long> (boxed) -> list of Python ints
+ Long -> int
+"""
+
+import traceback as _traceback
+from functools import lru_cache
+
+import pyarrow as pa
+
+from pyspark import cloudpickle
+from pyspark.sql.pandas.types import to_arrow_type
+from pyspark.sql.types import _parse_datatype_json_string
+
+_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:"
+
+
+@lru_cache(maxsize=128)
+def _load_udf(serialized_udf: bytes, return_type_json: str, timezone: str):
+ return (
+ cloudpickle.loads(serialized_udf),
+ to_arrow_type(_parse_datatype_json_string(return_type_json),
timezone=timezone),
+ )
+
+
+def _validate_result(result, expected_rows: int, expected_type: pa.DataType)
-> None:
+ if not isinstance(result, pa.Array):
+ raise TypeError(f"In-process UDF must return a pyarrow.Array, got
{type(result).__name__}")
+ if len(result) != expected_rows:
+ raise ValueError(f"In-process UDF returned {len(result)} rows;
expected {expected_rows}")
+ if result.type != expected_type:
+ raise TypeError(f"In-process UDF returned {result.type}; expected
{expected_type}")
+ result.validate()
+
+
+def _inprocess_invoke(
+ serialized_udf,
+ input_array_ptrs,
+ input_schema_ptrs,
+ output_array_ptr: int,
+ output_schema_ptr: int,
+ expected_rows: int,
+ return_type_json: str,
+ timezone: str,
+) -> None:
+ """Consume input CDI structs and export a validated, row-preserving result.
+
+ The caller owns the struct memory and releases any unconsumed exports on
failure.
+ Imported input arrays and exported output buffers follow Arrow's release
callbacks.
+ """
+ udf_key = bytes(b & 0xFF for b in serialized_udf)
+ udf_func, expected_type = _load_udf(udf_key, return_type_json, timezone)
+ if len(input_array_ptrs) != len(input_schema_ptrs):
+ raise ValueError("Mismatched input ArrowArray and ArrowSchema pointer
counts")
+ input_arrays = [
+ pa.Array._import_from_c(int(ap), int(sp))
+ for ap, sp in zip(input_array_ptrs, input_schema_ptrs)
+ ]
+ try:
+ result = udf_func(*input_arrays)
+ _validate_result(result, int(expected_rows), expected_type)
+ result._export_to_c(int(output_array_ptr), int(output_schema_ptr))
Review Comment:
Fixed by checking for offsets recursively, including nested children, and
materializing affected results with `pa.concat_arrays([result])` before CDI
export. Results without offsets keep the existing transfer path. Added runtime
and JVM integration coverage for slices, including validity bits, strings, and
nested child offsets.
##########
python/pyspark/inprocess/runtime.py:
##########
@@ -0,0 +1,96 @@
+#
+# 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 Python UDF runtime entry point.
+
+``_inprocess_invoke`` is imported into the jep SharedInterpreter's global
namespace
+during executor initialization (see ``InProcessPythonRuntime.initialize()``),
then called
+directly from the JVM via ``interp.invoke("_inprocess_invoke", ...)``.
+
+Both input and output use the Arrow C Data Interface (CDI). The JVM
pre-allocates
+ArrowArray/ArrowSchema C structs for every input column and for the output,
passing
+their native addresses as Python ints. Input arrays are reconstructed via
+``pa.Array._import_from_c`` (zero-copy). The output is written via
``arr._export_to_c``
+into the JVM-owned structs (zero-copy).
+
+jep type conversions (Java -> Python):
+ byte[] -> bytes (or sequence of signed ints; masked to
unsigned below)
+ List<Long> (boxed) -> list of Python ints
+ Long -> int
+"""
+
+import traceback as _traceback
+from functools import lru_cache
+
+import pyarrow as pa
+
+from pyspark import cloudpickle
+from pyspark.sql.pandas.types import to_arrow_type
+from pyspark.sql.types import _parse_datatype_json_string
+
+_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:"
+
+
+@lru_cache(maxsize=128)
+def _load_udf(serialized_udf: bytes, return_type_json: str, timezone: str):
+ return (
+ cloudpickle.loads(serialized_udf),
+ to_arrow_type(_parse_datatype_json_string(return_type_json),
timezone=timezone),
+ )
+
+
+def _validate_result(result, expected_rows: int, expected_type: pa.DataType)
-> None:
+ if not isinstance(result, pa.Array):
+ raise TypeError(f"In-process UDF must return a pyarrow.Array, got
{type(result).__name__}")
+ if len(result) != expected_rows:
+ raise ValueError(f"In-process UDF returned {len(result)} rows;
expected {expected_rows}")
+ if result.type != expected_type:
+ raise TypeError(f"In-process UDF returned {result.type}; expected
{expected_type}")
+ result.validate()
+
+
+def _inprocess_invoke(
+ serialized_udf,
+ input_array_ptrs,
+ input_schema_ptrs,
+ output_array_ptr: int,
+ output_schema_ptr: int,
+ expected_rows: int,
+ return_type_json: str,
+ timezone: str,
+) -> None:
+ """Consume input CDI structs and export a validated, row-preserving result.
+
+ The caller owns the struct memory and releases any unconsumed exports on
failure.
+ Imported input arrays and exported output buffers follow Arrow's release
callbacks.
+ """
+ udf_key = bytes(b & 0xFF for b in serialized_udf)
+ udf_func, expected_type = _load_udf(udf_key, return_type_json, timezone)
+ if len(input_array_ptrs) != len(input_schema_ptrs):
+ raise ValueError("Mismatched input ArrowArray and ArrowSchema pointer
counts")
+ input_arrays = [
+ pa.Array._import_from_c(int(ap), int(sp))
+ for ap, sp in zip(input_array_ptrs, input_schema_ptrs)
+ ]
+ try:
+ result = udf_func(*input_arrays)
+ _validate_result(result, int(expected_rows), expected_type)
+ result._export_to_c(int(output_array_ptr), int(output_schema_ptr))
+ except Exception:
Review Comment:
Both function registration/deserialization and invocation now catch
`BaseException` and convert it to a `RuntimeError` carrying the Python
traceback. This keeps `SystemExit` from reaching JEP's exception handler. Added
tests for `SystemExit` during invocation and deserialization. This exception
handling does not cover native process termination such as `os._exit()`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonUDF.scala:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.execution.python
+
+import org.apache.spark.sql.catalyst.expressions.{
+ Attribute, AttributeReference, AttributeSet, Expression, ExprId,
NamedExpression, Unevaluable
+}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode}
+import org.apache.spark.sql.types.DataType
+
+/**
+ * A cloudpickle-serialized Python function to be executed in-process via jep
+ * (Java Embedded Python).
+ *
+ * Distinct from [[org.apache.spark.sql.catalyst.expressions.PythonUDF]] which
uses an
+ * out-of-process Python worker connected via socket.
+ *
+ * Evaluated by [[InProcessArrowEvalExec]], which passes Arrow column buffers
to CPython
+ * as PyArrow arrays via native memory addresses (zero-copy input), then
imports the
+ * PyArrow result buffers through CDI without copying.
+ *
+ * @param name display name for plan explain output
+ * @param serializedFunc cloudpickle-serialized Python function bytes
+ * @param children input column expressions
+ * @param dataType declared return type (validated against the Arrow
result)
+ * @param udfDeterministic whether the UDF is deterministic
+ * @param resultId unique identifier for this UDF result
+ */
+case class InProcessPythonUDF(
+ name: String,
+ serializedFunc: Array[Byte],
+ children: Seq[Expression],
+ dataType: DataType,
+ udfDeterministic: Boolean = true,
+ resultId: ExprId = NamedExpression.newExprId)
+ extends Expression with Unevaluable {
Review Comment:
I replaced the custom expression with `PythonUDF` itself. It now matches
both the existing `isInstanceOf[PythonUDF]` guards and the `PYTHON_UDF` tree
pattern, including the runtime-filter guard. Added an integration regression
for partition filters to check that the UDF is not evaluated on the driver.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalExec.scala:
##########
@@ -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.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.c.{ArrowArray, ArrowSchema}
+import org.apache.arrow.util.AutoCloseables
+import org.apache.arrow.vector.VectorSchemaRoot
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
UnsafeProjection}
+import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.util.Utils
+
+/**
+ * Evaluates scalar Python UDFs using Arrow CDI in the executor process. Rows
(including
+ * computed UDF arguments) are written to Arrow once, then input and output
buffers cross
+ * the JVM/Python boundary without IPC serialization. The runtime owns the JEP
thread.
+ */
+case class InProcessArrowEvalExec(
+ udfs: Seq[InProcessPythonUDF],
+ resultAttrs: Seq[Attribute],
+ child: SparkPlan) extends UnaryExecNode {
+
+ override def output: Seq[Attribute] = child.output ++ resultAttrs
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val expressions = ArrayBuffer[Expression](child.output: _*)
+ val inputOrdinals = udfs.map { udf =>
+ udf.children.map {
+ case attr: Attribute if child.output.exists(_.exprId == attr.exprId) =>
+ child.output.indexWhere(_.exprId == attr.exprId)
+ case expr =>
+ expressions += expr
+ expressions.size - 1
+ }
+ }
+ // Synthetic names also allow joins with duplicate output column names.
+ val inputSchema = StructType(expressions.zipWithIndex.map { case (expr, i)
=>
+ StructField(s"_input$i", expr.dataType, expr.nullable)
+ }.toSeq)
+ val inputExpressions = expressions.toSeq
+ val childOutput = child.output
+ val resultOutput = output
+ val batchSize = conf.arrowMaxRecordsPerBatch
+ val maxBytes = conf.arrowMaxBytesPerBatch
+ val timeZoneId = conf.sessionLocalTimeZone
+
+ child.execute().mapPartitions { rows =>
+ val context = Option(TaskContext.get())
+ def checkCancellation(): Unit =
context.foreach(_.killTaskIfInterrupted())
+
+ val resultProjection = UnsafeProjection.create(resultOutput,
resultOutput)
+ val projectInput: InternalRow => InternalRow =
+ if (inputExpressions.size == childOutput.size) {
+ identity[InternalRow]
+ } else {
+ val projection = UnsafeProjection.create(inputExpressions,
childOutput)
+ projection.initialize(TaskContext.getPartitionId())
+ projection
+ }
+ val root = VectorSchemaRoot.create(
+ ArrowUtils.toArrowSchema(inputSchema, timeZoneId, false),
ArrowUtils.rootAllocator)
Review Comment:
Only UDF arguments are converted to Arrow now, so unrelated pass-through
structs retain their original Spark representation. For actual UDF arguments,
`failDuplicatedFieldNames` runs before Arrow schema/vector creation; duplicate
names in declared results are rejected during registration as well. Added tests
for a pass-through struct with duplicate names and rejection when such a struct
is a UDF argument.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalExec.scala:
##########
@@ -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.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.c.{ArrowArray, ArrowSchema}
+import org.apache.arrow.util.AutoCloseables
+import org.apache.arrow.vector.VectorSchemaRoot
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
UnsafeProjection}
+import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.util.Utils
+
+/**
+ * Evaluates scalar Python UDFs using Arrow CDI in the executor process. Rows
(including
+ * computed UDF arguments) are written to Arrow once, then input and output
buffers cross
+ * the JVM/Python boundary without IPC serialization. The runtime owns the JEP
thread.
+ */
+case class InProcessArrowEvalExec(
+ udfs: Seq[InProcessPythonUDF],
+ resultAttrs: Seq[Attribute],
+ child: SparkPlan) extends UnaryExecNode {
+
+ override def output: Seq[Attribute] = child.output ++ resultAttrs
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val expressions = ArrayBuffer[Expression](child.output: _*)
Review Comment:
Changed the operator to project only the UDF arguments into Arrow and keep
the original rows in a spillable `HybridRowQueue`, then join the results back
to those rows. This removes the Arrow conversion of unrelated columns. Added a
regression with a pass-through calendar interval whose microseconds would
overflow the Arrow nanosecond conversion, as well as the duplicate-field
pass-through case.
--
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]
