← PySpark

▶ Play from here · 59 min

Starts this lesson and continues through 7 more to the end of the course.

UDFs, Arrow, and pandas UDFs

The hierarchy

When you need a transformation, there is a strict preference order, and it is about the boundary from lesson 1:

1. Built-in function        → JVM only. No crossing. Optimiser can see it.
2. Higher-order / SQL expr  → still JVM.
3. pandas UDF (vectorised)  → crosses, but in Arrow batches
4. Plain Python UDF         → crosses, row by row
5. RDD mapPartitions        → crosses, and you lose the optimiser entirely

Always exhaust level 1 before descending. The pyspark.sql.functions module is large and people reach past it constantly — regexp_extract, split, explode, transform, aggregate, from_json, date_trunc, coalesce, when/otherwise cover an enormous amount of what people write UDFs for.

Why a plain Python UDF is slow — the mechanism

Three costs, and you should be able to name them:

1. Serialisation per row, out of the JVM into a Python worker and back.
2. A separate Python process per executor core, with its own memory, outside the JVM heap.
3. The optimiser goes blind. A UDF is an opaque function. Catalyst can't see inside it, so it can't push the predicate down, prune columns based on it, or reorder around it.

⚠️ Point 3 is the one people miss, and it's often the biggest. A UDF in a filter prevents predicate pushdown — so instead of the scan skipping partitions, Spark reads everything and then filters in Python. The UDF's own cost is dwarfed by the data you now read unnecessarily.

# BAD: opaque predicate, no pushdown, reads everything
df.filter(is_valid_region_udf(col("region")))

# GOOD: the optimiser can see this and push it to the scan
df.filter(col("region").isin("EU", "US", "APAC"))

Arrow — what it is and the configuration

"Apache Arrow is an in-memory columnar data format that is used in Spark to efficiently transfer data between JVM and Python processes."

The documented configuration keys (Apache Arrow in PySpark, verified 2026-08-12):

⚠️ Version sensitivity — read this before you copy a config from a blog post. These defaults have changed across Spark releases. The current documentation states Arrow is enabled by default for pandas conversion, and Arrow optimisation for regular Python UDFs is a separate, more recent setting. Check the docs for the Spark version your Glue or EMR release actually ships — those are usually behind the current release (lesson 5).

⚠️ The fallback setting is a trap in disguise. Silent fallback to the non-Arrow path means a job that quietly gets much slower without failing. If you depend on Arrow for your performance budget, you want to know when it falls back, not to sail past it.

Supported types are not universal:

"Currently, all Spark SQL data types are supported by Arrow-based conversion except ArrayType of TimestampType. MapType and ArrayType of nested StructType are only supported when using PyArrow 2.0.0 and above."

And the documented minimums: "For usage with pyspark.sql, the minimum supported versions of Pandas is 2.2.0 and PyArrow is 18.0.0."

pandas UDFs — the four types

A pandas UDF is vectorised: it receives a batch as a pandas.Series rather than a value per row. The four documented signatures:

⚠️ The Series-to-Scalar constraint is significant. The docs state it "does not support partial aggregation and all data for a group or window will be loaded into memory. Also, only unbounded window is supported."

All data for a group loaded into memory — so a skewed group will OOM the executor. If one key has 50 million rows, that group must fit in memory. Lesson 4 is about exactly that.

Iterator variants are the underused ones. They let you initialise something expensive once per batch-iterator rather than per batch — loading a model, opening a connection — which is often the difference between a viable and an unviable design.

pandas Function APIs

Distinct from pandas UDFs, and worth knowing by name:

mapInPandas is the one to remember — it's the only one here that can change the output row count, and it streams rather than materialising a whole group. When you genuinely need Python over large data, this is usually the right shape.

⚠️ Note that two of the three explicitly load an entire group into memory. That constraint is documented, not incidental, and it makes group size a correctness-of-design question rather than a tuning detail.

Choosing, in practice

Testing the assumption

Never argue about UDF cost — measure it. The comparison that settles it:

import time
def timed(label, fn):
    t = time.time(); fn(); print(label, round(time.time() - t, 1), "s")

timed("builtin",    lambda: df.withColumn("y", upper(col("x"))).write.mode("overwrite").parquet(a))
timed("pandas_udf", lambda: df.withColumn("y", pandas_upper(col("x"))).write.mode("overwrite").parquet(b))
timed("python_udf", lambda: df.withColumn("y", python_upper(col("x"))).write.mode("overwrite").parquet(c))

Write to storage rather than count() so you're measuring the real path, and use the same data each time. I'm deliberately not quoting a speedup multiplier here — published figures vary hugely with data shape, types, and version, and a number I can't reproduce for your workload is worse than no number. Run it on your data.

Check yourself

1. Name the three costs of a plain Python UDF.
2. Why is a UDF inside filter() especially bad?
3. What does spark.sql.execution.arrow.maxRecordsPerBatch control, and what's the documented default?
4. Which pandas UDF type loads all data for a group into memory, and why does that matter?
5. Which pandas Function API can change the number of output rows?

Answers

1. Per-row serialisation across the JVM/Python boundary; a separate Python process per executor core with memory outside the JVM heap; and loss of optimiser visibility, since Catalyst cannot see inside an opaque function.
2. Because it blocks predicate pushdown. Instead of the scan skipping partitions, Spark reads everything and filters in Python — the extra I/O usually costs more than the UDF itself.
3. The maximum number of rows per Arrow record batch. Documented default: 10,000.
4. Series-to-Scalar — the docs state it "does not support partial aggregation and all data for a group or window will be loaded into memory". A skewed key can therefore OOM an executor.
5. mapInPandas — documented as able to "return the output of arbitrary length". applyInPandas and the cogrouped variant both load a whole group into memory.

Teaching this section

• Lead with the hierarchy, not with Arrow. People want the pandas-UDF trick; the higher-value habit is exhausting pyspark.sql.functions first. Ask the room to name ten built-ins — most struggle past six, which is the point.
• Demo: the predicate-pushdown difference. Show explain() with a UDF in the filter and with an isin, and point at where PushedFilters disappears.
• Good discussion: "Which of our UDFs could be a built-in?" Bring three real ones from the codebase; usually at least two can.
• Watch for: copying Arrow config from blog posts. Have people check the docs for their Spark version — the version pinned by their Glue or EMR release, not the latest.
• Timing: 25 min presented, 15 hands-on with explain().