PySpark DataFrames, shuffles, and fast UDFs: pandas_udf + Arrow beat Python UDFs 10x

LEARN · DISTRIBUTED ML WITH PYSPARK

PySpark is often introduced as “Spark with Python,” but that description hides the most important performance detail: Python is not where Spark does most of its work.

The Spark engine is written primarily in Scala and runs on the JVM. When you use DataFrame operations, Spark can optimize your query, move data efficiently, and execute many transformations without involving Python at all.

Performance problems usually appear when:

Data has to move between machines.
A query creates large intermediate datasets.
Python code forces expensive serialization boundaries.
A user-defined function prevents Spark optimizations.

This article walks through the practical performance model behind PySpark DataFrames, explains why shuffles hurt, and shows why pandas_udf with Apache Arrow can outperform normal Python UDFs by an order of magnitude.

The execution model: DataFrames are more than tables

A PySpark DataFrame is not just a collection of Python objects. It is a logical query plan that Spark analyzes and optimizes before execution.

For example:

from pyspark.sql import functions as F

filtered = (
    df
    .filter(F.col("country") == "DE")
    .groupBy("category")
    .agg(F.sum("sales").alias("total_sales"))
)

Spark can inspect this plan and decide:

Which columns are actually needed.
Where filters should be applied.
How joins should be executed.
How many partitions should be used.

This optimization layer is one of the biggest reasons DataFrame APIs are usually faster than manually processing data.

The key optimizer component is Catalyst. It can rewrite logical plans into more efficient physical plans before Spark runs anything.

Narrow versus wide transformations

Spark transformations generally fall into two categories.

Narrow transformations

A narrow transformation does not require data movement between executors.

Examples:

select
filter
withColumn
simple expressions

A partition can be processed independently.

Example:

cleaned = (
    df
    .filter(F.col("amount") > 0)
    .select("customer_id", "amount")
)

Each executor processes its own partition.

Wide transformations

A wide transformation requires data exchange between executors. This operation is called a shuffle.

Common shuffle operations:

groupBy
join
distinct
orderBy
window functions

Example:

summary = (
    df
    .groupBy("customer_id")
    .agg(
        F.sum("amount").alias("revenue")
    )
)

Before Spark can calculate the sum per customer, all rows with the same customer_id must end up in the same partition.

That means Spark must:

Serialize rows.
Write shuffle files.
Transfer data over the network.
Read the shuffled data.
Perform the aggregation.

On large datasets, network transfer is often the limiting factor.

A practical DataFrame workflow

A typical analytics job might include filtering, aggregation, and window calculations.

Example dataset:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("sales-demo").getOrCreate()

sales = spark.createDataFrame(
    [
        ("DE", "A", 100),
        ("DE", "A", 200),
        ("FR", "B", 300),
        ("FR", "B", 150),
    ],
    ["country", "category", "amount"]
)

result = (
    sales
    .filter(F.col("amount") > 100)
    .groupBy("country", "category")
    .agg(
        F.sum("amount").alias("total_amount")
    )
)

window = Window.partitionBy("country").orderBy(
    F.desc("total_amount")
)

ranked = result.withColumn(
    "rank",
    F.row_number().over(window)
)

ranked.show()

This example contains both narrow and wide operations:

Filtering is narrow.
Grouping requires a shuffle.
The window operation may require another repartition and sort.

Understanding where these boundaries happen is essential for tuning.

Joins are one of the most common sources of expensive Spark jobs.

Consider two tables:

A large transaction table.
A small lookup table.

A normal join may shuffle both sides.

Example:

joined = transactions.join(
    customers,
    "customer_id"
)

If customers is small enough, broadcasting it is usually faster.

A broadcast join sends the small dataset to every executor instead of moving the large dataset.

from pyspark.sql.functions import broadcast

joined = transactions.join(
    broadcast(customers),
    "customer_id"
)

A broadcast join avoids a large shuffle, but it has a limit: the broadcasted table must fit comfortably in executor memory.

Do not broadcast a table just because it is smaller than the other side. A multi-gigabyte lookup table can still create memory pressure.

You can inspect query plans:

joined.explain("formatted")

Look for physical operators such as:

BroadcastHashJoin
SortMergeJoin

A BroadcastHashJoin usually indicates Spark avoided a large shuffle.

Python UDFs are convenient:

from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

def add_tax(value):
    return int(value * 1.19)

tax_udf = udf(add_tax, IntegerType())

result = df.withColumn(
    "price_with_tax",
    tax_udf("price")
)

But this introduces a boundary between the JVM and Python.

For every row, Spark must:

Convert JVM data into Python objects.
Send data to a Python worker.
Execute Python code.
Serialize results back.

For millions or billions of rows, this overhead becomes significant.

The problem is not that Python itself is always slow. The problem is moving data between Spark’s execution engine and Python one row at a time.

pandas_udf changes the execution model.

Instead of sending one row at a time, Spark uses Apache Arrow to transfer batches of data between the JVM and Python.

The flow becomes:

Spark JVM
    |
    | Arrow columnar batches
    v
Python + pandas
    |
    | Arrow columnar batches
    v
Spark JVM

This reduces serialization overhead dramatically.

Example:

import pandas as pd

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType


@pandas_udf(IntegerType())
def add_tax_vectorized(price: pd.Series) -> pd.Series:
    return (price * 1.19).astype("int")


result = df.withColumn(
    "price_with_tax",
    add_tax_vectorized("price")
)

The function receives a pandas Series, not a single value.

Spark processes batches of values at once.

A simple benchmark:

import time

from pyspark.sql.functions import udf, pandas_udf
from pyspark.sql.types import LongType
import pandas as pd


@udf(LongType())
def slow_double(value):
    return value * 2


@pandas_udf(LongType())
def fast_double(values: pd.Series) -> pd.Series:
    return values * 2


data = spark.range(0, 20_000_000)


start = time.time()

(
    data
    .withColumn("value", slow_double("id"))
    .agg({"value": "sum"})
    .collect()
)

print("Python UDF seconds:", time.time() - start)


start = time.time()

(
    data
    .withColumn("value", fast_double("id"))
    .agg({"value": "sum"})
    .collect()
)

print("pandas_udf seconds:", time.time() - start)

The exact improvement depends on:

Cluster size.
Data types.
Function complexity.
Serialization cost.

For simple transformations, vectorized UDFs can be many times faster. In real workloads, improvements of around 10x are common when replacing row-by-row Python UDFs with vectorized operations.

However, the fastest option is often no UDF at all.

For example:

optimized = df.withColumn(
    "value_times_two",
    F.col("value") * 2
)

Built-in Spark functions stay inside the optimized execution engine.

Before writing a UDF, check whether Spark already provides the operation.

Common examples:

from pyspark.sql import functions as F

result = df.select(
    F.lower("name").alias("normalized_name"),
    F.round("amount", 2).alias("rounded_amount"),
    F.when(
        F.col("amount") > 100,
        "high"
    ).otherwise("low").alias("category")
)

Built-in functions allow Spark to:

Optimize expressions.
Generate efficient JVM code.
Reduce Python overhead.
Push operations closer to the data.

A UDF should usually be the exception, not the default.

Reduce data before joining

Instead of joining huge tables immediately:

filtered_orders = orders.filter(
    F.col("status") == "completed"
)

result = filtered_orders.join(
    customers,
    "customer_id"
)

Filter early and select only required columns.

Avoid unnecessary repartitioning

This creates a shuffle:

df = df.repartition("customer_id")

Only repartition when you have a reason, such as preparing for a known expensive operation.

Use adaptive query execution

Modern Spark versions include Adaptive Query Execution (AQE), which can adjust plans during runtime.

Example configuration:

spark.conf.set(
    "spark.sql.adaptive.enabled",
    "true"
)

AQE can help with:

Better join strategies.
Handling skewed data.
Reducing unnecessary shuffle partitions.

A common production pattern is a pipeline that runs for hours because of one hidden shuffle.

One widely shared Spark optimization pattern is that teams have reduced multi-hour jobs to minutes by identifying a single expensive join or aggregation, then replacing it with a broadcast join, early filtering, or better partitioning.

The lesson is simple: adding more executors is not always the answer. If the job spends most of its time moving data across the network, more compute may only increase the amount of data competing for the same bottleneck.

Performance is not the only reason to understand Spark internals.

In 2022, Apache Spark addressed a vulnerability tracked as CVE-2022-33891. The issue affected certain Spark configurations where an attacker could exploit improper handling of authorization settings and potentially execute shell commands.

The important takeaway is that big data systems are production software, not just batch-processing tools. Secure configuration matters:

Do not expose Spark interfaces unnecessarily.
Keep Spark versions updated.
Review authentication and authorization settings.
Treat cluster access as privileged access.

Optimizing a cluster is only useful if the cluster is protected.

When a PySpark job is slow:

Check the physical plan with explain("formatted").
Look for unexpected shuffles.
Filter and select columns early.
Use broadcast joins for genuinely small tables.
Prefer built-in Spark SQL functions.
Replace row-based Python UDFs with pandas_udf when Python logic is unavoidable.
Measure before and after every optimization.

The biggest PySpark performance wins usually come from understanding where data moves, not from writing faster Python.

Start by opening your slowest Spark job, inspect its execution plan, find the largest shuffle, and remove one unnecessary data movement step. That single change can often deliver a bigger improvement than tuning dozens of configuration settings.

The execution model: DataFrames are more than tables

Narrow versus wide transformations

Narrow transformations

Wide transformations

A practical DataFrame workflow

Reduce data before joining

Avoid unnecessary repartitioning

Use adaptive query execution

Vulnerability Watch No15

ONNX: one model format to run everywhere, from server to browser

PySpark ML performance: partitions, data skew, caching, and the small-data trap