Spark batch scoring 20x slower than expected: the model was loaded inside every UDF call

Our nightly propensity scoring job applies a gradient boosting model (about 600 MB serialised, with its preprocessing pipeline) to 80 million customer rows. A benchmark on one machine scored 200,000 rows per second. On a 20-node Spark cluster we expected minutes. The job took 9 hours 20 minutes, and executors spent most of their time in what looked like I/O.

Reproducing it

The Spark UI showed tasks taking 3 to 6 minutes for 50,000 rows each. A py-spy profile of one Python worker:

  %Own   %Total  Function
  71.0%   71.0%  load (joblib/numpy_pickle.py)
  12.4%   12.4%  read (s3fs/core.py)
   9.8%    9.8%  predict (lightgbm/basic.py)

71 percent of time loading the model.

The UDF:

@pandas_udf("double")
def score(*cols: pd.Series) -> pd.Series:
    model = joblib.load(fsspec.open(MODEL_PATH).open())   # runs for every Arrow batch
    X = pd.concat(cols, axis=1)
    X.columns = FEATURES
    return pd.Series(model.predict_proba(X)[:, 1])

What we ruled out

• Data skew. Partition sizes were even.
• Cluster sizing. CPU was busy, but with deserialisation, not prediction.

Root cause

A scalar pandas UDF is called once per Arrow record batch, 10,000 rows by default. Loading the model inside the function meant loading 600 MB from object storage and deserialising it 8,000 times across the job, more than once per second per executor core. Prediction itself took a small fraction of the time.

What is happening under the hood

PySpark sends data from the JVM to Python workers in Arrow batches. For a scalar pandas UDF, your function runs on each batch independently and has no natural place for one-time setup. Variables defined at module level are created when the Python worker imports the function, but the worker may be reused or restarted, and large objects captured in closures are serialised with the function and sent to every task, which is also slow.

Two better patterns exist:

• Iterator of Series to iterator of Series pandas UDF: your function receives an iterator over all batches in a partition, so you can load the model once and then loop.
• mapInPandas: same idea at the DataFrame level, convenient when you need several columns and return several outputs.

MLflow's spark_udf uses the iterator pattern internally and caches the model per worker.

Environment

Minimal reproduction

Any scalar pandas UDF that loads a large object inside the function body shows load time proportional to the number of batches.

The fix

1. Iterator pandas UDF that loads once per partition, cached per Python worker.

from functools import lru_cache
from typing import Iterator, Tuple

@lru_cache(maxsize=1)
def get_model():
    with fsspec.open(MODEL_PATH, "rb") as f:
        return joblib.load(f)

@pandas_udf("double")
def score(batches: Iterator[Tuple[pd.Series, ...]]) -> Iterator[pd.Series]:
    model = get_model()                          # once per worker process
    for cols in batches:
        X = pd.concat(cols, axis=1)
        X.columns = FEATURES
        yield pd.Series(model.predict_proba(X)[:, 1])

scored = df.withColumn("propensity", score(*[F.col(c) for c in FEATURES]))

lru_cache keeps the model for the life of the Python worker, which Spark reuses across tasks by default (spark.python.worker.reuse=true).

2. Larger Arrow batches for this job, since the model is fast on big batches and memory allowed it:

spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", 50000)

3. Ship the model to local disk once. Executors download the model to local storage at startup through spark.files, so loading reads from disk rather than S3.

Alternatives and trade-offs

• mlflow.pyfunc.spark_udf if the model is in MLflow. Handles caching and environment for you.
• Broadcast the model. Works for small models; for hundreds of megabytes, broadcasting and deserialising per task can still be expensive.
• Score outside Spark with a dedicated batch inference service reading Parquet directly. Simpler for very large models on GPUs.

How to test the fix

On a 2-million-row sample: old UDF 14 minutes, new UDF 41 seconds. A test asserts that get_model is called at most once per worker by counting calls through a metric.

Confirming in production

• The nightly job now finishes in 28 minutes instead of 9 hours 20 minutes.
• The cluster runs for a fraction of the time, which cut the job's compute cost by about 90 percent.
• Scores are identical to the old job's output on a sampled comparison.

Regression guard

• Job duration alert at twice the trailing median.

• A code review guideline for UDFs: no I/O or heavy initialisation inside per-batch functions.

• Profiling of new batch jobs on a sample before they run at full size.

• The shared scoring library exposes one score_with_model(df, model_uri) helper built on the iterator pattern, and new jobs use it instead of writing their own UDFs.

Security implications

None directly. The executors' role can read only the model prefix and the input and output tables; caching the model on local disk is fine because executor volumes are ephemeral and encrypted.

Prevention checklist

• Never load models inside scalar UDF bodies.
• Use iterator pandas UDFs or mapInPandas and cache models per worker.
• Tune Arrow batch size for inference workloads.
• Profile a sample before scaling a job up.
• Compare per-row cost on Spark with a single-machine benchmark.

References

• PySpark pandas UDFs, iterator of Series — https://spark.apache.org/docs/latest/api/python/user_guide/sql/arrow_pandas.html#iterator-of-series-to-iterator-of-series
• PySpark mapInPandas — https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.mapInPandas.html
• Spark Arrow configuration — https://spark.apache.org/docs/latest/api/python/user_guide/sql/arrow_pandas.html#setting-arrow-batch-size
• MLflow spark_udf — https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#mlflow.pyfunc.spark_udf