[SPARK-59993][PYTHON] Collect worker and Arrow-optimized Python UDF timing metrics

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

@@ -477,21 +477,12 @@ private[spark] object BasePythonRunner extends Logging {
     }
   }
 
-  // Values used to update the existing Python SQL metrics and task spill metrics.
-  private[python] case class WorkerMetrics(
-      bootTimestampMs: Long,
-      initTimestampMs: Long,
-      finishTimestampMs: Long,
-      pythonExecutionDurationMs: Long,
-      memoryBytesSpilled: Long,
-      diskBytesSpilled: Long)
-
   private lazy val workerMetricsMapper = new ObjectMapper()
     .enable(JsonParser.Feature.STRICT_DUPLICATE_DETECTION)
     .enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS)
 
-  /** Read and validate worker metrics after METRICS_DATA; ignore additional fields. */
-  private[python] def readWorkerMetrics(stream: DataInputStream): WorkerMetrics = {
+  /** Decode every METRICS_DATA field as a named int64; reject other values. */
+  private[python] def readWorkerMetrics(stream: DataInputStream): Map[String, Long] = {
     val length = stream.readInt()
     if (length <= 0) {
       throw new SparkException(s"Invalid Python worker report length: $length")
@@ -507,21 +498,13 @@ private[spark] object BasePythonRunner extends Logging {
       throw new SparkException("Expected a Python worker JSON object")
     }
 
-    def metricValue(name: String): Long = {
-      val value = report.get(name)
-      if (value == null || !value.isIntegralNumber || !value.canConvertToLong) {
-        throw new SparkException(s"Missing or invalid Python worker metric: $name")
+    report.fields().asScala.map { entry =>
+      val value = entry.getValue
+      if (!value.isIntegralNumber || !value.canConvertToLong) {
+        throw new SparkException(s"Invalid Python worker metric: ${entry.getKey}")
       }
-      value.longValue()
-    }
-
-    WorkerMetrics(
-      metricValue("bootTimestampMs"),
-      metricValue("initTimestampMs"),
-      metricValue("finishTimestampMs"),
-      metricValue("pythonExecutionDurationMs"),
-      metricValue("memoryBytesSpilled"),
-      metricValue("diskBytesSpilled"))
+      entry.getKey -> value.longValue()
+    }.toMap
   }
 
   /**
@@ -1155,17 +1138,32 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
 
     protected def handleMetricsData(): Unit = {
       val workerMetrics = BasePythonRunner.readWorkerMetrics(stream)
-      val bootTime = workerMetrics.bootTimestampMs
-      val initTime = workerMetrics.initTimestampMs
-      val finishTime = workerMetrics.finishTimestampMs
-      val pythonExecutionDurationMs = workerMetrics.pythonExecutionDurationMs
+      def requiredMetric(name: String): Long = workerMetrics.getOrElse(name,
+        throw new SparkException(s"Missing Python worker metric: $name"))
+
+      // Read required inputs before updating any SQL or task metrics.
+      val bootTime = requiredMetric("bootTimestampMs")
+      val initTime = requiredMetric("initTimestampMs")
+      val finishTime = requiredMetric("finishTimestampMs")
+      val pythonExecutionDurationMs = requiredMetric("pythonExecutionDurationMs")
+      val memoryBytesSpilled = requiredMetric("memoryBytesSpilled")
+      val diskBytesSpilled = requiredMetric("diskBytesSpilled")
+
       // A reused Python worker records bootTime before waiting for this task, so it can precede
       // startTime. Use the later timestamp to exclude the worker's idle time from initialization.
       val pythonWorkerInitializationStartTime = math.max(startTime, bootTime)
       val boot = pythonWorkerInitializationStartTime - startTime
       val init = initTime - pythonWorkerInitializationStartTime
-      val finish = finishTime - initTime
       val total = finishTime - startTime
+      val finish = finishTime - initTime
+      val calculatedMetrics = Map(
+        "pythonBootTime" -> boot,
+        "pythonInitTime" -> init,
+        "pythonTotalTime" -> total,
+        "pythonProcessingTime" -> pythonExecutionDurationMs)
+      // Select registered report values; calculated timings take precedence over raw aliases.
+      val sqlMetrics = workerMetrics.filter { case (name, _) => metrics.contains(name) } ++
+        calculatedMetrics
 
       // Format data size for readability
       val dataKB = totalDataReceived / 1024.0
@@ -1181,12 +1179,11 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
         log"finish = ${MDC(TIME, finish)} - " +
         log"Batches: ${MDC(COUNT, batchesProcessed)}, Data: ${MDC(SIZE, dataStr)} - " +
         log"${MDC(TASK_NAME, taskIdentifier(context))}")
-      metrics.get("pythonBootTime").foreach(_.add(boot))
-      metrics.get("pythonInitTime").foreach(_.add(init))
-      metrics.get("pythonTotalTime").foreach(_.add(total))
-      metrics.get("pythonProcessingTime").foreach(_.add(pythonExecutionDurationMs))
-      context.taskMetrics().incMemoryBytesSpilled(workerMetrics.memoryBytesSpilled)
-      context.taskMetrics().incDiskBytesSpilled(workerMetrics.diskBytesSpilled)
+      sqlMetrics.foreach { case (name, value) =>
+        metrics.get(name).foreach(_.add(value))
+      }
+      context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
+      context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
     }
 
     protected def handlePythonException(): PythonException = {

core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala

@@ -720,7 +720,13 @@ class BasePythonRunnerSuite extends SparkFunSuite with SharedSparkContext {
     "{\"bootTimestampMs\":1250,\"initTimestampMs\":2500," +
       "\"finishTimestampMs\":3750,\"pythonExecutionDurationMs\":42," +
       "\"memoryBytesSpilled\":7,\"diskBytesSpilled\":9}"
-  private val expectedMetrics = BasePythonRunner.WorkerMetrics(1250L, 2500L, 3750L, 42L, 7L, 9L)
+  private val expectedMetrics = Map(
+    "bootTimestampMs" -> 1250L,
+    "initTimestampMs" -> 2500L,
+    "finishTimestampMs" -> 3750L,
+    "pythonExecutionDurationMs" -> 42L,
+    "memoryBytesSpilled" -> 7L,
+    "diskBytesSpilled" -> 9L)
 
   private def framedMetricsStream(json: String, declaredLength: Option[Int] = None,
       endMarkers: Boolean = true): DataInputStream = {
@@ -745,41 +751,32 @@ class BasePythonRunnerSuite extends SparkFunSuite with SharedSparkContext {
   }
 
   test("SPARK-59773: read worker metrics and leave the following sections aligned") {
-    val withExtra = metricsReport.dropRight(1) + ",\"futureMetric\":{\"value\":[1,2]}}"
+    val withExtra = metricsReport.dropRight(1) + ",\"futureMetric\":17,\"zeroMetric\":0}"
     val stream = framedMetricsStream(withExtra)
-    assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics)
+    val result = BasePythonRunner.readWorkerMetrics(stream)
+    assert(result == expectedMetrics ++ Map("futureMetric" -> 17L, "zeroMetric" -> 0L))
     checkEndMarkers(stream)
   }
 
-  test("SPARK-59773: worker metrics require all current fields") {
-    val fields = Seq(
-      "bootTimestampMs" -> "\"bootTimestampMs\":1250,",
-      "initTimestampMs" -> "\"initTimestampMs\":2500,",
-      "finishTimestampMs" -> "\"finishTimestampMs\":3750,",
-      "pythonExecutionDurationMs" -> "\"pythonExecutionDurationMs\":42,",
-      "memoryBytesSpilled" -> "\"memoryBytesSpilled\":7,",
-      "diskBytesSpilled" -> ",\"diskBytesSpilled\":9")
-    fields.foreach { case (name, field) =>
-      val stream = framedMetricsStream(metricsReport.replace(field, ""))
-      val error = intercept[SparkException] {
-        BasePythonRunner.readWorkerMetrics(stream)
-      }
-      assert(error.getMessage.contains(name))
-      checkEndMarkers(stream)
-    }
+  test("SPARK-59993: decode worker metric names independently of lifecycle and SQL consumers") {
+    val stream = framedMetricsStream("{\"futureDuration\":42}")
+    val result = BasePythonRunner.readWorkerMetrics(stream)
+    assert(result == Map("futureDuration" -> 42L))
+    checkEndMarkers(stream)
   }
 
-  test("SPARK-59773: worker metrics reject non-int64 values") {
+  test("SPARK-59773: worker metric decoder rejects non-int64 values") {
     val invalidValues = Seq("null", "true", "\"42\"", "42.0", "[]", "{}",
       "9223372036854775808", "-9223372036854775809")
+    val reportWithExtra = metricsReport.dropRight(1) + ",\"futureMetric\":17}"
     val fields = Seq("pythonExecutionDurationMs" -> 42,
-      "memoryBytesSpilled" -> 7, "diskBytesSpilled" -> 9)
+      "memoryBytesSpilled" -> 7, "diskBytesSpilled" -> 9, "futureMetric" -> 17)
     for {
       (name, originalValue) <- fields
       value <- invalidValues
     } {
       val field = "\"" + name + "\":"
-      val json = metricsReport.replace(field + originalValue, field + value)
+      val json = reportWithExtra.replace(field + originalValue, field + value)
       val stream = framedMetricsStream(json)
       val error = intercept[SparkException] {
         BasePythonRunner.readWorkerMetrics(stream)
@@ -797,8 +794,8 @@ class BasePythonRunnerSuite extends SparkFunSuite with SharedSparkContext {
         .replace("\"diskBytesSpilled\":9", "\"diskBytesSpilled\":" + value)
       val stream = framedMetricsStream(json)
       assert(BasePythonRunner.readWorkerMetrics(stream) ==
-        expectedMetrics.copy(pythonExecutionDurationMs = value,
-          memoryBytesSpilled = value, diskBytesSpilled = value))
+        expectedMetrics ++ Map("pythonExecutionDurationMs" -> value,
+          "memoryBytesSpilled" -> value, "diskBytesSpilled" -> value))
       checkEndMarkers(stream)
     }
   }
@@ -830,9 +827,10 @@ class BasePythonRunnerSuite extends SparkFunSuite with SharedSparkContext {
   }
 
   test("SPARK-59773: worker metrics accept a larger report with an additional field") {
-    val json = metricsReport.dropRight(1) + ",\"future\":\"" + ("x" * (70 * 1024)) + "\"}"
+    val json = metricsReport.dropRight(1) + ",\"futureMetric\":17" + (" " * (70 * 1024)) + "}"
     val stream = framedMetricsStream(json)
-    assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics)
+    val result = BasePythonRunner.readWorkerMetrics(stream)
+    assert(result == expectedMetrics + ("futureMetric" -> 17L))
     checkEndMarkers(stream)
   }
 }

python/pyspark/eval_handlers/_arrow.py

@@ -28,6 +28,7 @@
 from __future__ import annotations
 
 import itertools
+import time
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
@@ -56,10 +57,12 @@
     ArrowTableToRowsConversion,
     LocalDataToArrowConversion,
 )
+from pyspark.sql.pandas.serializers import ArrowStreamSerializer
 from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
 from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
 from pyspark.sql.types import StructField, StructType
 from pyspark.util import PythonEvalType, fail_on_stopiteration
+from pyspark.worker_metrics import WorkerMetrics
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -855,6 +858,7 @@ def __init__(
                     ),
                 )
             )
+        assert self._udf_infos, "At least one Arrow UDF expected here."
         self._col_names = [f"_{i}" for i in range(len(udfs))]
 
         # --- Input preparation ---
@@ -865,10 +869,28 @@ def __init__(
             for f in eval_conf.input_type  # type: ignore[union-attr]
         ]
 
+    @property
+    def serializer(self) -> ArrowStreamSerializer:
+        return ArrowStreamSerializer(
+            write_start_stream=True,
+            collect_timing=self._runner_conf.detailed_timing_enabled,
+        )
+
     def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
         import pyarrow as pa
 
         runner_conf = self._runner_conf
+        metrics = WorkerMetrics()
+        if runner_conf.detailed_timing_enabled:
+            phase_timers = (
+                metrics.measure("pythonInputPreparationTime"),
+                metrics.measure("pythonUDFExecutionTime"),
+                metrics.measure("pythonOutputPreparationTime"),
+            )
+            metrics.set("pythonNumTimingReports", 1)
+            metrics.set("pythonNumTimedBatches", 0)
+        else:
+            phase_timers = None
 
         @fail_on_stopiteration
         def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
@@ -879,42 +901,77 @@ def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
             with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
                 return list(pool.map(lambda row: udf_func(*row), rows))
 
-        for input_batch in data:
-            num_rows = input_batch.num_rows
-
-            # --- Input: Arrow -> Python columns ---
-            columns = [
-                (
-                    [conv(v) for v in ArrowTableToRowsConversion._to_pylist(col)]
-                    if conv is not None
-                    else ArrowTableToRowsConversion._to_pylist(col)
-                )
-                for col, conv in zip(input_batch.itercolumns(), self._arrow_to_py_converters)
-            ]
-            if not columns:
-                columns = [[_NoValue] * num_rows]
-
-            # --- Process: evaluate each UDF row-by-row ---
-            output_arrays = []
-            for udf_func, offsets, zero_arg, arrow_return_type, result_conv in self._udf_infos:
-                rows = (
-                    [() for _ in range(num_rows)]
-                    if zero_arg
-                    else list(zip(*[columns[o] for o in offsets]))
-                )
-                results = _evaluate_batch_udf(udf_func, rows)
-                verify_result_row_count(len(results), num_rows)
-
-                # --- Output: Python -> Arrow ---
-                converted = (
-                    [result_conv(r) for r in results] if result_conv is not None else results
-                )
-                try:
-                    arr = pa.array(converted, type=arrow_return_type)
-                except pa.lib.ArrowInvalid:
-                    arr = pa.array(converted).cast(
-                        target_type=arrow_return_type, safe=runner_conf.safecheck
+        # Keep phase totals and the batch count local, then publish them once per task.
+        input_duration_ns = 0
+        udf_duration_ns = 0
+        output_duration_ns = 0
+        num_timed_batches = 0
+        clock = time.perf_counter_ns
+        num_udfs = len(self._udf_infos)
+        try:
+            for input_batch in data:
+                if phase_timers is not None:
+                    num_timed_batches += 1
+                    input_start_ns = clock()
+                num_rows = input_batch.num_rows
+
+                # --- Input: Arrow -> Python columns ---
+                columns = [
+                    (
+                        [conv(v) for v in ArrowTableToRowsConversion._to_pylist(col)]
+                        if conv is not None
+                        else ArrowTableToRowsConversion._to_pylist(col)
                     )
-                output_arrays.append(arr)
-
-            yield pa.RecordBatch.from_arrays(output_arrays, self._col_names)
+                    for col, conv in zip(input_batch.itercolumns(), self._arrow_to_py_converters)
+                ]
+                if not columns:
+                    columns = [[_NoValue] * num_rows]
+
+                # --- Process: evaluate each UDF row-by-row ---
+                output_arrays: list[pa.Array] = []
+                for udf_func, offsets, zero_arg, arrow_return_type, result_conv in self._udf_infos:
+                    if phase_timers is not None and output_arrays:
+                        # Reuse the batch start for the first UDF, avoiding another clock read.
+                        # Later UDFs start input after the loop transition.
+                        input_start_ns = clock()
+                    rows = (
+                        [() for _ in range(num_rows)]
+                        if zero_arg
+                        else list(zip(*[columns[o] for o in offsets]))
+                    )
+                    if phase_timers is not None:
+                        udf_start_ns = clock()
+                        input_duration_ns += udf_start_ns - input_start_ns
+                    results = _evaluate_batch_udf(udf_func, rows)
+                    if phase_timers is not None:
+                        output_start_ns = clock()
+                        udf_duration_ns += output_start_ns - udf_start_ns
+                    verify_result_row_count(len(results), num_rows)
+
+                    # --- Output: Python -> Arrow ---
+                    converted = (
+                        [result_conv(r) for r in results] if result_conv is not None else results
+                    )
+                    try:
+                        arr = pa.array(converted, type=arrow_return_type)
+                    except pa.lib.ArrowInvalid:
+                        arr = pa.array(converted).cast(
+                            target_type=arrow_return_type, safe=runner_conf.safecheck
+                        )
+                    output_arrays.append(arr)
+                    # All UDF results for this batch are ready; time its assembly as output.
+                    if len(output_arrays) == num_udfs:
+                        output_batch = pa.RecordBatch.from_arrays(output_arrays, self._col_names)
+                    if phase_timers is not None:
+                        output_duration_ns += clock() - output_start_ns
+
+                yield output_batch
+                del output_batch
+        finally:
+            # Publish the count and completed intervals if the generator closes early or raises.
+            if phase_timers is not None:
+                metrics.set("pythonNumTimedBatches", num_timed_batches)
+                input_timer, udf_timer, output_timer = phase_timers
+                input_timer.add_duration_ns(input_duration_ns)
+                udf_timer.add_duration_ns(udf_duration_ns)
+                output_timer.add_duration_ns(output_duration_ns)