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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
