[SPARK-59993][PYTHON] Collect worker and Arrow-optimized Python UDF timing metrics by Spenserrrr · Pull Request #59246 · apache/spark

Spenserrrr · 2026-10-06T00:30:12Z

What changes were proposed in this pull request?

Reset WorkerMetrics per task and decode its numeric report into SQL metrics.
Report seven Arrow scalar UDF phases: raw receive, Arrow deserialization, input preparation, UDF execution, output preparation, Arrow serialization, and raw send; count timed tasks, batches, and pipelined tasks.
A default-off config gates timing. The Arrow handler owns phase timers and accumulates the batch count locally; its serializer times raw I/O beneath buffers.

Why are the changes needed?

pythonProcessingTime covers the worker loop, hiding its phases.

Does this PR introduce any user-facing change?

Yes. Enabling the config adds seven durations and two counts to Arrow UDF SQL metrics.
Receive/send measure raw I/O during Arrow processing; buffer work stays in Arrow ser/de. Prefetched input and final flush are excluded.

How was this patch tested?

Focused tests (52 tests, 54 subtests), 20 counter lifecycle cases, Ruff, format, and mypy passed; a B20 worker benchmark measured 0.11-0.12 us per batch saved; JVM integration awaits CI.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Codex (GPT-6).

Spenserrrr · 2026-10-06T01:08:24Z

Hi @gaogaotiantian, could you take a look when you have time? This adds scalar pandas UDF phase timings and generic worker metric handling.
@Yicong-Huang, feel free to take a look if interested, especially for the integration with your pandas eval-handler changes. Thanks a lot!

HyukjinKwon

Review summary

The current revision addresses the earlier inline feedback. It uses a worker-level WorkerMetrics collector, renames the counter to increment, adds to_dict, and no longer threads a metrics argument through read_udfs. Three P3 items remain. The required-field check that moved into handleMetricsData no longer has a test. Four of the five phase metrics are asserted only with >= 0, which also passes if they are never reported. The readWorkerMetrics Scaladoc describes a tolerance the decoder no longer has. None of these blocks the change. Separately, the per-batch timing is unconditional on the default Arrow-optimized UDF path. Given the small-batch slowdown in the PR description, that tradeoff is raised as a question below.

Cross-cutting observations

Splitting SPARK-59773's single validated decoder into a generic int64 decoder plus a consumer-side required-field check left both halves of the old contract under-specified. The decoder's doc still reads as tolerant of other fields, and the required-field check lost its only test.
Separation: The split itself is the PR's accepted design. Each half needs its own small correction (a doc rewrite, a test), so the shared origin establishes no common owner or coordinated repair.
The new tests pin the generic layers well: the int64 decoder, WorkerMetrics, and the serializer timers in isolation. They do not pin the two integration points where the metrics actually flow: JVM handleMetricsData's required-field enforcement and the worker's read_udfs phase wiring.
Separation: The two integration points are in different processes and suites and guard different invariants, so the shared pattern doesn't merge the repairs or change their P3 severity.

Findings

3 total: 0 P0, 0 P1, 0 P2, 3 P3.

Nit (P3)

Required worker-metric check lost its only test — core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:130 — see inline.
Four phase metrics have no failing test signal — python/pyspark/sql/tests/arrow/test_arrow_python_udf.py:602 — see inline.
readWorkerMetrics doc misstates the int64-only contract — core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:220 — see inline.

PR description suggestions

The 'Does this PR introduce any user-facing change?' section says No, but the PR adds new SQL metrics shown in the SQL UI: five phase timings and two counts on ArrowEvalPython, plus pythonNumPipelinedTasks on every Python operator. Consider listing them there.

Decision challenges

Should phase timing be always-on for the default Python UDF path?

Phase timing is enabled for every non-legacy SQL_ARROW_BATCHED_UDF task. That is the default for regular Python UDFs, since spark.sql.execution.pythonUDF.arrow.enabled defaults to true, and there is no way to turn it off. For k UDFs, each Arrow batch now enters 4 + 3k timer scopes, each with two perf_counter_ns calls. The PR description's own A/B reports +22.5% median job time with one-row batches and consistent slowdowns at 100-row batches, with mixed results at 1,000 and 10,000 rows. Is that cost acceptable for always-on diagnostics? If so, could the PR say so explicitly? If not, would you gate collection behind a config, or reduce the per-batch timing work (for example, fewer scopes per batch)?

Generated by Omnigent on Databricks.

HyukjinKwon · 2026-10-08T20:39:12Z

    checkEndMarkers(stream)
  }

-  test("SPARK-59773: worker metrics require all current fields") {


Nit (P3): This test was the only coverage for the missing-field failure. That check now lives in requiredMetric inside handleMetricsData (PythonRunner.scala:895-903), and no test reaches it. The remaining tests call readWorkerMetrics, which accepts any subset of fields. If requiredMetric regressed to getOrElse(name, 0L), every test would still pass, and boot/init/total/processing times or spill bytes would silently become 0. Could the six lookups move into a small private[python] helper on BasePythonRunner? Then this test could drop each field in turn and assert the Missing Python worker metric: <name> failure again.

Verification:

Regression: For each of the six required fields, a report that omits it is rejected with a SparkException whose message names the field.

Behavior: A complete report still yields the same required values, so boot, init, total and processing times and spill increments are unchanged.

HyukjinKwon · 2026-10-08T20:39:13Z

+                        metrics.apply("pythonNumTimedBatches").value(), (num_rows + 1) // 2
+                    )
+                    phases = [metrics.apply(name).value() for name in phase_names]
+                    self.assertTrue(all(value >= 0 for value in phases))


Nit (P3): value >= 0 also holds for a metric the worker never reported, because a timing SQLMetric that was never updated reads 0. If collect_timing were dropped from read_udfs (worker.py:1835-1838), or the input_timer/output_timer scopes were removed, this test would still pass. ArrowStreamTimingTests builds the serializer with collect_timing=True itself, so it doesn't cover that wiring either. One deterministic option is to assert not metrics.apply(name).isZero() for each phase on a non-empty task, since any worker update clears isZero, even a value of 0. Then add a patched-perf_counter_ns check that the input, UDF and output scopes each accumulate time.

Verification:

Regression: On the non-legacy Arrow-batched UDF path, each of the five phase metrics is observed as reported by the worker for a non-empty task.

Behavior: With a controlled clock, the input preparation, UDF execution and output preparation phases each accumulate a positive duration and exclude the work done between their scopes.

HyukjinKwon · 2026-10-08T20:39:13Z


-  /** Read and validate worker metrics after METRICS_DATA; ignore additional fields. */
-  private[python] def readWorkerMetrics(stream: DataInputStream): WorkerMetrics = {
+  /** Decode all numeric fields after METRICS_DATA without knowing their metric names. */


Nit (P3): This reads as if non-numeric or non-integral fields were skipped, which is what the old 'ignore additional fields' doc promised. But the loop below throws Invalid Python worker metric: <key> for any value that is not an int64 integer, including 42.0 and true. Someone adding a string or float diagnostic to the worker report would make every UDF task fail. Maybe something like: /** Decode every field after METRICS_DATA as an int64 by name; any other value fails the report. */

Verification:

Inspection: The decoder's doc comment describes the same accepted value domain and failure behavior that the decoder enforces.

gaogaotiantian

I have some questions about how we collect time. After we agreed on that, I can do a final scan of the full PR.

gaogaotiantian · 2026-10-06T21:50:31Z

-  /** Read and validate worker metrics after METRICS_DATA; ignore additional fields. */
-  private[python] def readWorkerMetrics(stream: DataInputStream): WorkerMetrics = {
+  /** Decode all numeric fields after METRICS_DATA without knowing their metric names. */
+  private[python] def readWorkerMetrics(stream: DataInputStream): Map[String, Long] = {


Is there a restriction here for all the values to be the same type? Is it possible for us to emit string in the future?

Currently, all the metrics are numeric because the values are fed into SQLMetrics, and that accumulates only Long values. The downstream QPL operator map also uses string metric names with int64 values. Thus, I let readWorkerMetrics only accept Long. A string value would require a separately designed path.

gaogaotiantian · 2026-10-10T03:10:10Z

+            try:
+                for input_batch in data:
+                    input_start_ns = clock()
+                    # --- Input: Arrow -> Python columns ---


This comment should either be before clock() or after increment. It's weird to have it in the middle.

Moved the input-section comment before clock() so it labels the interval being timed. Thanks for catching this!

gaogaotiantian · 2026-10-10T03:19:02Z

+                        if output_arrays:
+                            # One boundary ends the previous output and starts the next input.
+                            input_start_ns = clock()
+                            output_duration_ns += input_start_ns - output_start_ns


I don't quite understand what are we trying to do here. What's this value? output_start_ns is from the previous iteration of udf func? This is super difficult to follow. Why do we need this? Doesn't this include some random time outside of this function?

Yeah it is a bit hard to follow. It reused one clock reading to end the previous UDF’s output and start the next UDF’s input. I do this to save a clock read, and the tradeoff is that it includes the loop transition to output time. I removed the placeholder output_start_ns = input_start_ns and now record each UDF’s output time in its own iteration. With one UDF, the worker still makes four clock reads per batch. With multiple UDFs, it uses one additional read to leave the loop time out. The final UDF’s output includes batch assembly. Do you think this is better?

…readers

Spenserrrr marked this pull request as ready for review October 6, 2026 01:07