From d09e2b5a3d79a0df7fa3e642c632078417b7dd2b Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 30 Aug 2026 18:05:29 +0000 Subject: [PATCH] fix: align Spark 4.2 Python worker configuration --- .github/workflows/pyarrow_udf_test.yml | 12 ++++++++---- .../python/CometArrowPythonRunner.scala | 7 +++++++ .../python/CometArrowPythonRunnerBase.scala | 19 ++++++++++++++----- 3 files changed, 29 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pyarrow_udf_test.yml b/.github/workflows/pyarrow_udf_test.yml index 8fff5d414cd..6af67aa7e36 100644 --- a/.github/workflows/pyarrow_udf_test.yml +++ b/.github/workflows/pyarrow_udf_test.yml @@ -38,6 +38,9 @@ on: - "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.0/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala" + - "spark/src/main/spark-4.1/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala" + - "spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala" - "spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/Spark4xMapInBatchSupport.scala" - "spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala" - "spark/src/test/resources/pyspark/conftest.py" @@ -66,16 +69,17 @@ jobs: fail-fast: false matrix: include: - # Run a real Python worker against each 4.x runner. The 4.0 and 4.1 subclasses differ in - # constructor shape and `writeUDF` arity, so both need worker-level coverage; a wiring - # divergence in one would otherwise slip through. (4.2 is a preview with no released - # pyspark, so it stays compile-only via the pr_build matrix.) + # Run a real Python worker against each 4.x runner. Their constructor, command framing, + # and `writeUDF` shapes differ, so a wiring divergence would otherwise slip through. - name: Spark 4.0 maven_profiles: "-Pspark-4.0 -Pscala-2.13" pyspark: "4.0.4" - name: Spark 4.1 maven_profiles: "-Pspark-4.1" pyspark: "4.1.3" + - name: Spark 4.2 + maven_profiles: "-Pspark-4.2" + pyspark: "4.2.0" container: # Pinned to the Debian 12 (bookworm) base so the system `python3` is 3.11. The default # `amd64/rust` image is Debian 13 (trixie) which ships Python 3.13 and no python3.11 apt diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala index adc4d185c71..789e4b97a96 100644 --- a/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala +++ b/spark/src/main/spark-4.2/org/apache/spark/sql/execution/python/CometArrowPythonRunner.scala @@ -50,6 +50,13 @@ class CometArrowPythonRunner( override protected def workerConf: Map[String, String] = pythonRunnerConf + // Spark 4.2 writes runnerConf and evalConf before writeCommand. Pass Comet's settings through the + // native slot and do not emit the legacy map inside the command, which the worker would interpret + // as the number of UDFs. + override protected def runnerConf: Map[String, String] = super.runnerConf ++ workerConf + + override protected def writeWorkerConf(dataOut: DataOutputStream): Unit = () + override protected def writeUDF(dataOut: DataOutputStream): Unit = PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets) } diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index 943287cde1f..07281d63671 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -76,6 +76,19 @@ private[python] trait CometArrowPythonRunnerBase /** Version-specific UDF command serialization. */ protected def writeUDF(dataOut: DataOutputStream): Unit + /** + * Write the worker configuration where Spark 4.0 and 4.1 workers expect it. Spark 4.2 moved + * this map into [[BasePythonRunner.runnerConf]], so its subclass overrides this hook with a + * no-op. + */ + protected def writeWorkerConf(dataOut: DataOutputStream): Unit = { + dataOut.writeInt(workerConf.size) + for ((key, value) <- workerConf) { + PythonRDD.writeUTF(key, dataOut) + PythonRDD.writeUTF(value, dataOut) + } + } + /** * 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 @@ -126,11 +139,7 @@ private[python] trait CometArrowPythonRunnerBase 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) - } + writeWorkerConf(dataOut) writeUDF(dataOut) }