You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/05/06 00:28:23 UTC
[PR] feat: Add support for accelerated PyArrow UDFs [experimental] [datafusion-comet]
andygrove opened a new pull request, #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## How are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4426341963
@mbutrovich thanks for the thorough pass — full mapping below.
**Framing.** PR description updated to "drop two `UnsafeProjection` copies and keep the stage columnar". #4240 still tracks full Arrow-to-Arrow round-trip elimination.
**Bigger items**
1. Conf + class renamed to `spark.comet.exec.pyarrowUdf.enabled` and `CometMapInBatchExec`.
2. Shared 4.x bits live in `spark-4.x/.../Spark4xMapInBatchSupport.scala` (matchers, `RunnerInputs`, `runnerInputs` helper). Each minor's `ShimCometMapInBatch` is just the per-version `ArrowPythonRunner` factory.
3. Scala suite replaced with per-version `CometMapInBatchSuite` (`spark/src/test/spark-{3.5,4.x}`) that constructs a `PythonMapInArrowExec` / `MapInArrowExec` over a stub `CometPlan` leaf and asserts `EliminateRedundantTransitions` rewrites it (plus the disabled-conf negative).
4. `CometPlan` mixed into `CometMapInBatchExec`.
**Smaller things**
5. 5-tuple replaced with `MapInBatchInfo` case class.
6. Single `match` on the matcher result in `EliminateRedundantTransitions`.
7. Unreachable raw `ColumnarToRowExec` branch removed from `extractColumnarChild`; reasoning documented in the scaladoc above it.
8. Why-comment added on the `doExecute` fallback explaining the row transition reintroduction.
9. `computeArrowPython` parameter count cut by deriving `timeZoneId` / `largeVarTypes` / `pythonRunnerConf` / `jobArtifactUUID` inside the shim — but that introduced the NPE @wForget caught on the executor (`SQLConf.ConfigReader` is driver-only). The latest commit (`6f5aca3b3`) moves resolution back to the driver in `doExecuteColumnar` and passes a serializable `RunnerInputs` case class into `computeArrowPython`, so the param count is still small without dereferencing `SQLConf` in the task closure.
10. `local*` prefixes dropped.
11. Code-restating comments dropped; the struct-wrap and barrier notes kept.
**Test infra**
12. Jar resolution consolidated into `spark/src/test/resources/pyspark/conftest.py`; pytest and the benchmark script both import `resolve_comet_jar` from there. The workflow no longer needs the inline `ls | grep`.
13. `pyarrow_udf_test.yml` switched from `paths-ignore` to a feature-file allowlist; build flavor switched to debug (no `-Prelease`, no `cargo --release`).
14. CI workflow pinned to Spark 4.0 (covers the 4.x shim path). Spark 3.5 still gets matcher coverage through the existing `CometMapInBatchSuite` runs in `pr_build_linux` / `pr_build_macos`. Compile-only 4.1/4.2 isn't wired up here yet — happy to add a matrix if you want it; my read was that the existing 4.1 build job already provides that, but let me know.
15. Benchmark caveat (passthrough UDF + `local[2]` measures Python fork/IPC, real UDFs shrink the delta further) included in the PR description.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3228508601
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
+ * (internally: rows -> Arrow -> Python -> Arrow -> rows)
+ *
+ * Optimized flow: CometNativeExec (Arrow) -> CometPythonMapInArrowExec (batch.rowIterator() ->
+ * Arrow -> Python -> Arrow columnar output)
+ *
+ * This eliminates:
+ * 1. The UnsafeProjection in ColumnarToRow (expensive copy) 2. The output Arrow->Row conversion
+ * (keeps Python output as ColumnarBatch)
+ */
+case class CometPythonMapInArrowExec(
+ func: Expression,
+ output: Seq[Attribute],
+ child: SparkPlan,
+ isBarrier: Boolean,
+ pythonEvalType: Int)
+ extends UnaryExecNode
Review Comment:
Added
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3198287997
##########
common/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -314,6 +314,18 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_PYTHON_MAP_IN_ARROW_ENABLED: ConfigEntry[Boolean] =
Review Comment:
maybe `COMET_PYARROW_SUPPORT` or something like that?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4509884836
**Item 7 (CometArrowPythonRunner dedupe across 4.1/4.2).** Took a closer look. The shared surface ends up smaller than the diff suggests because Spark's `BaseArrowPythonRunner` itself changed shape between minors:
- 4.1 constructor: `(funcs, evalType, argOffsets, schema, timeZoneId, largeVarTypes, workerConf, pythonMetrics, jobArtifactUUID, sessionUUID)`
- 4.2 constructor: `(funcs, evalType, argOffsets, schema, timeZoneId, largeVarTypes, pythonMetrics, jobArtifactUUID, sessionUUID)` — `workerConf` dropped, overridden via `runnerConf` instead
A shared `CometArrowPythonRunnerBase` would have to live in `spark-4.x/` to be visible to both minors, but `BaseArrowPythonRunner` isn't on the classpath there at a single signature, so the base would still need per-minor constructors and `super(...)` calls. Net saving: the two trait mix-ins (`CometColumnarPythonInput` + `BasicPythonArrowOutput`) factor out, which is roughly 10 lines for one more abstraction layer. I'd rather keep the per-minor pattern that mirrors `Spark4xMapInBatchSupport` already in the repo. Open to revisiting if 4.3 lands and the two minors converge on a stable signature.
**Item 11 (fuzz harness).** Borrowing the `CometCodegenFuzzSuite` shape into `test_pyarrow_udf.py` is the right next step but I'd like to do it in a follow-up so this PR doesn't grow further. Filed as #4384.
For this round I've added the targeted gaps in commit a52032176:
- Decimal precision sweep `{1, 9, 17, 18, 19, 28, 38}` × scale `{0, half, max}` — covers the long-backed / `FixedSizeBinary` boundary at 18/19
- Null density sweep `{0, 0.01, 0.5, 0.99, 1.0}` for validity-buffer memcpy
- Multi-batch per partition (4000 rows, `maxRecordsPerBatch=16`) so the persistent IPC root is exercised across ~250 batches
- 50-column mixed schema for the flattened-tree address walk
- Mid-stream zero-row batch (validity sizing + `setValueCount(0)`)
- Transforming array UDF (reverse each list) — catches symmetric encode/decode mistakes a passthrough would invert
Also filed #4383 for item 3 (drop the per-batch Comet→Spark buffer copy via direct `ArrowRecordBatch` assembly). That one supersedes the bulk-copy framing #4294 was built on; the doc rewrite in b69da2ebe explains why.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3546872995
##########
spark/src/main/spark-4.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala:
##########
@@ -0,0 +1,54 @@
+/*
+ * 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 java.io.DataOutputStream
+
+import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions}
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * Comet's Arrow Python runner for Spark 4.0. The Arrow IPC exchange lives in
+ * [[CometArrowPythonRunnerBase]]; this subclass only supplies the Spark 4.0 constructor shape and
+ * UDF command serialization.
+ */
+class CometArrowPythonRunner(
+ funcs: Seq[(ChainedPythonFunctions, Long)],
+ evalType: Int,
+ argOffsets: Array[Array[Int]],
+ override val schema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
Review Comment:
All three per-version runners take `timeZoneId: String` and `largeVarTypes: Boolean`, but neither is referenced. The base builds its schema from the source fields, and the rule already falls back to vanilla when `useLargeVarTypes` is on. Since these are brand-new classes, dropping the dead params (and the corresponding `RunnerInputs` fields) would keep the constructors honest. This is also the tell for comment #1: `sessionLocalTimeZone` is resolved on the driver and threaded all the way here, then ignored.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,386 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ // The runner's input schema is a single struct column ("struct") whose children are the
+ // user's input columns (see `schema` above). Cast once here rather than at each use site.
+ private lazy val inputStructType = schema.head.dataType.asInstanceOf[StructType]
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val childFields = inputStructType.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = inputStructType.fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
Review Comment:
The input Arrow schema is built from the source vector's `Field` (`getValueVector.getField`), and the empty-input branch hardcodes `"UTC"`. Vanilla Spark builds the schema it sends to Python with `ArrowUtils.toArrowSchema(schema, timeZoneId)`, so a `TimestampType` reaches the worker as `Timestamp(MICROSECOND, sessionLocalTimeZone)`. Comet normalizes timestamps to UTC internally, so under a non-UTC `spark.sql.session.timeZone` the worker would see a `UTC` label where vanilla Spark presents the session zone.
A passthrough UDF round-trips fine, which is why the current tests pass, but a tz-sensitive `mapInPandas` UDF or a `mapInArrow` UDF that reads `field.type.tz` could diverge from vanilla. `test_map_in_arrow_date_and_timestamp` only runs under the default session timezone. Could we add a non-UTC `spark.sql.session.timeZone` case with a tz-aware UDF to confirm parity, or document it as a known limitation?
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -0,0 +1,112 @@
+# 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.
+
+name: PyArrow UDF Tests
+
+concurrency:
+ group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }}
+ cancel-in-progress: true
+
+on:
+ push:
+ branches:
+ - main
+ paths: &feature-paths
+ - "pom.xml"
+ - "common/pom.xml"
+ - "common/src/main/scala/org/apache/comet/CometConf.scala"
+ - "spark/pom.xml"
+ - "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala"
+ - "spark/src/main/scala/org/apache/spark/sql/comet/shims/MapInBatchInfo.scala"
+ - "spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
+ - "spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
+ - "spark/src/main/spark-4.0/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
+ - "spark/src/main/spark-4.1/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
+ - "spark/src/main/spark-4.2/org/apache/spark/sql/comet/shims/ShimCometMapInBatch.scala"
+ - "spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala"
+ - "spark/src/test/resources/pyspark/conftest.py"
+ - "spark/src/test/resources/pyspark/test_pyarrow_udf.py"
+ - "spark/src/test/spark-3.5/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
+ - "spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala"
+ - ".github/workflows/pyarrow_udf_test.yml"
+ pull_request:
+ paths: *feature-paths
+ workflow_dispatch:
+
+permissions:
+ contents: read
+
+env:
+ RUST_VERSION: stable
+ RUST_BACKTRACE: 1
+ RUSTFLAGS: "-Clink-arg=-fuse-ld=bfd"
+
+jobs:
+ pyarrow-udf:
+ name: PyArrow UDF (Spark 4.0, JDK 17, Python 3.11)
Review Comment:
The 4.1 and 4.2 runners differ from 4.0 in the config param handling (`workerConf` as `override val` vs `pythonRunnerConf` + `override def`) and in `writeUDF` arity, and those files landed under `[skip ci]`. The pytest workflow only runs a real Python worker against 4.0, so a wiring divergence in the 4.1/4.2 runners would still pass CI. A single 4.1 or 4.2 smoke run would close that gap. Not a blocker for an experimental drop if out of scope.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4391005940
I assume not only speed is improved, would also be interesting to check memory metrics
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3198079384
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -0,0 +1,115 @@
+# 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.
+
+name: PyArrow UDF Tests
+
+concurrency:
+ group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }}
+ cancel-in-progress: true
+
+on:
+ push:
+ branches:
+ - main
+ paths-ignore:
+ - "benchmarks/**"
+ - "doc/**"
+ - "docs/**"
+ - "**.md"
+ - "dev/changelog/*.md"
+ - "native/core/benches/**"
+ - "native/spark-expr/benches/**"
+ - "spark/src/test/scala/org/apache/spark/sql/benchmark/**"
+ - "spark/src/main/scala/org/apache/comet/GenerateDocs.scala"
+ pull_request:
+ paths-ignore:
+ - "benchmarks/**"
+ - "doc/**"
+ - "docs/**"
+ - "**.md"
+ - "dev/changelog/*.md"
+ - "native/core/benches/**"
+ - "native/spark-expr/benches/**"
+ - "spark/src/test/scala/org/apache/spark/sql/benchmark/**"
+ - "spark/src/main/scala/org/apache/comet/GenerateDocs.scala"
+ workflow_dispatch:
+
+permissions:
+ contents: read
+
+env:
+ RUST_VERSION: stable
+ RUST_BACKTRACE: 1
+ RUSTFLAGS: "-Clink-arg=-fuse-ld=bfd"
+
+jobs:
+ pyarrow-udf:
+ name: PyArrow UDF (Spark 3.5, JDK 17, Python 3.11)
Review Comment:
Should this be Spark 4 now?
I'm assuming that this is enabled for only one version of Spark because it is experimental?
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
+ * (internally: rows -> Arrow -> Python -> Arrow -> rows)
+ *
+ * Optimized flow: CometNativeExec (Arrow) -> CometPythonMapInArrowExec (batch.rowIterator() ->
+ * Arrow -> Python -> Arrow columnar output)
+ *
+ * This eliminates:
Review Comment:
+1 !
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
+ * (internally: rows -> Arrow -> Python -> Arrow -> rows)
+ *
+ * Optimized flow: CometNativeExec (Arrow) -> CometPythonMapInArrowExec (batch.rowIterator() ->
+ * Arrow -> Python -> Arrow columnar output)
+ *
+ * This eliminates:
+ * 1. The UnsafeProjection in ColumnarToRow (expensive copy) 2. The output Arrow->Row conversion
+ * (keeps Python output as ColumnarBatch)
+ */
+case class CometPythonMapInArrowExec(
+ func: Expression,
+ output: Seq[Attribute],
+ child: SparkPlan,
+ isBarrier: Boolean,
+ pythonEvalType: Int)
+ extends UnaryExecNode
Review Comment:
Other Comet operators extend `CometPlan`. The idea was that all Comet specific common behavior (say some Comet specific metrics, for example) can be in a single place.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3212277309
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometMapInBatch
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * Comet replacement for Spark's `MapInBatchExec` family (`PythonMapInArrowExec` /
+ * `MapInArrowExec` in 4.1+ / `MapInPandasExec`). Accepts columnar input directly from a Comet
+ * child instead of going through the per-row `UnsafeProjection` that `ColumnarToRowExec` applies,
+ * and keeps the Python runner output as `ColumnarBatch` so downstream Comet operators consume it
+ * natively.
+ *
+ * What this eliminates: two `UnsafeProjection` copies (input and output) and the row transition
+ * between Comet and the Python operator. The internal row-to-Arrow IPC re-encoding inside
+ * `ArrowPythonRunner` is unchanged; full round-trip elimination is tracked in #4240.
+ */
+case class CometMapInBatchExec(
+ func: Expression,
+ output: Seq[Attribute],
+ child: SparkPlan,
+ isBarrier: Boolean,
+ pythonEvalType: Int)
+ extends UnaryExecNode
+ with CometPlan
+ with PythonSQLMetrics
+ with ShimCometMapInBatch {
+
+ override def supportsColumnar: Boolean = true
+
+ override def producedAttributes: AttributeSet = AttributeSet(output)
+
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+
+ override lazy val metrics: Map[String, SQLMetric] = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
+ "numOutputBatches" -> SQLMetrics.createMetric(sparkContext, "number of output batches"),
+ "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows")) ++
+ pythonMetrics
+
+ // Fallback for row-consuming parents (e.g. a top-level `collect()` that produces rows).
+ // Wraps this columnar exec in `ColumnarToRowExec`, reintroducing exactly the row transition
+ // this operator otherwise eliminates. Only fires when nothing downstream consumes columnar.
+ override def doExecute(): RDD[InternalRow] = {
+ ColumnarToRowExec(this).doExecute()
+ }
+
+ override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ val numOutputRows = longMetric("numOutputRows")
+ val numOutputBatches = longMetric("numOutputBatches")
+ val numInputRows = longMetric("numInputRows")
+
+ val pythonUDF = func.asInstanceOf[PythonUDF]
+ val outputAttrs = output
+ val childSchema = child.schema
+ val batchSize = conf.arrowMaxRecordsPerBatch
+ val evalType = pythonEvalType
+ val sqlConf = conf
+ val metricsCopy = pythonMetrics
+
+ val inputRDD = child.executeColumnar()
+
+ def processPartition(batches: Iterator[ColumnarBatch]): Iterator[ColumnarBatch] = {
+ val context = TaskContext.get()
+ val argOffsets = Array(Array(0))
+
+ val rowIter = batches.flatMap { batch =>
+ numInputRows += batch.numRows()
+ batch.rowIterator().asScala
+ }
+
+ val contextAwareIterator = new ContextAwareIterator(context, rowIter)
+
+ // Wrap rows as a struct, matching MapInBatchEvaluatorFactory behavior
+ val wrappedIter = contextAwareIterator.map(InternalRow(_))
+
+ val batchIter =
+ if (batchSize > 0) new BatchIterator(wrappedIter, batchSize) else Iterator(wrappedIter)
+
+ val columnarBatchIter = computeArrowPython(
+ pythonUDF,
+ evalType,
+ argOffsets,
+ StructType(Array(StructField("struct", childSchema))),
+ sqlConf,
Review Comment:
`sqlConf` object should not be passed in the task calculation logic.
```
E : org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 2.0 failed 1 times, most recent failure: Lost task 0.0 in stage 2.0 (TID 3) (836a973dfdc3 executor driver): java.lang.NullPointerException: Cannot invoke "org.apache.spark.internal.config.ConfigReader.get(String)" because "reader" is null
E at org.apache.spark.internal.config.ConfigEntry.readString(ConfigEntry.scala:96)
E at org.apache.spark.internal.config.ConfigEntryWithDefaultFunction.readFrom(ConfigEntry.scala:173)
E at org.apache.spark.sql.internal.SQLConf.getConf(SQLConf.scala:6661)
E at org.apache.spark.sql.internal.SQLConf.sessionLocalTimeZone(SQLConf.scala:6231)
E at org.apache.spark.sql.comet.shims.Spark4xMapInBatchSupport.runnerInputs(Spark4xMapInBatchSupport.scala:77)
E at org.apache.spark.sql.comet.shims.Spark4xMapInBatchSupport.runnerInputs$(Spark4xMapInBatchSupport.scala:74)
E at org.apache.spark.sql.comet.CometMapInBatchExec.runnerInputs(CometMapInBatchExec.scala:48)
E at org.apache.spark.sql.comet.shims.ShimCometMapInBatch.computeArrowPython(ShimCometMapInBatch.scala:43)
E at org.apache.spark.sql.comet.shims.ShimCometMapInBatch.computeArrowPython$(ShimCometMapInBatch.scala:33)
E at org.apache.spark.sql.comet.CometMapInBatchExec.computeArrowPython(CometMapInBatchExec.scala:48)
E at org.apache.spark.sql.comet.CometMapInBatchExec.processPartition$1(CometMapInBatchExec.scala:119)
E at org.apache.spark.sql.comet.CometMapInBatchExec.$anonfun$doExecuteColumnar$6(CometMapInBatchExec.scala:137)
E at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2(RDD.scala:901)
E at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2$adapted(RDD.scala:901)
E at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3198282722
##########
common/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -314,6 +314,18 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_PYTHON_MAP_IN_ARROW_ENABLED: ConfigEntry[Boolean] =
Review Comment:
MAP is confusing IMO, so many things related to map in 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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3501879306
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
+ childrenForceNullable)
+ }.asJava
+ // Force the field nullable where allowed. Comet's FFI-imported vectors may carry a
+ // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema
+ // and does not round-trip Spark's nullability), and the worker rejects a null value under a
+ // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field
+ // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source
+ // has no nulls.
+ val ft = field.getFieldType
+ val nullable = forceNullable || ft.isNullable
+ val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata)
+ new Field(name, newFt, newChildren)
+ }
+
+ /**
+ * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes
+ * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then
+ * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just
+ * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded /
+ * unshaded type crosses.
+ */
+ private def copyVector(src: FieldVector, dst: FieldVector): Unit = {
Review Comment:
`copyVector` calls `dst.allocateNew(...)` on every child for every batch. In Arrow Java `BaseFixedWidthVector.allocateNew(int)` calls `clear()` first, which releases both buffers before reallocating (`BaseFixedWidthVector.java:292` -> `:296` -> `releaseBuffer` at `:240-241`), and `BaseVariableWidthVector` does the same. So the only thing actually persistent across batches is the `StructVector` container. Every leaf buffer is freed and re-malloc'd per batch. On a 50-column batch that is ~100 free/malloc pairs per batch.
Can we make the root genuinely persistent? Allocate once on the first batch, then for subsequent batches use `reset()` (`BaseFixedWidthVector.java:225`, "doesn't release any memory") plus a capacity check, reallocating only when the incoming batch does not fit. That is the `setValueCount(0)` + grow-on-demand path, and it removes the per-batch churn that the design is otherwise paying for. While you are in here, the `test_map_in_arrow_multi_batch_per_partition` docstring (`test_pyarrow_udf.py`, approx lines 437-442) says it exercises "the persistent destination IPC root over multiple batches," which is not true today since the child buffers reallocate. Once the reuse lands the docstring becomes accurate. If for some reason the reuse cannot work, the docstring needs correcting either way.
##########
spark/src/test/resources/pyspark/test_pyarrow_udf.py:
##########
@@ -0,0 +1,1177 @@
+#!/usr/bin/env python3
+# 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.
+
+"""
+Pytest-driven integration tests for Comet's PyArrow UDF acceleration.
+
+Each test runs against two execution paths:
+ - "accelerated": spark.comet.exec.pyarrowUdf.enabled=true
+ (plan should contain CometMapInBatch and no ColumnarToRow)
+ - "fallback": spark.comet.exec.pyarrowUdf.enabled=false
+ (plan should contain vanilla PythonMapInArrow / MapInArrow)
+
+Usage:
+ # Build Comet first:
+ make
+
+ # Then either let the test discover the jar from spark/target, or pass it
+ # explicitly via COMET_JAR:
+ export COMET_JAR=$PWD/spark/target/comet-spark-spark3.5_2.12-0.16.0-SNAPSHOT.jar
+
+ pip install pyspark==3.5.8 pyarrow pandas pytest
+ pytest -v spark/src/test/resources/pyspark/test_pyarrow_udf.py
+"""
+
+import datetime as dt
+import os
+from decimal import Decimal
+
+import pyarrow as pa
+import pytest
+from pyspark.sql import SparkSession, types as T
+
+from conftest import resolve_comet_jar
+
+
+@pytest.fixture(scope="session")
+def spark():
+ jar = resolve_comet_jar()
+ # PYSPARK_SUBMIT_ARGS is consumed when pyspark launches its JVM. Setting
+ # --jars puts the Comet jar on both driver and executor classpaths so the
+ # CometPlugin can be loaded.
+ os.environ["PYSPARK_SUBMIT_ARGS"] = (
+ f"--jars {jar} --driver-class-path {jar} pyspark-shell"
+ )
+ session = (
+ SparkSession.builder.master("local[2]")
+ .appName("comet-pyarrow-udf-tests")
+ .config("spark.plugins", "org.apache.spark.CometPlugin")
+ .config("spark.comet.enabled", "true")
+ .config("spark.comet.exec.enabled", "true")
+ # spark.comet.exec.shuffle.enabled defaults to true, and
+ # CometSparkSessionExtensions.isCometLoaded refuses to register Comet's rules
+ # at all when shuffle is on but spark.shuffle.manager is not the Comet manager.
+ # These tests do not need Comet shuffle, so disable it explicitly to keep
+ # Comet's scan and exec rules active without configuring shuffle.
+ .config("spark.comet.exec.shuffle.enabled", "false")
+ .config("spark.memory.offHeap.enabled", "true")
+ .config("spark.memory.offHeap.size", "2g")
+ .getOrCreate()
+ )
+ try:
+ yield session
+ finally:
+ session.stop()
+
+
+@pytest.fixture(params=[True, False], ids=["accelerated", "fallback"])
+def accelerated(request, spark) -> bool:
+ spark.conf.set(
+ "spark.comet.exec.pyarrowUdf.enabled",
+ "true" if request.param else "false",
+ )
+ return request.param
+
+
+def _executed_plan(df) -> str:
+ return df._jdf.queryExecution().executedPlan().toString()
+
+
+def _assert_plan_matches_mode(
+ plan: str, accelerated: bool, vanilla_node: str = "MapInArrow"
+) -> None:
+ if accelerated:
+ assert "CometMapInBatch" in plan, (
+ f"expected CometMapInBatch in accelerated plan, got:\n{plan}"
+ )
+ assert "ColumnarToRow" not in plan, (
+ f"unexpected ColumnarToRow in accelerated plan:\n{plan}"
+ )
+ else:
+ assert "CometMapInBatch" not in plan, (
+ f"unexpected CometMapInBatch in fallback plan:\n{plan}"
+ )
+ assert vanilla_node in plan, (
+ f"expected {vanilla_node} in fallback plan, got:\n{plan}"
+ )
+
+
+def test_map_in_arrow_doubles_value(spark, tmp_path, accelerated):
+ data = [(i, float(i * 1.5), f"name_{i}") for i in range(100)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value", "name"]).write.parquet(src)
+
+ def double_value(iterator):
+ for batch in iterator:
+ pdf = batch.to_pandas()
+ pdf["value"] = pdf["value"] * 2
+ yield pa.RecordBatch.from_pandas(pdf)
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("name", T.StringType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInArrow(double_value, schema)
+
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == len(data)
+ for row, original in zip(rows, data):
+ assert row["id"] == original[0]
+ assert abs(row["value"] - original[1] * 2) < 1e-6
+ assert row["name"] == original[2]
+
+
+# All other tests use the default `vanilla_node="MapInArrow"`. The mapInPandas tests below
+# pass `MapInPandas` explicitly. The substring is the same on Spark 3.5 (PythonMapInArrowExec)
+# and Spark 4.x (MapInArrowExec) since the latter is a substring of the former.
+
+
+def test_map_in_arrow_changes_schema(spark, tmp_path, accelerated):
+ data = [(i, float(i)) for i in range(50)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def add_computed_column(iterator):
+ for batch in iterator:
+ pdf = batch.to_pandas()
+ pdf["squared"] = pdf["value"] ** 2
+ pdf["label"] = pdf["id"].apply(lambda x: f"item_{x}")
+ yield pa.RecordBatch.from_pandas(pdf)
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("squared", T.DoubleType()),
+ T.StructField("label", T.StringType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInArrow(add_computed_column, schema)
+
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == 50
+ for i, row in enumerate(rows):
+ assert abs(row["squared"] - float(i) ** 2) < 1e-6
+ assert row["label"] == f"item_{i}"
+
+
+def test_map_in_pandas_doubles_value(spark, tmp_path, accelerated):
+ data = [(i, float(i * 1.5)) for i in range(100)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def double_value(iterator):
+ for pdf in iterator:
+ pdf = pdf.copy()
+ pdf["value"] = pdf["value"] * 2
+ yield pdf
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInPandas(double_value, schema)
+
+ _assert_plan_matches_mode(
+ _executed_plan(result_df), accelerated, vanilla_node="MapInPandas"
+ )
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == len(data)
+ for row, original in zip(rows, data):
+ assert row["id"] == original[0]
+ assert abs(row["value"] - original[1] * 2) < 1e-6
+
+
+def test_map_in_pandas_changes_schema(spark, tmp_path, accelerated):
+ data = [(i, float(i)) for i in range(50)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def add_squared(iterator):
+ for pdf in iterator:
+ pdf = pdf.copy()
+ pdf["squared"] = pdf["value"] ** 2
+ yield pdf
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("squared", T.DoubleType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInPandas(add_squared, schema)
+
+ _assert_plan_matches_mode(
+ _executed_plan(result_df), accelerated, vanilla_node="MapInPandas"
+ )
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == 50
+ for i, row in enumerate(rows):
+ assert abs(row["squared"] - float(i) ** 2) < 1e-6
+
+
+def test_map_in_arrow_preserves_nulls(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("name", T.StringType()),
+ ]
+ )
+ rows = [
+ (1, "a"),
+ (2, None),
+ (None, "c"),
+ (None, None),
+ (5, "e"),
+ ]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(rows, schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ # Pure Arrow passthrough so nulls survive without a pandas roundtrip
+ # (pandas would coerce null longs to NaN floats).
+ for batch in iterator:
+ yield batch
+
+ result_df = spark.read.parquet(src).mapInArrow(passthrough, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = {(r["id"], r["name"]) for r in result_df.collect()}
+ assert out == set(rows)
+
+
+def test_map_in_arrow_empty_input(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ ]
+ )
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame([(1, 1.0), (2, 2.0)], schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ for batch in iterator:
+ yield batch
+
+ # Filter all rows out so the operator sees an empty stream from CometScan.
+ result_df = (
+ spark.read.parquet(src).where("id < 0").mapInArrow(passthrough, schema_in)
+ )
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ assert result_df.count() == 0
+
+
+def test_map_in_arrow_python_exception_propagates(spark, tmp_path, accelerated):
+ schema_in = T.StructType([T.StructField("id", T.LongType())])
+ data = [(i,) for i in range(10)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, schema_in).write.parquet(src)
+
+ sentinel = "boom-from-pyarrow-udf"
+
+ def boom(iterator):
+ for _batch in iterator:
+ raise ValueError(sentinel)
+ # Unreachable, but mapInArrow requires the callable to be a generator.
+ yield # pragma: no cover
+
+ result_df = spark.read.parquet(src).mapInArrow(boom, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ with pytest.raises(Exception) as exc_info:
+ result_df.collect()
+ assert sentinel in str(exc_info.value), (
+ f"expected sentinel {sentinel!r} in exception, got: {exc_info.value}"
+ )
+
+
+def test_map_in_arrow_decimal_type(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("amount", T.DecimalType(18, 6)),
+ ]
+ )
+ rows = [
+ (1, Decimal("123.456789")),
+ (2, Decimal("0.000001")),
+ (3, Decimal("-99999999.999999")),
+ (4, None),
+ ]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(rows, schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ for batch in iterator:
+ yield batch
+
+ result_df = spark.read.parquet(src).mapInArrow(passthrough, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = {(r["id"], r["amount"]) for r in result_df.collect()}
+ assert out == set(rows)
+
+
+@pytest.mark.parametrize(
+ "precision,scale",
+ [
+ (1, 0),
+ (9, 0),
+ (9, 4),
+ (17, 8),
+ (18, 0),
+ (18, 18),
+ (19, 0),
+ (28, 14),
+ (38, 0),
+ (38, 18),
+ (38, 38),
+ ],
+)
+def test_map_in_arrow_decimal_precision_sweep(
Review Comment:
The docstring says `BaseFixedWidthVector` handles precision <= 18 as "long-backed (8 bytes)" and precision >= 19 as "16-byte FixedSizeBinary." The Arrow `DecimalVector` that `copyVector` touches is always 16 bytes wide regardless of precision. The 8-byte long-backed form is Spark's `UnsafeRow` encoding, a layer this Arrow copy never sees. The sweep is good coverage. Please fix the rationale so it does not point at a buffer-width boundary that does not exist on the Arrow path.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
+ childrenForceNullable)
+ }.asJava
+ // Force the field nullable where allowed. Comet's FFI-imported vectors may carry a
+ // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema
+ // and does not round-trip Spark's nullability), and the worker rejects a null value under a
+ // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field
+ // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source
+ // has no nulls.
+ val ft = field.getFieldType
+ val nullable = forceNullable || ft.isNullable
+ val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata)
+ new Field(name, newFt, newChildren)
+ }
+
+ /**
+ * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes
+ * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then
+ * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just
+ * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded /
+ * unshaded type crosses.
+ */
+ private def copyVector(src: FieldVector, dst: FieldVector): Unit = {
+ val valueCount = src.getValueCount
+
+ dst match {
+ case bfwv: BaseFixedWidthVector =>
+ bfwv.allocateNew(valueCount)
+ case bvwv: BaseVariableWidthVector =>
+ bvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
+ case blvwv: BaseLargeVariableWidthVector =>
+ blvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
+ case _ =>
+ dst.setInitialCapacity(valueCount)
+ dst.allocateNew()
+ }
+
+ val srcBufs = src.getFieldBuffers
+ val dstBufs = dst.getFieldBuffers
+ require(
+ srcBufs.size == dstBufs.size,
+ s"buffer count mismatch for ${dst.getField}: src=${srcBufs.size}, dst=${dstBufs.size}")
+ var b = 0
+ while (b < srcBufs.size) {
Review Comment:
The manual `while (b < srcBufs.size)` reads more cleanly as `srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) => d.setBytes(0, s, 0, s.readableBytes) }`, with the buffer-count `require` moved to a size check beforehand. Drops a nesting level.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
Review Comment:
The writer constructs `ArrowStreamWriter(writeRoot, null, channel)` with no compression codec, so `spark.sql.execution.arrow.compression.codec` is ignored on this path. The fallback path honors it. The Python `ArrowStreamReader` auto-detects compression from the IPC message metadata, so an uncompressed stream still reads correctly, but a user who sets the codec gets it silently dropped when the rewrite fires. Please resolve the codec from `SQLConf` and pass it through the `VectorUnloader` the writer uses, matching the fallback behavior.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
Review Comment:
The read path increments `pythonNumRowsReceived` but nothing increments `pythonDataReceived`, while the writer does increment `pythonDataSent`. Stock `BasicPythonArrowOutput` populates `pythonDataReceived` as it reads frames. As written the accelerated path reports 0 bytes received, so the metric silently diverges from the fallback path depending on whether the rewrite fired. Please track the bytes read here so the metric matches the fallback. The metrics map in `CometMapInBatchExec` is where it surfaces to users.
##########
spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala:
##########
@@ -130,6 +154,48 @@ case class EliminateRedundantTransitions(session: SparkSession) extends Rule[Spa
}
}
+ /**
+ * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python
+ * UDF operator can consume directly. By the time this rule runs the earlier
+ * `hasCometNativeChild` arm has already rewritten any `ColumnarToRowExec` over a Comet columnar
+ * source to one of the Comet variants, so vanilla `ColumnarToRowExec` cannot reach here on a
+ * Comet-driven plan and is intentionally not handled.
+ */
+ private def extractColumnarChild(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case CometColumnarToRowExec(child) => Some(child)
+ case CometNativeColumnarToRowExec(child) => Some(child)
+ // Chained `mapInArrow(udf1).mapInArrow(udf2)`: by the time the outer operator is visited
+ // (transformUp is bottom-up) the inner one has already become a `CometMapInBatchExec`, which
+ // is itself columnar. There is no row transition between them to strip, so consume its
+ // columnar output directly. Its flattened output vectors are `CometVector`s, exactly what
+ // `CometMapInBatchExec`'s input path expects.
+ case child: CometMapInBatchExec => Some(child)
+ case _ => None
+ }
+
+ /**
+ * Matches the plans this rule should rewrite to `CometMapInBatchExec`. Single extractor used in
+ * the `transformUp` arm above so the matchers and conf reads run once per visited plan. Returns
+ * `(info, columnarChild)` where `columnarChild` is the Comet columnar producer that
+ * `CometMapInBatchExec` will consume directly. Returns `None` (and the arm misses) when the
+ * conf is off, when `useLargeVarTypes` forces the fallback, when the plan is not one of the
+ * version-shimmed MapInArrow / MapInPandas operators, or when the child is not a Comet
+ * columnar-to-row transition we can strip.
+ */
+ private object EligibleMapInBatch {
+ def unapply(plan: SparkPlan): Option[(MapInBatchInfo, SparkPlan)] = {
+ if (!CometConf.COMET_PYARROW_UDF_ENABLED.get()) {
+ None
+ } else if (arrowUseLargeVarTypes(plan.conf)) {
+ None
Review Comment:
The rule skips the rewrite entirely when `useLargeVarTypes` is on, because `copyVector` does a raw `setBytes` that cannot bridge Comet's 4-byte offsets to the destination's 8-byte offsets. This PR already taught `CometPlainVector` to read both offset widths, so the boundary is half-solved: the read side handles both, the write side gates the whole optimization off. Rather than gate, can the var-width arm of `copyVector` widen the offsets when the destination is a `BaseLargeVariableWidthVector` and the source is a `BaseVariableWidthVector` (read each 4-byte offset, write the 8-byte equivalent, copy the data buffer as-is)? That lets the rewrite always fire and removes a planner special case that would otherwise become dead code the moment the copy understands the wider offset.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
Review Comment:
A null-named FFI child gets the positional placeholder `_$idx`. A real sibling literally named `_0` would collide. It only applies to positional FFI fields with null names so it is unlikely, but a one-line comment noting the placeholder assumes no real field uses the `_N` form would save the next reader the head-scratch.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
Review Comment:
`schema.head.dataType.asInstanceOf[StructType]` is cast in two mutually exclusive arms. Hoist it to a single `lazy val` on the Writer so the cast lives in one place.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3223009855
##########
common/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -314,6 +314,18 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_PYTHON_MAP_IN_ARROW_ENABLED: ConfigEntry[Boolean] =
Review Comment:
Thanks, updated docs to clarify.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4433133720
> @comphead @parthchandra @mbutrovich @wForget this is ready for another round of reviews whenever you have time
The PR description is updated with the new scope
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3228501856
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -0,0 +1,115 @@
+# 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.
+
+name: PyArrow UDF Tests
+
+concurrency:
+ group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }}
+ cancel-in-progress: true
+
+on:
+ push:
+ branches:
+ - main
+ paths-ignore:
+ - "benchmarks/**"
+ - "doc/**"
+ - "docs/**"
+ - "**.md"
+ - "dev/changelog/*.md"
+ - "native/core/benches/**"
+ - "native/spark-expr/benches/**"
+ - "spark/src/test/scala/org/apache/spark/sql/benchmark/**"
+ - "spark/src/main/scala/org/apache/comet/GenerateDocs.scala"
+ pull_request:
+ paths-ignore:
+ - "benchmarks/**"
+ - "doc/**"
+ - "docs/**"
+ - "**.md"
+ - "dev/changelog/*.md"
+ - "native/core/benches/**"
+ - "native/spark-expr/benches/**"
+ - "spark/src/test/scala/org/apache/spark/sql/benchmark/**"
+ - "spark/src/main/scala/org/apache/comet/GenerateDocs.scala"
+ workflow_dispatch:
+
+permissions:
+ contents: read
+
+env:
+ RUST_VERSION: stable
+ RUST_BACKTRACE: 1
+ RUSTFLAGS: "-Clink-arg=-fuse-ld=bfd"
+
+jobs:
+ pyarrow-udf:
+ name: PyArrow UDF (Spark 3.5, JDK 17, Python 3.11)
Review Comment:
With the latest changes, this feature is now only supported for Spark 4.x. If there is user demand, we can create a separate PR to support for 3.x
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3206709449
##########
spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala:
##########
@@ -98,6 +101,26 @@ case class EliminateRedundantTransitions(session: SparkSession) extends Rule[Spa
case CometNativeColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) =>
sparkToColumnar.child
case CometSparkToColumnarExec(child: CometSparkToColumnarExec) => child
+ // Replace MapInBatchExec (PythonMapInArrowExec / MapInArrowExec / MapInPandasExec) that has
+ // a ColumnarToRow child with CometPythonMapInArrowExec to avoid the unnecessary
+ // Arrow->Row->Arrow round-trip. The matchers are version-shimmed: Spark 3.4 returns None
+ // (it lacks the required APIs) and Spark 4.1+ matches the renamed `MapInArrowExec`.
+ case p: SparkPlan
Review Comment:
The current PR only includes optimizations on jvm side, splitting it into a separate module might be a bit overengineering.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4920095557
Merged. Thanks @mbutrovich!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3506358563
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
Review Comment:
Fixed in 5802e7ca6. The reader now meters `reader.bytesRead()` around `loadNextBatch()` into `pythonDataReceived`, matching `BasicPythonArrowOutput`.
##########
spark/src/test/resources/pyspark/test_pyarrow_udf.py:
##########
@@ -0,0 +1,1177 @@
+#!/usr/bin/env python3
+# 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.
+
+"""
+Pytest-driven integration tests for Comet's PyArrow UDF acceleration.
+
+Each test runs against two execution paths:
+ - "accelerated": spark.comet.exec.pyarrowUdf.enabled=true
+ (plan should contain CometMapInBatch and no ColumnarToRow)
+ - "fallback": spark.comet.exec.pyarrowUdf.enabled=false
+ (plan should contain vanilla PythonMapInArrow / MapInArrow)
+
+Usage:
+ # Build Comet first:
+ make
+
+ # Then either let the test discover the jar from spark/target, or pass it
+ # explicitly via COMET_JAR:
+ export COMET_JAR=$PWD/spark/target/comet-spark-spark3.5_2.12-0.16.0-SNAPSHOT.jar
+
+ pip install pyspark==3.5.8 pyarrow pandas pytest
+ pytest -v spark/src/test/resources/pyspark/test_pyarrow_udf.py
+"""
+
+import datetime as dt
+import os
+from decimal import Decimal
+
+import pyarrow as pa
+import pytest
+from pyspark.sql import SparkSession, types as T
+
+from conftest import resolve_comet_jar
+
+
+@pytest.fixture(scope="session")
+def spark():
+ jar = resolve_comet_jar()
+ # PYSPARK_SUBMIT_ARGS is consumed when pyspark launches its JVM. Setting
+ # --jars puts the Comet jar on both driver and executor classpaths so the
+ # CometPlugin can be loaded.
+ os.environ["PYSPARK_SUBMIT_ARGS"] = (
+ f"--jars {jar} --driver-class-path {jar} pyspark-shell"
+ )
+ session = (
+ SparkSession.builder.master("local[2]")
+ .appName("comet-pyarrow-udf-tests")
+ .config("spark.plugins", "org.apache.spark.CometPlugin")
+ .config("spark.comet.enabled", "true")
+ .config("spark.comet.exec.enabled", "true")
+ # spark.comet.exec.shuffle.enabled defaults to true, and
+ # CometSparkSessionExtensions.isCometLoaded refuses to register Comet's rules
+ # at all when shuffle is on but spark.shuffle.manager is not the Comet manager.
+ # These tests do not need Comet shuffle, so disable it explicitly to keep
+ # Comet's scan and exec rules active without configuring shuffle.
+ .config("spark.comet.exec.shuffle.enabled", "false")
+ .config("spark.memory.offHeap.enabled", "true")
+ .config("spark.memory.offHeap.size", "2g")
+ .getOrCreate()
+ )
+ try:
+ yield session
+ finally:
+ session.stop()
+
+
+@pytest.fixture(params=[True, False], ids=["accelerated", "fallback"])
+def accelerated(request, spark) -> bool:
+ spark.conf.set(
+ "spark.comet.exec.pyarrowUdf.enabled",
+ "true" if request.param else "false",
+ )
+ return request.param
+
+
+def _executed_plan(df) -> str:
+ return df._jdf.queryExecution().executedPlan().toString()
+
+
+def _assert_plan_matches_mode(
+ plan: str, accelerated: bool, vanilla_node: str = "MapInArrow"
+) -> None:
+ if accelerated:
+ assert "CometMapInBatch" in plan, (
+ f"expected CometMapInBatch in accelerated plan, got:\n{plan}"
+ )
+ assert "ColumnarToRow" not in plan, (
+ f"unexpected ColumnarToRow in accelerated plan:\n{plan}"
+ )
+ else:
+ assert "CometMapInBatch" not in plan, (
+ f"unexpected CometMapInBatch in fallback plan:\n{plan}"
+ )
+ assert vanilla_node in plan, (
+ f"expected {vanilla_node} in fallback plan, got:\n{plan}"
+ )
+
+
+def test_map_in_arrow_doubles_value(spark, tmp_path, accelerated):
+ data = [(i, float(i * 1.5), f"name_{i}") for i in range(100)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value", "name"]).write.parquet(src)
+
+ def double_value(iterator):
+ for batch in iterator:
+ pdf = batch.to_pandas()
+ pdf["value"] = pdf["value"] * 2
+ yield pa.RecordBatch.from_pandas(pdf)
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("name", T.StringType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInArrow(double_value, schema)
+
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == len(data)
+ for row, original in zip(rows, data):
+ assert row["id"] == original[0]
+ assert abs(row["value"] - original[1] * 2) < 1e-6
+ assert row["name"] == original[2]
+
+
+# All other tests use the default `vanilla_node="MapInArrow"`. The mapInPandas tests below
+# pass `MapInPandas` explicitly. The substring is the same on Spark 3.5 (PythonMapInArrowExec)
+# and Spark 4.x (MapInArrowExec) since the latter is a substring of the former.
+
+
+def test_map_in_arrow_changes_schema(spark, tmp_path, accelerated):
+ data = [(i, float(i)) for i in range(50)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def add_computed_column(iterator):
+ for batch in iterator:
+ pdf = batch.to_pandas()
+ pdf["squared"] = pdf["value"] ** 2
+ pdf["label"] = pdf["id"].apply(lambda x: f"item_{x}")
+ yield pa.RecordBatch.from_pandas(pdf)
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("squared", T.DoubleType()),
+ T.StructField("label", T.StringType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInArrow(add_computed_column, schema)
+
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == 50
+ for i, row in enumerate(rows):
+ assert abs(row["squared"] - float(i) ** 2) < 1e-6
+ assert row["label"] == f"item_{i}"
+
+
+def test_map_in_pandas_doubles_value(spark, tmp_path, accelerated):
+ data = [(i, float(i * 1.5)) for i in range(100)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def double_value(iterator):
+ for pdf in iterator:
+ pdf = pdf.copy()
+ pdf["value"] = pdf["value"] * 2
+ yield pdf
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInPandas(double_value, schema)
+
+ _assert_plan_matches_mode(
+ _executed_plan(result_df), accelerated, vanilla_node="MapInPandas"
+ )
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == len(data)
+ for row, original in zip(rows, data):
+ assert row["id"] == original[0]
+ assert abs(row["value"] - original[1] * 2) < 1e-6
+
+
+def test_map_in_pandas_changes_schema(spark, tmp_path, accelerated):
+ data = [(i, float(i)) for i in range(50)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, ["id", "value"]).write.parquet(src)
+
+ def add_squared(iterator):
+ for pdf in iterator:
+ pdf = pdf.copy()
+ pdf["squared"] = pdf["value"] ** 2
+ yield pdf
+
+ schema = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ T.StructField("squared", T.DoubleType()),
+ ]
+ )
+ result_df = spark.read.parquet(src).mapInPandas(add_squared, schema)
+
+ _assert_plan_matches_mode(
+ _executed_plan(result_df), accelerated, vanilla_node="MapInPandas"
+ )
+
+ rows = result_df.orderBy("id").collect()
+ assert len(rows) == 50
+ for i, row in enumerate(rows):
+ assert abs(row["squared"] - float(i) ** 2) < 1e-6
+
+
+def test_map_in_arrow_preserves_nulls(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("name", T.StringType()),
+ ]
+ )
+ rows = [
+ (1, "a"),
+ (2, None),
+ (None, "c"),
+ (None, None),
+ (5, "e"),
+ ]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(rows, schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ # Pure Arrow passthrough so nulls survive without a pandas roundtrip
+ # (pandas would coerce null longs to NaN floats).
+ for batch in iterator:
+ yield batch
+
+ result_df = spark.read.parquet(src).mapInArrow(passthrough, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = {(r["id"], r["name"]) for r in result_df.collect()}
+ assert out == set(rows)
+
+
+def test_map_in_arrow_empty_input(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("value", T.DoubleType()),
+ ]
+ )
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame([(1, 1.0), (2, 2.0)], schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ for batch in iterator:
+ yield batch
+
+ # Filter all rows out so the operator sees an empty stream from CometScan.
+ result_df = (
+ spark.read.parquet(src).where("id < 0").mapInArrow(passthrough, schema_in)
+ )
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ assert result_df.count() == 0
+
+
+def test_map_in_arrow_python_exception_propagates(spark, tmp_path, accelerated):
+ schema_in = T.StructType([T.StructField("id", T.LongType())])
+ data = [(i,) for i in range(10)]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(data, schema_in).write.parquet(src)
+
+ sentinel = "boom-from-pyarrow-udf"
+
+ def boom(iterator):
+ for _batch in iterator:
+ raise ValueError(sentinel)
+ # Unreachable, but mapInArrow requires the callable to be a generator.
+ yield # pragma: no cover
+
+ result_df = spark.read.parquet(src).mapInArrow(boom, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ with pytest.raises(Exception) as exc_info:
+ result_df.collect()
+ assert sentinel in str(exc_info.value), (
+ f"expected sentinel {sentinel!r} in exception, got: {exc_info.value}"
+ )
+
+
+def test_map_in_arrow_decimal_type(spark, tmp_path, accelerated):
+ schema_in = T.StructType(
+ [
+ T.StructField("id", T.LongType()),
+ T.StructField("amount", T.DecimalType(18, 6)),
+ ]
+ )
+ rows = [
+ (1, Decimal("123.456789")),
+ (2, Decimal("0.000001")),
+ (3, Decimal("-99999999.999999")),
+ (4, None),
+ ]
+ src = str(tmp_path / "src.parquet")
+ spark.createDataFrame(rows, schema_in).write.parquet(src)
+
+ def passthrough(iterator):
+ for batch in iterator:
+ yield batch
+
+ result_df = spark.read.parquet(src).mapInArrow(passthrough, schema_in)
+ _assert_plan_matches_mode(_executed_plan(result_df), accelerated)
+
+ out = {(r["id"], r["amount"]) for r in result_df.collect()}
+ assert out == set(rows)
+
+
+@pytest.mark.parametrize(
+ "precision,scale",
+ [
+ (1, 0),
+ (9, 0),
+ (9, 4),
+ (17, 8),
+ (18, 0),
+ (18, 18),
+ (19, 0),
+ (28, 14),
+ (38, 0),
+ (38, 18),
+ (38, 38),
+ ],
+)
+def test_map_in_arrow_decimal_precision_sweep(
Review Comment:
Fixed. Rewrote the rationale: the Arrow `DecimalVector` on this path is always 16 bytes wide, and the 8-byte long-backed form is Spark's `UnsafeRow` encoding the Arrow copy never sees. The docstring now says the sweep guards the precision/scale extremes and the 18/19 point where Spark changes its own representation, rather than an Arrow buffer-width boundary.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
+ childrenForceNullable)
+ }.asJava
+ // Force the field nullable where allowed. Comet's FFI-imported vectors may carry a
+ // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema
+ // and does not round-trip Spark's nullability), and the worker rejects a null value under a
+ // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field
+ // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source
+ // has no nulls.
+ val ft = field.getFieldType
+ val nullable = forceNullable || ft.isNullable
+ val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata)
+ new Field(name, newFt, newChildren)
+ }
+
+ /**
+ * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes
+ * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then
+ * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just
+ * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded /
+ * unshaded type crosses.
+ */
+ private def copyVector(src: FieldVector, dst: FieldVector): Unit = {
+ val valueCount = src.getValueCount
+
+ dst match {
+ case bfwv: BaseFixedWidthVector =>
+ bfwv.allocateNew(valueCount)
+ case bvwv: BaseVariableWidthVector =>
+ bvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
+ case blvwv: BaseLargeVariableWidthVector =>
+ blvwv.allocateNew(src.getDataBuffer.readableBytes, valueCount)
+ case _ =>
+ dst.setInitialCapacity(valueCount)
+ dst.allocateNew()
+ }
+
+ val srcBufs = src.getFieldBuffers
+ val dstBufs = dst.getFieldBuffers
+ require(
+ srcBufs.size == dstBufs.size,
+ s"buffer count mismatch for ${dst.getField}: src=${srcBufs.size}, dst=${dstBufs.size}")
+ var b = 0
+ while (b < srcBufs.size) {
Review Comment:
Done, switched to `srcBufs.asScala.zip(dstBufs.asScala).foreach { case (s, d) => d.setBytes(0, s, 0, s.readableBytes) }` with the buffer-count `require` kept ahead of it.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
Review Comment:
Done. Hoisted to a single `lazy val inputStructType` on the Writer; both the empty-stream and first-batch arms use it.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
Review Comment:
Added a comment noting the positional `_$idx` placeholder is only applied to null-named FFI children and assumes no real sibling uses the `_N` form.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4433124320
@comphead @parthchandra @mbutrovich @wForget this is ready for another round of reviews whenever you have time
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3206658362
##########
spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometPythonMapInArrow.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.comet.shims
+
+import org.apache.spark.{JobArtifactSet, TaskContext}
+import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, PythonUDF}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.python.{ArrowPythonRunner, MapInPandasExec, PythonMapInArrowExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+trait ShimCometPythonMapInArrow {
+
+ protected def matchMapInArrow(
+ plan: SparkPlan): Option[(Expression, Seq[Attribute], SparkPlan, Boolean, Int)] =
+ plan match {
+ case p: PythonMapInArrowExec =>
+ Some((p.func, p.output, p.child, p.isBarrier, p.func.asInstanceOf[PythonUDF].evalType))
+ case _ => None
+ }
+
+ protected def matchMapInPandas(
+ plan: SparkPlan): Option[(Expression, Seq[Attribute], SparkPlan, Boolean, Int)] =
+ plan match {
+ case p: MapInPandasExec =>
+ Some((p.func, p.output, p.child, p.isBarrier, p.func.asInstanceOf[PythonUDF].evalType))
Review Comment:
> p.func.asInstanceOf[PythonUDF].evalType
Should we fix it to `PythonEvalType.SQL_MAP_PANDAS_ITER_UDF` like in spark?
https://github.com/apache/spark/blob/bb72aef62854ac35aad11cd7b4813784a539d289/sql/core/src/main/scala/org/apache/spark/sql/execution/python/MapInPandasExec.scala#L37
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add support for accelerated PyArrow UDFs [experimental] [datafusion-comet]
Posted by "github-advanced-security[bot] (via GitHub)" <gi...@apache.org>.
github-advanced-security[bot] commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3192364109
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -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.
+
+name: PyArrow UDF Tests
+
+concurrency:
+ group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }}
+ cancel-in-progress: true
+
+on:
+ push:
+ branches:
+ - main
+ paths:
+ - "spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala"
+ - "spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala"
+ - "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/test/resources/pyspark/test_pyarrow_udf.py"
+ - ".github/workflows/pyarrow_udf_test.yml"
+ - "native/**"
+ pull_request:
+ paths:
+ - "spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala"
+ - "spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala"
+ - "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/test/resources/pyspark/test_pyarrow_udf.py"
+ - ".github/workflows/pyarrow_udf_test.yml"
+ - "native/**"
+ workflow_dispatch:
+
+env:
+ RUST_VERSION: stable
+ RUST_BACKTRACE: 1
+ RUSTFLAGS: "-Clink-arg=-fuse-ld=bfd"
+
+jobs:
+ pyarrow-udf:
+ name: PyArrow UDF (Spark 3.5, JDK 17, Python 3.11)
+ runs-on: ubuntu-latest
+ container:
+ image: amd64/rust
+ env:
+ JAVA_TOOL_OPTIONS: "--add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED"
+ steps:
+ - uses: actions/checkout@v6
+
+ - name: Setup Rust & Java toolchain
+ uses: ./.github/actions/setup-builder
+ with:
+ rust-version: ${{ env.RUST_VERSION }}
+ jdk-version: 17
+
+ - name: Cache Maven dependencies
+ uses: actions/cache@v5
+ with:
+ path: |
+ ~/.m2/repository
+ /root/.m2/repository
+ key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-pyarrow-udf
+ restore-keys: |
+ ${{ runner.os }}-java-maven-
+
+ - name: Build Comet (release, Spark 3.5 / Scala 2.12)
+ run: |
+ cd native && cargo build --release
+ cd .. && ./mvnw -B -Prelease install -DskipTests -Pspark-3.5 -Pscala-2.12
+
+ - name: Install Python 3.11 and pip
+ run: |
+ apt-get update
+ apt-get install -y --no-install-recommends python3.11 python3.11-venv python3-pip
+ python3.11 -m venv /tmp/venv
+ /tmp/venv/bin/pip install --upgrade pip
+ /tmp/venv/bin/pip install "pyspark==3.5.8" "pyarrow>=14" pandas pytest
+
+ - name: Run PyArrow UDF pytest
+ run: |
+ jar=$(ls "$PWD"/spark/target/comet-spark-spark3.5_2.12-*-SNAPSHOT.jar \
+ | grep -v sources | grep -v tests | head -n1)
+ echo "Using $jar"
+ COMET_JAR="$jar" /tmp/venv/bin/python -m pytest -v \
+ spark/src/test/resources/pyspark/test_pyarrow_udf.py
Review Comment:
## CodeQL / Workflow does not contain permissions
Actions job or workflow does not limit the permissions of the GITHUB_TOKEN. Consider setting an explicit permissions block, using the following as a minimal starting point: {{contents: read}}
[Show more details](https://github.com/apache/datafusion-comet/security/code-scanning/46)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3198296936
##########
spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala:
##########
@@ -98,6 +101,26 @@ case class EliminateRedundantTransitions(session: SparkSession) extends Rule[Spa
case CometNativeColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) =>
sparkToColumnar.child
case CometSparkToColumnarExec(child: CometSparkToColumnarExec) => child
+ // Replace MapInBatchExec (PythonMapInArrowExec / MapInArrowExec / MapInPandasExec) that has
+ // a ColumnarToRow child with CometPythonMapInArrowExec to avoid the unnecessary
+ // Arrow->Row->Arrow round-trip. The matchers are version-shimmed: Spark 3.4 returns None
+ // (it lacks the required APIs) and Spark 4.1+ matches the renamed `MapInArrowExec`.
+ case p: SparkPlan
Review Comment:
thinking if it actually makes sense to have python/udf support in separate module? so user can build comet with python support or not
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3506361536
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
+ arrowWriter.start()
+ }
+
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ while (currentGroup == null || !currentGroup.hasNext) {
+ if (!inputIterator.hasNext) {
+ if (arrowWriter == null) {
+ // No input batch was ever produced (e.g. an upstream filter removed every row).
+ // Still emit a valid, empty Arrow IPC stream so the Python worker's
+ // ArrowStreamReader reads a schema and then sees zero batches, instead of failing
+ // on an absent stream ("Invalid IPC stream: negative continuation token"). There is
+ // no sample batch, so derive the schema from the Spark input schema. The timezone is
+ // irrelevant here because no rows are exchanged.
+ val inner = schema.head.dataType.asInstanceOf[StructType]
+ val childFields = inner.fields.toSeq.map(f =>
+ Utils.toArrowField(f.name, f.dataType, nullable = true, "UTC"))
+ startWriter(childFields, dataOut)
+ }
+ arrowWriter.end()
+ return false
+ }
+ currentGroup = inputIterator.next()
+ }
+
+ val cometBatch = currentGroup.next()
+ val startData = dataOut.size()
+
+ if (arrowWriter == null) {
+ // Build the destination struct root once, sized to the first batch's child fields.
+ // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
+ // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
+ // column names from the input schema (the worker reads columns by name, and shaded
+ // Arrow rejects a null field name). The field types and child structure are kept as-is
+ // so copyVector still walks the source and destination trees in lockstep.
+ val childNames = schema.head.dataType.asInstanceOf[StructType].fieldNames
+ val childFields = (0 until cometBatch.numCols()).map { i =>
+ val vecField =
+ cometBatch.column(i).asInstanceOf[CometDecodedVector].getValueVector.getField
+ renamed(vecField, childNames(i), forceNullable = true)
+ }
+ startWriter(childFields, dataOut)
+ }
+
+ var i = 0
+ while (i < cometBatch.numCols()) {
+ val src = cometBatch
+ .column(i)
+ .asInstanceOf[CometDecodedVector]
+ .getValueVector
+ .asInstanceOf[FieldVector]
+ val dst = structVec.getChildByOrdinal(i).asInstanceOf[FieldVector]
+ copyVector(src, dst)
+ i += 1
+ }
+ val numRows = cometBatch.numRows()
+ structVec.setValueCount(numRows)
+ // Mark every row of the struct non-null (all-1 validity). The validity buffer is freshly
+ // allocated and zero-initialised, so without this Python would see an all-null struct.
+ val validityBytes = (numRows + 7) / 8
+ Platform.setMemory(
+ structVec.getValidityBuffer.memoryAddress(),
+ 0xff.toByte,
+ validityBytes)
+ writeRoot.setRowCount(numRows)
+ arrowWriter.writeBatch()
+
+ pythonMetrics("pythonDataSent") += dataOut.size() - startData
+ true
+ }
+ }
+ }
+
+ override protected def newReaderIterator(
+ stream: DataInputStream,
+ writer: Writer,
+ startTime: Long,
+ env: SparkEnv,
+ worker: PythonWorker,
+ pid: Option[Int],
+ releasedOrClosed: AtomicBoolean,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue)
+ private var reader: ArrowStreamReader = _
+ private var root: VectorSchemaRoot = _
+ private var batchLoaded = true
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (reader != null) {
+ reader.close(false)
+ }
+ allocator.close()
+ }
+
+ protected override def read(): ColumnarBatch = {
+ if (writer.exception.isDefined) {
+ throw writer.exception.get
+ }
+ try {
+ if (reader != null && batchLoaded) {
+ batchLoaded = reader.loadNextBatch()
+ if (batchLoaded) {
+ // Re-wrap the (reloaded) field vectors fresh each batch, mirroring Comet's
+ // StreamReader, so each ColumnarBatch reflects the current buffers.
+ val vectors: Array[ColumnVector] = root.getFieldVectors.asScala.map { vector =>
+ CometVector.getVector(vector, null).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(root.getRowCount)
+ pythonMetrics("pythonNumRowsReceived") += root.getRowCount
+ batch
+ } else {
+ reader.close(false)
+ allocator.close()
+ read()
+ }
+ } else {
+ stream.readInt() match {
+ case SpecialLengths.START_ARROW_STREAM =>
+ reader = new ArrowStreamReader(stream, allocator)
+ root = reader.getVectorSchemaRoot()
+ read()
+ case SpecialLengths.TIMING_DATA =>
+ handleTimingData()
+ read()
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ throw handlePythonException()
+ case SpecialLengths.END_OF_DATA_SECTION =>
+ handleEndOfDataSection()
+ null
+ }
+ }
+ } catch handleException
+ }
+ }
+ }
+
+ /**
+ * Rebuild `field` with `name`, preserving its Arrow type and child structure. Any nested child
+ * whose name Comet's FFI import left null is given a positional placeholder so shaded Arrow can
+ * materialize the struct. Keeping the type and structure intact means the destination tree
+ * still mirrors the Comet source tree for [[copyVector]].
+ */
+ private def renamed(field: Field, name: String, forceNullable: Boolean): Field = {
+ // A Map's descendants must keep their original nullability: Arrow requires the entries struct
+ // (and its key) to be non-nullable, and `MapVector.createVector` rejects a nullable entries
+ // struct. Stop forcing nullable once we enter a Map subtree.
+ val childrenForceNullable = forceNullable && !field.getType.isInstanceOf[ArrowType.Map]
+ val children = field.getChildren
+ val newChildren =
+ if (children.isEmpty) children
+ else
+ children.asScala.zipWithIndex.map { case (child, idx) =>
+ renamed(
+ child,
+ if (child.getName == null) s"_$idx" else child.getName,
+ childrenForceNullable)
+ }.asJava
+ // Force the field nullable where allowed. Comet's FFI-imported vectors may carry a
+ // non-nullable Arrow `Field` even for columns that contain nulls (Comet uses positional schema
+ // and does not round-trip Spark's nullability), and the worker rejects a null value under a
+ // non-nullable field (`from_pandas(pdf, schema=batch.schema)` raises). Marking the field
+ // nullable is a safe superset; `copyVector` fills an all-valid validity buffer when the source
+ // has no nulls.
+ val ft = field.getFieldType
+ val nullable = forceNullable || ft.isNullable
+ val newFt = new FieldType(nullable, ft.getType, ft.getDictionary, ft.getMetadata)
+ new Field(name, newFt, newChildren)
+ }
+
+ /**
+ * Copy a Comet column into the destination FieldVector. Walks both trees in lockstep: sizes
+ * each destination node from the source, copies every buffer with `ArrowBuf.setBytes`, then
+ * sets value counts bottom-up so `setValueCount` does not rewrite the offset bytes we just
+ * copied. Both source and destination are Comet's (shaded) Arrow vectors, so no shaded /
+ * unshaded type crosses.
+ */
+ private def copyVector(src: FieldVector, dst: FieldVector): Unit = {
Review Comment:
Deferred the buffer reuse to #4383, which removes the `copyVector` bulk-copy path entirely (the per-batch `allocateNew` goes away with it), so a `reset()` + grow-on-demand rework here would be throwaway against code that is being deleted. I did correct the `test_map_in_arrow_multi_batch_per_partition` docstring in 5802e7ca6: it no longer claims leaf-buffer persistence (only the struct container is reused today) and now points at #4383.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -0,0 +1,378 @@
+/*
+ * 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 java.io.{DataInputStream, DataOutputStream}
+import java.nio.channels.Channels
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.{BaseFixedWidthVector, BaseLargeVariableWidthVector, BaseVariableWidthVector, FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{LargeListVector, ListVector, StructVector}
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths}
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+import org.apache.spark.unsafe.Platform
+
+import org.apache.comet.CometArrowAllocator
+import org.apache.comet.vector.{CometDecodedVector, CometVector}
+
+/**
+ * Shared base for Comet's Arrow Python runners (Spark 4.0 / 4.1 / 4.2).
+ *
+ * Unlike a stock `ArrowPythonRunner`, this does not extend Spark's `PythonArrowInput` /
+ * `BasicPythonArrowOutput` traits. Those traits expose Spark's Arrow types (`VectorSchemaRoot`,
+ * `Schema`) in their members, and the packaged `comet-spark` jar relocates `org.apache.arrow` to
+ * `org.apache.comet.shaded.arrow`, so mixing them in produces a class whose synthetic Arrow
+ * members no longer match Spark's unshaded trait contract (an `AbstractMethodError` at runtime).
+ *
+ * Instead it extends only the Arrow-agnostic `BasePythonRunner` and performs the Arrow IPC
+ * exchange itself using Comet's (shaded) Arrow. The Python worker only ever sees a standard Arrow
+ * IPC byte stream, which is version-neutral, so nothing crosses the shaded/unshaded boundary:
+ * - Input: each Comet `ColumnarBatch` is copied into a shaded struct root and written to the
+ * worker with a shaded `ArrowStreamWriter`.
+ * - Output: the worker's Arrow IPC is read with a shaded `ArrowStreamReader` straight into
+ * `CometVector`s, which is exactly what `CometMapInBatchExec` and downstream native operators
+ * consume.
+ *
+ * `BasePythonRunner` has the same shape across Spark 4.0/4.1/4.2; only the subclass constructor
+ * arguments and `writeUDF` differ, so those stay in the per-version subclasses.
+ */
+private[python] trait CometArrowPythonRunnerBase
+ extends BasePythonRunner[Iterator[ColumnarBatch], ColumnarBatch] {
+
+ /** Worker configuration written to the Python worker before execution. */
+ protected def workerConf: Map[String, String]
+
+ /** Comet's Python SQL metrics (data sent/received, rows). */
+ protected def pythonMetrics: Map[String, SQLMetric]
+
+ /** Version-specific UDF command serialization. */
+ protected def writeUDF(dataOut: DataOutputStream): Unit
+
+ /**
+ * Input schema as Comet hands it to the runner: a single non-nullable struct named "struct"
+ * whose children are the user's input columns. Comet's FFI-imported vectors carry Arrow
+ * `Field`s with null names (Comet uses positional schema), so these names are the source of
+ * truth for the field names written into the IPC stream that the Python worker reads by name.
+ */
+ protected def schema: StructType
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(funcs.head.funcs.head.pythonExec)
+
+ override val faultHandlerEnabled: Boolean = SQLConf.get.pythonUDFWorkerFaulthandlerEnabled
+ override val idleTimeoutSeconds: Long = SQLConf.get.pythonUDFWorkerIdleTimeoutSeconds
+ override val hideTraceback: Boolean = SQLConf.get.pysparkHideTraceback
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+
+ override protected def newWriter(
+ env: SparkEnv,
+ worker: PythonWorker,
+ inputIterator: Iterator[Iterator[ColumnarBatch]],
+ partitionIndex: Int,
+ context: TaskContext): Writer = {
+ new Writer(env, worker, inputIterator, partitionIndex, context) {
+
+ private val allocator =
+ CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue)
+ private var currentGroup: Iterator[ColumnarBatch] = _
+ private var arrowWriter: ArrowStreamWriter = _
+ private var writeRoot: VectorSchemaRoot = _
+ private var structVec: StructVector = _
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ if (writeRoot != null) {
+ writeRoot.close()
+ }
+ allocator.close()
+ }
+
+ protected override def writeCommand(dataOut: DataOutputStream): Unit = {
+ // handleMetadataBeforeExec: write the worker config as key/value string pairs.
+ dataOut.writeInt(workerConf.size)
+ for ((k, v) <- workerConf) {
+ PythonRDD.writeUTF(k, dataOut)
+ PythonRDD.writeUTF(v, dataOut)
+ }
+ writeUDF(dataOut)
+ }
+
+ /** Build the destination struct root and start the writer from the given child fields. */
+ private def startWriter(childFields: Seq[Field], dataOut: DataOutputStream): Unit = {
+ val structField =
+ new Field(
+ "struct",
+ new FieldType(false, ArrowType.Struct.INSTANCE, null),
+ childFields.asJava)
+ structVec = structField.createVector(allocator).asInstanceOf[StructVector]
+ writeRoot = new VectorSchemaRoot(Seq[FieldVector](structVec).asJava)
+ arrowWriter = new ArrowStreamWriter(writeRoot, null, Channels.newChannel(dataOut))
Review Comment:
Deferred to #4383. One correctness note on the "fallback honors it" framing: on Spark 4.0 the fallback does *not* compress input to Python either. `PythonArrowInput` there builds `new ArrowStreamWriter(root, null, dataOut)` and writes via `arrowWriter.writeBatch()`, with no `VectorUnloader`/codec. The codec-on-input path only exists on 4.1+, where it is applied through `VectorUnloader.getRecordBatch()` + `MessageSerializer`. #4383 moves this operator onto that same unloader assembly, which is exactly where the codec drops in and where per-version behavior can be matched cleanly. Applying it in the current `copyVector`/`ArrowStreamWriter` writer would be throwaway against that change.
##########
spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala:
##########
@@ -130,6 +154,48 @@ case class EliminateRedundantTransitions(session: SparkSession) extends Rule[Spa
}
}
+ /**
+ * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python
+ * UDF operator can consume directly. By the time this rule runs the earlier
+ * `hasCometNativeChild` arm has already rewritten any `ColumnarToRowExec` over a Comet columnar
+ * source to one of the Comet variants, so vanilla `ColumnarToRowExec` cannot reach here on a
+ * Comet-driven plan and is intentionally not handled.
+ */
+ private def extractColumnarChild(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case CometColumnarToRowExec(child) => Some(child)
+ case CometNativeColumnarToRowExec(child) => Some(child)
+ // Chained `mapInArrow(udf1).mapInArrow(udf2)`: by the time the outer operator is visited
+ // (transformUp is bottom-up) the inner one has already become a `CometMapInBatchExec`, which
+ // is itself columnar. There is no row transition between them to strip, so consume its
+ // columnar output directly. Its flattened output vectors are `CometVector`s, exactly what
+ // `CometMapInBatchExec`'s input path expects.
+ case child: CometMapInBatchExec => Some(child)
+ case _ => None
+ }
+
+ /**
+ * Matches the plans this rule should rewrite to `CometMapInBatchExec`. Single extractor used in
+ * the `transformUp` arm above so the matchers and conf reads run once per visited plan. Returns
+ * `(info, columnarChild)` where `columnarChild` is the Comet columnar producer that
+ * `CometMapInBatchExec` will consume directly. Returns `None` (and the arm misses) when the
+ * conf is off, when `useLargeVarTypes` forces the fallback, when the plan is not one of the
+ * version-shimmed MapInArrow / MapInPandas operators, or when the child is not a Comet
+ * columnar-to-row transition we can strip.
+ */
+ private object EligibleMapInBatch {
+ def unapply(plan: SparkPlan): Option[(MapInBatchInfo, SparkPlan)] = {
+ if (!CometConf.COMET_PYARROW_UDF_ENABLED.get()) {
+ None
+ } else if (arrowUseLargeVarTypes(plan.conf)) {
+ None
Review Comment:
Deferred to #4383. Agreed the `useLargeVarTypes` gate becomes dead code once the copy understands 8-byte offsets, but the offset widening belongs in the reworked copy path (#4383 assembles the `ArrowRecordBatch` directly and is where the offset handling gets re-answered) rather than in the `copyVector` version that is being removed. The current gate is a safe fallback to vanilla Spark, not a correctness gap, so I left it in place for this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3206677431
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
+ * (internally: rows -> Arrow -> Python -> Arrow -> rows)
+ *
+ * Optimized flow: CometNativeExec (Arrow) -> CometPythonMapInArrowExec (batch.rowIterator() ->
+ * Arrow -> Python -> Arrow columnar output)
+ *
+ * This eliminates:
+ * 1. The UnsafeProjection in ColumnarToRow (expensive copy) 2. The output Arrow->Row conversion
+ * (keeps Python output as ColumnarBatch)
+ */
+case class CometPythonMapInArrowExec(
Review Comment:
Does this support both `MapInArrowExec` and `MapInPandasExec`? Would it be better to name it `CometMapInBatchExec`?
refer: https://github.com/apache/spark/blob/bb72aef62854ac35aad11cd7b4813784a539d289/sql/core/src/main/scala/org/apache/spark/sql/execution/python/MapInBatchExec.scala#L36
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3206665356
##########
spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimCometPythonMapInArrow.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.comet.shims
+
+import org.apache.spark.{JobArtifactSet, TaskContext}
+import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, PythonUDF}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.python.{ArrowPythonRunner, MapInPandasExec, PythonMapInArrowExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+trait ShimCometPythonMapInArrow {
+
+ protected def matchMapInArrow(
+ plan: SparkPlan): Option[(Expression, Seq[Attribute], SparkPlan, Boolean, Int)] =
+ plan match {
+ case p: PythonMapInArrowExec =>
+ Some((p.func, p.output, p.child, p.isBarrier, p.func.asInstanceOf[PythonUDF].evalType))
Review Comment:
> p.func.asInstanceOf[PythonUDF].evalType
Should we fix it to `PythonEvalType.SQL_MAP_ARROW_ITER_UDF` like in spark?
https://github.com/apache/spark/blob/d94025b4a370ffd237f17f3d797ac200b1ae4ae9/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonMapInArrowExec.scala#L35
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "wForget (via GitHub)" <gi...@apache.org>.
wForget commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4404135616
Thanks @andygrove , overall it looks good to me, just some minor comments.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4398007703
I am reviewing this morning, please do not merge.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3198299635
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
Review Comment:
this would be nice to have also in `pyarrow-udfs`
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometPythonMapInArrowExec.scala:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometPythonMapInArrow
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * An optimized version of Spark's MapInBatchExec (PythonMapInArrowExec / MapInPandasExec) that
+ * accepts columnar input directly from Comet operators, avoiding unnecessary Arrow -> Row ->
+ * Arrow conversions.
+ *
+ * Normal Spark flow: CometNativeExec (Arrow) -> ColumnarToRow -> PythonMapInArrowExec
Review Comment:
this would be nice to have also in `pyarrow-udfs.md`
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#discussion_r3223007462
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometMapInBatchExec.scala:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{ContextAwareIterator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.comet.shims.ShimCometMapInBatch
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.python.{BatchIterator, PythonSQLMetrics}
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * Comet replacement for Spark's `MapInBatchExec` family (`PythonMapInArrowExec` /
+ * `MapInArrowExec` in 4.1+ / `MapInPandasExec`). Accepts columnar input directly from a Comet
+ * child instead of going through the per-row `UnsafeProjection` that `ColumnarToRowExec` applies,
+ * and keeps the Python runner output as `ColumnarBatch` so downstream Comet operators consume it
+ * natively.
+ *
+ * What this eliminates: two `UnsafeProjection` copies (input and output) and the row transition
+ * between Comet and the Python operator. The internal row-to-Arrow IPC re-encoding inside
+ * `ArrowPythonRunner` is unchanged; full round-trip elimination is tracked in #4240.
+ */
+case class CometMapInBatchExec(
+ func: Expression,
+ output: Seq[Attribute],
+ child: SparkPlan,
+ isBarrier: Boolean,
+ pythonEvalType: Int)
+ extends UnaryExecNode
+ with CometPlan
+ with PythonSQLMetrics
+ with ShimCometMapInBatch {
+
+ override def supportsColumnar: Boolean = true
+
+ override def producedAttributes: AttributeSet = AttributeSet(output)
+
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+
+ override lazy val metrics: Map[String, SQLMetric] = Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
+ "numOutputBatches" -> SQLMetrics.createMetric(sparkContext, "number of output batches"),
+ "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows")) ++
+ pythonMetrics
+
+ // Fallback for row-consuming parents (e.g. a top-level `collect()` that produces rows).
+ // Wraps this columnar exec in `ColumnarToRowExec`, reintroducing exactly the row transition
+ // this operator otherwise eliminates. Only fires when nothing downstream consumes columnar.
+ override def doExecute(): RDD[InternalRow] = {
+ ColumnarToRowExec(this).doExecute()
+ }
+
+ override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ val numOutputRows = longMetric("numOutputRows")
+ val numOutputBatches = longMetric("numOutputBatches")
+ val numInputRows = longMetric("numInputRows")
+
+ val pythonUDF = func.asInstanceOf[PythonUDF]
+ val outputAttrs = output
+ val childSchema = child.schema
+ val batchSize = conf.arrowMaxRecordsPerBatch
+ val evalType = pythonEvalType
+ val sqlConf = conf
+ val metricsCopy = pythonMetrics
+
+ val inputRDD = child.executeColumnar()
+
+ def processPartition(batches: Iterator[ColumnarBatch]): Iterator[ColumnarBatch] = {
+ val context = TaskContext.get()
+ val argOffsets = Array(Array(0))
+
+ val rowIter = batches.flatMap { batch =>
+ numInputRows += batch.numRows()
+ batch.rowIterator().asScala
+ }
+
+ val contextAwareIterator = new ContextAwareIterator(context, rowIter)
+
+ // Wrap rows as a struct, matching MapInBatchEvaluatorFactory behavior
+ val wrappedIter = contextAwareIterator.map(InternalRow(_))
+
+ val batchIter =
+ if (batchSize > 0) new BatchIterator(wrappedIter, batchSize) else Iterator(wrappedIter)
+
+ val columnarBatchIter = computeArrowPython(
+ pythonUDF,
+ evalType,
+ argOffsets,
+ StructType(Array(StructField("struct", childSchema))),
+ sqlConf,
Review Comment:
Thanks for catching this @wForget. I have addressed this now.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Add experimental support for accelerated PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4234:
URL: https://github.com/apache/datafusion-comet/pull/4234#issuecomment-4426698780
Moving this to draft while I implement https://github.com/apache/datafusion-comet/issues/4240 as part of this PR. I don't think it makes sense to ship this feature without it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org