suaroman

Writing 10 min read

Following up on the Spark 4.2.0 UDF hang

In my last post I left one problem unexplained: ordinary Python UDFs that never finished on cached or Parquet data. This weekend I took a closer look. Here is how I narrowed it down, what was going on, and where the fix stands.

By Gregory J. Suarez

This weekend, I decided to take a further look at the problem I left open at the end of Spark 4.2 turns on Arrow for ordinary Python UDFs. On Spark 4.2.0, an ordinary Python UDF never finished when its input came from a cached DataFrame or a Parquet file. I had stack traces from both sides, a few guesses, and a note to myself that it needed more troubleshooting.

The short version: it is a real bug in Spark 4.2.0, and it had already been found and fixed upstream. The fix is in 4.2.1, which is not released yet. The more interesting part turned out to be how a bug like this ended up in the 4.2.0 release without any test catching it. Along the way I also ran into a second, smaller problem that I don’t think anyone has reported yet.

This post follows the steps in the order I took them, including the ones where I got something wrong.

Step 1: make it small, and make it stop

A hang is awkward to study. Nothing fails, so there is no error message, and every attempt leaves a stuck JVM behind. So the first thing I wrote was a script that reproduces the problem with as little as possible and always exits. It reads 100 rows three ways, applies lambda x: x + 1, and a watchdog cancels the job if it has not finished in time. Stripped down, it is this:

import os, sys
os.environ.setdefault("PYSPARK_PYTHON", sys.executable)
from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.master("local[1]").getOrCreate()
plus_one = F.udf(lambda x: x + 1, "long")      # an ordinary UDF; Arrow is the 4.2 default

spark.range(100).select(plus_one("id")).collect()          # fine

df = spark.range(100).cache()
df.count()
df.select(plus_one("id")).collect()                         # never returns on 4.2.0
Spark 4.2.0
  range    OK    sum=5050  (2.0s)
  cache    HUNG  no result after 20s, job cancelled
  parquet  HUNG  no result after 20s, job cancelled

On Spark 4.1.3 all three finish.

Step 2: map the edges

With a script I could trust, I could ask a lot of small questions quickly. I turned it into a matrix of variations and ran each one on Spark 4.2.0 and on 4.1.3. The full results are further down. On 4.2.0 the pattern was:

  • Hangs: cached input, Parquet input, and ORC input, which I had not tried before.
  • Doesn’t hang: Arrow UDFs switched off, useArrow=False on the one UDF, a pandas_udf, an arrow_udf, the vectorized readers switched off, and everything on 4.1.3, even with Arrow UDFs switched on.
  • A surprise: in the first post I wrote that setting spark.sql.execution.arrow.pythonUDF.columnarInput.enabled to false did not help. That is only half true. It doesn’t help for cached input, but it does work for Parquet.
  • Another surprise: the same UDF applied to id * 2 instead of id doesn’t hang.

So something about where the data comes from and what the UDF is given decided it. Cache, Parquet and ORC have one thing in common: Spark can hand all three over as columns rather than rows.

Step 3: read the code the stack traces point to

The JVM stack trace I saved last time runs through a class called ColumnarArrowEvalPythonEvaluatorFactory. It is new in 4.2.0. It came with a change that lets a Python UDF take columnar data straight from the operator below it, instead of Spark turning the columns into rows and then back into Arrow columns for Python. A sensible optimization.

Inside it there are three paths. Two of them, one for cached and Parquet/ORC data and one for Arrow-backed data sources, hand the data to a new runner: the piece of JVM code that talks to the Python worker. The third path, used when a UDF’s argument is an expression like id * 2, uses the older runner. That already explained the id * 2 result.

Before any data, a runner sends the worker a short description of the job: some settings, then the list of UDFs and their arguments, then the function itself. For an ordinary Arrow UDF it also sends the input schema, the types of the columns that are coming. Side by side, this is where the two runners differ (condensed):

// The older runner (ArrowPythonRunner.scala): the schema goes with the settings
override protected def evalConf: Map[String, String] =
  if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) super.evalConf ++ Map("input_type" -> schema.json)
  else super.evalConf

// The new runner (ColumnarArrowPythonRunner.scala): the schema goes in front of the UDF list
override protected def writeUDF(dataOut: DataOutputStream): Unit = {
  if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
    PythonWorkerUtils.writeUTF(schema.json, dataOut)
  }
  PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas)
}

The 4.2.0 worker only looks for the schema in the settings. If this was the bug, it would also explain the rest of my list. SQL_ARROW_BATCHED_UDF is the type of an ordinary Arrow UDF, so pandas and Arrow UDFs never reach that line. With Arrow UDFs off, Spark uses a different operator altogether. And 4.1.3 has neither the new runner nor the move of the schema. The figure below follows the route for each case. Try “Cached” with the switch on and off.

The route Which runner talks to Python?
Input
Switch
Build
The path from the input, through the UDF operator, to one of two runners INPUT df.cache() can produce rows or columns ColumnarToRow added by the planner no converter THE UDF OPERATOR ArrowEvalPythonExec goes columnar if child.supportsColumnar ROW RUNNER turns rows into Arrow batches ✓ schema sent with the settings COLUMNAR RUNNER · NEW IN 4.2.0 sends the columns as they are ✗ schema written into the UDF list

 

 

 

The route follows the Spark 4.2.0 source: ArrowEvalPythonExec, the planner's ColumnarToRow rule, and the two runners. “Cached” means a cached DataFrame whose columns are all numbers or booleans; the cache only hands over columns for those. Every “measured” line is from a real run.

Step 4: watch the conversation

Reading code gives you a theory. I wanted to see the bytes. PySpark lets you choose which Python module a worker runs (spark.python.worker.module), so I wrote a small wrapper that records every byte the worker reads, and which worker function asked for it, and then hands over to the normal worker. I ran the plain query and the cached query and compared the two recordings.

What the worker read Same Spark, two runners, one misplaced field

settings (eval_conf) input schema UDF description the function itself

spark.range input, Spark 4.2.0 (row runner)

 

THE JVM WROTE THE WORKER READ

 

Cached input, Spark 4.2.0 (columnar runner)

 

THE JVM WROTE THE WORKER READ

 

Real bytes, captured with a small wrapper around the PySpark worker that records every read. Each lane starts at the settings section and covers the UDF description that follows; bytes are shown in hex, text as characters, and long runs are folded into “…”. In the hung lane the worker is stuck inside one read after 21 bytes, so the faded bytes after that point were never captured: they are what the JVM sends next, taken from the fixed run, which uses the same schema and the same function.

Here is what the worker on the cached query does. The settings section says zero entries, which is correct: the new runner put nothing there. The next four bytes are the schema’s length, 86, but the worker expects the number of UDFs, so it now believes there are 86 of them. It reads the start of the schema text, {"ty, as an argument count of just over two billion, a quote mark as “yes, this argument has a name”, and stru as the length of that name: 1,937,011,317 bytes. Then it waits for 1.9 GB that will never arrive. The JVM has already sent everything it has, about 6 KB for this task, and is waiting for results. Neither side has a timeout at that point, so both wait forever.

That also corrects something I wrote in the first post. I said the worker was “still waiting to be told which function to run”. It had been told. It misread the message and was waiting for the rest of a 1.9 GB argument name.

Step 5: how did the bug reach a release without being caught?

The git history of Spark’s source code tells a simple story. Two changes went into Spark in April, three days apart. The first moved the schema out of the UDF list and into the settings, on both the JVM and the Python side. The second added the new runner, which writes the schema the old way. Each change made sense against the code it was written for. They touched different files, so git merged them without complaint, and the new runner’s tests used pandas UDFs, which never reach the line that writes the schema. That kind of problem is hard to catch in review, because each change looks right on its own.

Then I searched Spark’s JIRA for the runner’s name, and found that this had already been reported and fixed: SPARK-58241, by Liang-Chi Hsieh, on July 21 and 22. The fix is the one you would expect: the new runner sends the schema with the settings. It went into the 4.2 branch eleven days after 4.2.0 was tagged, and is marked for 4.2.1.

I would have found that sooner had I searched JIRA for the class name before reading the code. On the other hand, because I had read the code first, I understood the fix the moment I saw it.

How it reached a release Two sensible changes, three days apart

code change release report and fix me

  1. SPARK-56340: the input schema moves into the settings

    The JVM now sends the UDF input schema with the other settings, and the Python worker stops looking for it in the UDF description.

  2. 3 days later

    SPARK-56350: Python UDFs can take columns directly

    A new runner sends columnar input straight to Python. For ordinary Arrow UDFs it still writes the schema in the old place. The two changes touch different files, so they merge without a conflict, and the new tests use pandas UDFs, which skip that line.

  3. 3 months later

    v4.2.0 is tagged

    Both changes are in it.

  4. PySpark 4.2.0 on PyPI

    Arrow-optimized Python UDFs are now the default (SPARK-54555).

  5. 1 week later

    SPARK-58241 is reported

    Found on the main branch, where a new Arrow-based cache made the hang easy to hit.

  6. Fixed, and backported to the 4.2 branch

    The new runner sends the schema with the settings. Fix version: 4.2.1.

  7. 2 months later

    v4.2.1-rc1 is cut

    The first release candidate with the fix. Not on PyPI yet.

  8. My first post

    An ordinary UDF on cached or Parquet input never finishes. I could not explain it.

  9. This weekend

    This post.

All dates are 2026, from the Apache Spark git history, JIRA and PyPI.

One detail in the upstream notes did not match what I saw. The note on the 4.2 backport says that on the 4.2 branch the bug is only reachable through Arrow-backed data source connectors. In my runs a plain df.cache() or a Parquet read was enough, because the path for cached and Parquet/ORC data uses the same new runner. That matters for who should look out for 4.2.1: anyone on 4.2.0 who applies an ordinary Python UDF directly to columns read from Parquet or ORC, or from a cached DataFrame of numeric columns.

Step 6: the switch that didn’t switch

That left the question from Step 2: why does turning columnarInput.enabled off help for Parquet but not for a cache? The UDF operator decides like this, in 4.2.0 and still on Spark’s main branch today:

override def supportsColumnar: Boolean =
  child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled

override protected def doExecute(): RDD[InternalRow] = {
  if (child.supportsColumnar) {   // asks its input, not itself
    // ...the columnar path

With the switch off, the operator tells the planner that it wants rows. A Parquet scan can only produce columns, so the planner puts a ColumnarToRow converter in between, and the operator’s input now produces rows. A cache can produce either, so the planner adds nothing. When the query runs, the operator asks its input whether it can produce columns, hears yes, and takes the columnar path anyway.

On 4.2.1-rc1 this no longer causes a hang, because the columnar path now works. But the switch still does nothing for cached data: a thread dump taken in the middle of a query shows the columnar path with the switch on and with it off. The fix is small: ask supportsColumnar instead of child.supportsColumnar. I could not find a report of this one.

Step 7: try the fix

To check that I had the whole story, I built Spark’s SQL module from the 4.2.0 source with both changes:

 // ColumnarArrowPythonRunner.scala: send the schema with the settings, like the older runner
+  override protected def evalConf: Map[String, String] =
+    if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) super.evalConf ++ Map("input_type" -> schema.json)
+    else super.evalConf
+
   override protected def writeUDF(dataOut: DataOutputStream): Unit = {
-    if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
-      PythonWorkerUtils.writeUTF(schema.json, dataOut)
-    }
     PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas)
   }

 // ArrowEvalPythonExec.scala: let the switch decide, and decide once
-  override def supportsColumnar: Boolean =
+  override lazy val supportsColumnar: Boolean =
     child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled
   ...
-    if (child.supportsColumnar) {
+    if (supportsColumnar) {

The first change matches the upstream fix. The lazy val in the second makes the operator decide once, so the planner and the running query cannot disagree if someone changes the setting in between. I swapped the new jar into a pip-installed PySpark 4.2.0 and ran the matrix on six builds: 4.1.3, 4.2.0 from PyPI, my own build of 4.2.0 with no changes, 4.2.0 with only the switch fix, 4.2.0 with both fixes, and the official 4.2.1 release candidate.

The test matrix 15 cases, six builds, 90 runs

finished, right answer no result after 30 s

Case 4.1.3 Arrow UDFs switched on 4.2.0 from PyPI 4.2.0, rebuilt my own build, no changes 4.2.0 + switch fix only the second fix 4.2.0 + both fixes my test build 4.2.1-rc1 official release candidate
01 spark.range input (the control) 3.0 s 2.9 s 2.9 s 3.0 s 2.8 s 2.9 s
02 Cached DataFrame 2.5 s hung hung hung 2.1 s 2.3 s
03 Parquet file 1.9 s hung hung hung 1.6 s 1.6 s
04 ORC file 2.7 s hung hung hung 2.2 s 2.4 s
05 Cached, Arrow UDFs off 2.3 s 1.7 s 1.7 s 0.9 s 1.9 s 2.0 s
06 Parquet, Arrow UDFs off 1.7 s 2.2 s 2.3 s 1.0 s 1.4 s 1.6 s
07 Cached, useArrow=False on the UDF 1.4 s 1.9 s 1.9 s 1.0 s 1.5 s 1.6 s
08 Cached, pandas_udf 1.9 s 1.3 s 1.4 s 1.8 s 2.7 s 2.8 s
09 Cached, arrow_udf 3.5 s 1.6 s 1.9 s 1.7 s 1.7 s 1.9 s
10 Cached, columnarInput switch off 2.0 s hung hung 1.8 s 2.4 s 2.2 s
11 Parquet, columnarInput switch off 1.9 s 1.7 s 1.9 s 1.9 s 3.1 s 2.4 s
12 Cached, cache vectorized reader off 3.4 s 1.5 s 1.8 s 1.2 s 2.0 s 2.5 s
13 Parquet, Parquet vectorized reader off 2.9 s 1.6 s 1.9 s 1.3 s 1.9 s 2.9 s
14 Cached, UDF input is id * 2 2.9 s 1.5 s 1.8 s 1.2 s 1.8 s 2.7 s
15 Cached, one column is a string 1.5 s 1.1 s 1.2 s 1.2 s 1.2 s 1.1 s
Each cell is one run: 100 rows, local[1], a 30-second watchdog. A cell is green only if the query finished and returned the right sum. 4.1.3 runs with Arrow UDFs switched on, so it takes the same Arrow path. “Rebuilt” is my own build of 4.2.0 with no changes, there to show that building it myself is not what fixes anything.

Every build without the first fix hangs on cached, Parquet and ORC input, and every build with it finishes. The switch-only build shows the second fix on its own: there, case 10 (cached, switch off) finishes while the plain cached case still hangs. I also wrote a small regression test with six cases, each with a 60-second watchdog so that a regression fails instead of hanging the test run. All six fail on 4.2.0 and pass on my build and on 4.2.1-rc1.

If you are on Spark 4.2.0

Until 4.2.1 is out, any one of these avoids the hang:

  • Cached input: set spark.sql.inMemoryColumnarStorage.enableVectorizedReader to false.
  • Parquet or ORC input: set spark.sql.execution.arrow.pythonUDF.columnarInput.enabled to false (for these it works), or switch off that format’s vectorized reader.
  • One UDF: useArrow=False puts just that UDF back on pickle.
  • The whole session: spark.sql.execution.pythonUDF.arrow.enabled=false, which also gives up the speedup I measured last time.
  • Or use a pandas_udf or arrow_udf, which are not affected.

There is one more option, if you run your own Spark: apply the fix to 4.2.0 yourself, which is what I did to test it in Step 7. The upstream fix changes one file, so you build only Spark’s SQL module from the v4.2.0 source with that change (about two minutes on my machine; Maven downloads the other modules) and replace spark-sql_2.13-4.2.0.jar with the result. The catch is that you now run a custom build. The jar has to be replaced everywhere Spark runs, driver and executors alike, and you have to remember to undo it when you upgrade. On a managed platform you usually can’t replace Spark’s jars at all. For most people, one of the settings above and then 4.2.1 is the simpler path.

What’s next

I’d like to pass two things back to the Spark project: a note on SPARK-58241 that plain cached and Parquet input are enough to trigger it on 4.2.0, together with the regression test, which runs on the 4.2 branch as it is; and a new ticket for the switch.

How I tested

Same machine as before, Spark in local mode with one core, 100 rows per case, and a 30-second watchdog on every query. Spark 4.2.0 and 4.1.3 came from PyPI, and 4.2.1-rc1 from the Apache release candidate area, with its checksum checked. OpenJDK 21, Python 3.11, PyArrow 25.0.1, pandas 2.3.3. My builds are Spark’s sql/core module built with Maven from the v4.2.0 tag. The hang looked the same, with the same stacks, on Java 17 with Python 3.11 (the first post) and on Java 21 with Python 3.14 (an earlier lab), so it doesn’t depend on the setup.

Sources