Spark Core · Module 2 · Deep dive
Arrow and vectorised execution
Why pandas UDFs are so much faster than plain Python UDFs, what Apache Arrow actually does, and the three different things people mean by "vectorised" in Spark.
Row layout vs columnar layout
The same three-row table can sit in memory two ways. Row layout keeps each record together, which suits reading or writing whole records. Columnar layout keeps each column together, which suits analytics: most queries touch a few columns of many rows.
amount.
- Fewer bytes read: a query on
amounttouches only theamountblock. - CPU-friendly: values of one type sit next to each other, so the CPU cache is used well and one instruction can process several values at once.
- Better compression: similar values side by side compress well. That's why Parquet and ORC files are columnar.
What Apache Arrow is
Apache Arrow is a standard columnar format for data in memory, defined the same way for every language. A batch of Arrow data (a record batch) is a set of column buffers plus a schema. Because Java, Python, C++ and others all understand the same layout, they can hand data to each other without converting it value by value, often without copying it at all.
Parquet is columnar on disk; Arrow is columnar in memory. pandas and NumPy are columnar in memory too, which is why converting Arrow to pandas is cheap.
Crossing the JVM ↔ Python boundary
Spark's engine runs in the JVM; your Python function runs in separate Python worker processes on each executor. Every Python UDF means moving data across that process boundary and back. How it's moved makes the difference.
| Plain Python UDF | pandas UDF (Arrow) | |
|---|---|---|
| Unit of transfer | One row, pickled | One Arrow record batch, ~10,000 rows by default |
| Conversions | Every value converted to a Python object and back | Column buffers handed over; pandas wraps them |
| Your function is called | Once per row, through the Python interpreter | Once per batch, with whole columns as pandas Series |
Plan node in explain() | BatchEvalPython | ArrowEvalPython |
Batches of rows are buffered in the Python worker's memory, which counts against the executor's overhead, not its heap. If pandas UDF stages get containers killed, lower the batch size or raise overhead (see Module 1).
What "vectorised" means
Vectorised code applies one operation to a whole array of values at once, instead of looping over values one at a time in an interpreted language. In pandas and NumPy, s * 1.05 runs a single loop in compiled C over the whole column, often using CPU instructions that process several numbers per step (SIMD). A plain Python function pays the interpreter's overhead on every single value.
import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.functions import udf, pandas_udf
@udf("double")
def tax_row(amount): # called once per row, in the Python interpreter
return amount * 1.05
@pandas_udf("double")
def tax_vec(amount: pd.Series) -> pd.Series: # called once per batch of ~10,000 rows
return amount * 1.05 # one compiled loop over the whole array
df.withColumn("t", tax_row("amount")) # slowest: pickle + interpreter per row
df.withColumn("t", tax_vec("amount")) # much faster: Arrow + vectorised NumPy
df.withColumn("t", F.col("amount") * 1.05) # fastest: no Python at alls.apply(lambda x: …) or loops over rows in Python, you get Arrow's cheaper transfer but still pay the interpreter per value. The big win needs whole-column pandas or NumPy operations inside the function.Where Spark uses Arrow
| API | Function shape | Use it for |
|---|---|---|
df.toPandas() / spark.createDataFrame(pdf) | Whole DataFrame, via the driver | Moving small results between Spark and pandas |
| pandas UDF, Series → Series | def f(s: pd.Series) -> pd.Series | Column transforms in select / withColumn |
| pandas UDF, Iterator of Series | def f(it: Iterator[pd.Series]) -> Iterator[pd.Series] | Expensive setup once per partition, e.g. loading a model |
| pandas UDF, Series → scalar | def f(s: pd.Series) -> float | Custom aggregations in groupBy().agg(); each group in memory |
mapInPandas | Iterator[pd.DataFrame] → Iterator[pd.DataFrame] | Whole-row logic over batches, no grouping |
applyInPandas | pd.DataFrame → pd.DataFrame, per group | Per-group logic; whole group in one worker |
| Arrow-optimised Python UDF | @udf(…, useArrow=True), row-wise function | Keeping existing row-wise code but speeding up the transfer (Spark 3.5+) |
pandas UDFs, mapInPandas and applyInPandas always use Arrow. The arrow.pyspark.enabled setting controls only toPandas and createDataFrame from pandas.
Configs worth knowing
| Config | Default | Effect |
|---|---|---|
spark.sql.execution.arrow.pyspark.enabled | false* | Use Arrow for toPandas and createDataFrame(pdf) |
spark.sql.execution.arrow.pyspark.fallback.enabled | true | Silently fall back to the slow path if Arrow fails (see production) |
spark.sql.execution.arrow.maxRecordsPerBatch | 10000 | Rows per Arrow batch for pandas UDFs; lower it if Python workers run out of memory |
spark.sql.execution.pythonUDF.arrow.enabled | false | Use Arrow for all regular Python UDFs (Spark 3.5+) |
spark.sql.parquet.enableVectorizedReader | true | Read Parquet in columnar batches (see below) |
* Off in open-source Spark; enabled by default on Databricks.
Three things called "vectorised" in Spark
| Term | What runs on batches of columns | Where you'll meet it |
|---|---|---|
| pandas UDFs / Arrow | Your Python function, with NumPy and pandas doing the loops | PySpark code |
| Vectorised Parquet / ORC reader | Decoding files into column batches (4,096 rows by default) instead of row by row | Every scan stage; shows as ColumnarToRow in explain() when batches are turned into rows for the next operators |
| Photon | The whole query engine, written in C++, processing column batches with SIMD | Databricks; plan nodes named Photon… |
For contrast, open-source Spark's Tungsten whole-stage codegen fuses the operators of a stage into one compiled Java function that processes one row at a time in a tight loop. It's fast because it removes per-operator overhead, but it's row-oriented, not vectorised. That's the main design difference Photon changes.
Say it in an interview
Quick check
Does setting arrow.pyspark.enabled=false make pandas UDFs stop using Arrow?
mapInPandas and applyInPandas always use Arrow. That setting only affects toPandas() and createDataFrame from a pandas DataFrame.Your pandas UDF calls s.apply(my_func). Is it vectorised?
apply calls a Python function per value, so you still pay interpreter overhead per row. Rewrite with whole-Series operations (arithmetic, np.where, .str methods) to get the real speed-up.A pandas UDF stage keeps getting containers killed. Which Arrow setting can help?
spark.sql.execution.arrow.maxRecordsPerBatch: smaller batches mean less Python-worker memory per task. Raising executor memory overhead also helps.