All posts

Spark performance optimization: seven techniques, and how to prove each one worked

Most Spark tuning advice is a list of settings with no way to tell whether any of them fired. These seven techniques each come with the code, the plan line that proves the optimization applied, and the failure mode that silently turns it off again.

13 min read Spark

TL;DR

  • Prove it from the plan, not the clock. A wall-clock comparison on a busy cluster cannot separate your change from the noise. explain("formatted") names the optimization that fired: PartitionFilters, PushedFilters, ReadSchema, BroadcastHashJoin, AQEShuffleRead.
  • select * costs you column pruning. The same filter read 2 columns with an explicit projection and 5 columns without one, from the identical query.
  • A join is broadcast or not depending on whether the build side fits spark.sql.autoBroadcastJoinThreshold. Disabling it turned one BroadcastExchange into a SortMergeJoin with two Exchange nodes.
  • A Python UDF inserts BatchEvalPython into the plan and splits whole-stage code generation around it. The built-in equivalent stays inside one generated stage.
  • Read the plan that actually ran. Calling count() to inspect an adaptive plan builds a different query and reports isFinalPlan=false, describing a plan that never executed.

Why is Spark tuning so hard to verify?

The usual tuning cycle is: change a setting, re-run the job, watch the clock, keep the change if the number went down. On a shared cluster that loop cannot tell you anything. The input was cached differently, the executors landed on different hosts, another job was competing for the same disks, and the difference you measured is smaller than the variance between two runs of the unchanged job.

Worse, most optimizations are conditional. Predicate pushdown applies only if the filter is expressible in the file format. A broadcast join happens only if the build side is small enough at planning time. Partition pruning needs the filter to reference the partition column directly. Each of these silently does nothing when its condition is not met, and a slow job looks exactly the same whether the optimization was skipped or applied.

So this post is organised around a different question. Not “is it faster?”, but “did the optimization actually fire?” That question has a definite answer, it is visible in the query plan, and it does not move when the cluster is busy.

Seven techniques follow. Each one has the code, the line in the plan that proves it applied, and the mistake that turns it off.

The dataset these examples run against

Everything below was run on Spark 4.1.3 in local[4] mode against two Parquet tables written for the purpose. They are small enough to reproduce and shaped like something real rather than an id, value toy.

orders = (spark.range(0, 4_000_000)
          .withColumn("order_id", F.col("id"))
          # 70% of rows land on 10 customers: a realistic heavy-hitter skew
          .withColumn("customer_id",
                      F.when(F.rand(42) < 0.7, (F.rand(7) * 10).cast("int"))
                       .otherwise((F.rand(9) * 200000).cast("int")))
          .withColumn("order_date",
                      F.date_add(F.lit("2026-01-01").cast("date"), (F.col("id") % 30).cast("int")))
          .withColumn("amount", (F.rand(3) * 1000).cast("decimal(10,2)"))
          .withColumn("status",
                      F.when(F.col("id") % 17 == 0, F.lit("CANCELLED")).otherwise(F.lit("SHIPPED")))
          .withColumn("notes", F.concat(F.lit("note-"), F.col("id").cast("string")))
          .drop("id"))
orders.write.mode("overwrite").partitionBy("order_date").parquet("/work/data/orders")

What that produced, measured rather than assumed:

Table Rows Files Bytes on disk
orders 4,000,000 120 across 30 date partitions 66,125,993
customers 200,000 4 812,232

The customers table being well under a megabyte matters later: it is why the join in technique 3 is broadcast without anyone asking for it.

How do you read the proof?

One method underlies every technique here.

df.explain("formatted")

The formatted plan prints each operator with its own block of attributes, and five attribute names carry most of the signal:

In the plan What it proves
PartitionFilters directory-level pruning; whole partitions never opened
PushedFilters the filter was handed to the Parquet reader
ReadSchema exactly which columns leave the scan
BroadcastHashJoin / SortMergeJoin which join strategy was chosen
AQEShuffleRead adaptive execution changed the shuffle at runtime

Read the plan that ran, not a plan. For anything adaptive, the physical plan printed before execution is a prediction. The plan that ran is available only after the query has run, from the same QueryExecution object:

agg = orders.groupBy("status").agg(F.count("*").alias("n"))
rows = agg.collect()                 # executes THIS QueryExecution
print(agg._jdf.queryExecution().executedPlan().toString())

Using count() here instead of collect() is a trap worth naming, because it looks equivalent and is not. count() builds a different query, so the plan you then inspect is one that never executed. Doing exactly that while writing this post produced AdaptiveSparkPlan isFinalPlan=false, describing a plan Spark had discarded. The collect() version reports isFinalPlan=true.

1. Prune partitions before anything is opened

The cheapest data to process is data never read. When a table is partitioned on disk and the filter names the partition column, Spark resolves the filter against directory names and never opens the rest.

orders.filter(F.col("order_date") == "2026-01-05").select("order_id", "amount")

The proof is one line in the scan block:

(1) Scan parquet
Location: InMemoryFileIndex [file:/work/data/orders]
PartitionFilters: [isnotnull(order_date#5), (order_date#5 = 2026-01-05)]
ReadSchema: struct<order_id:bigint,amount:decimal(10,2)>

PartitionFilters carrying your predicate means 29 of the 30 date directories were never listed or opened.

How it silently stops working. Wrap the partition column in a function and the filter is no longer a partition filter, because Spark cannot invert the function to decide which directories match. F.year(F.col("order_date")) == 2026 moves the predicate out of PartitionFilters and into a post-scan Filter, and every partition gets read. Compare the partition column to a literal, and do the arithmetic on the other side of the comparison.

2. Push the filter down, and take only the columns you need

Two separate optimizations live in the same scan block, and one of them is undone by a habit almost everyone has.

orders.filter(F.col("status") == "CANCELLED").select("order_id", "status")
PushedFilters: [IsNotNull(status), EqualTo(status,CANCELLED)]
ReadSchema: struct<order_id:bigint,status:string>

PushedFilters means the Parquet reader itself skips row groups whose statistics rule out a match, so the rows never reach Spark. ReadSchema with two fields means only two columns were decoded.

Now the same query without the projection:

orders.filter(F.col("status") == "CANCELLED")
PushedFilters: [IsNotNull(status), EqualTo(status,CANCELLED)]
ReadSchema: struct<order_id:bigint,customer_id:int,amount:decimal(10,2),status:string,notes:string>

The pushdown is unchanged. The ReadSchema went from 2 columns to 5, including notes, the widest column in the table. Column pruning is driven by what you project, and select * switches it off while leaving every other optimization looking healthy.

This is the single cheapest change in this post: name your columns.

3. Let the small side be broadcast

A join between a large table and a small one does not need a shuffle. If the small side fits in memory, Spark ships a copy to every executor and joins locally.

orders.join(customers, "customer_id").select("order_id", "segment")
+- Project (8)
   +- BroadcastHashJoin Inner BuildRight (7)
      +- BroadcastExchange (6)

BroadcastHashJoin with BuildRight means customers was the side broadcast. It qualified because its 812,232 bytes are comfortably under spark.sql.autoBroadcastJoinThreshold, which defaults to 10 MiB.

Disabling the threshold on the identical query gives the contrast:

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)
+- Project (11)
   +- SortMergeJoin Inner (10)
      :  +- Exchange (4)
         +- Exchange (8)

One BroadcastExchange became a SortMergeJoin with two Exchange nodes: both sides of the join now shuffle. That is the cost the broadcast was avoiding, and it is why this is usually the highest-value line in a plan to check.

How it silently stops working. The decision uses the estimated size at planning time. If the small side is the result of earlier transformations rather than a table read, the estimate can exceed the threshold even when the real output is tiny, and the plan quietly falls back to a sort-merge join. The estimate comes from statistics, so a table that has never had statistics collected can be mis-sized in either direction. Where you are certain, say so explicitly with F.broadcast(customers), which bypasses the estimate.

The full strategy-selection logic is its own subject, covered in Spark joins in depth.

4. Let adaptive execution finish the plan

spark.sql.shuffle.partitions is a number chosen before Spark knows how much data there is. Adaptive query execution re-plans after each shuffle writes, using the sizes it actually observed.

agg = orders.groupBy("status").agg(F.count("*").alias("n"))
agg.collect()
print(agg._jdf.queryExecution().executedPlan().toString())

With spark.sql.adaptive.enabled on:

AdaptiveSparkPlan isFinalPlan=true
   +- *(2) HashAggregate(keys=[status#3], ...)
      +- AQEShuffleRead coalesced
            +- Exchange hashpartitioning(status#3, 200), ENSURE_REQUIREMENTS

And with it off, the same query:

*(2) HashAggregate(keys=[status#3], ...)
+- Exchange hashpartitioning(status#3, 200), ENSURE_REQUIREMENTS

AQEShuffleRead coalesced is the proof. The Exchange still says 200, because that is what was configured and written; the coalescing happens on the read side, where AQE merged those 200 partitions into far fewer tasks after seeing how little data each held. The aggregate produces two rows, so 200 reduce tasks were 198 tasks’ worth of scheduling overhead for nothing.

isFinalPlan=true is the other half of the proof, and it is the reason the collect() discipline from earlier matters. Without it you are reading a prediction.

The re-planning AQE performs, and what it declines to touch, is covered in adaptive query execution.

5. Keep the work inside the JVM

A Python UDF is not just slower arithmetic. It changes the shape of the plan.

upper_udf = F.udf(lambda s: s.upper() if s else None, StringType())
orders.select(upper_udf(F.col("status")).alias("s"))
* Project (5)
+- BatchEvalPython (4)
   +- * Project (3)
      +- * ColumnarToRow (2)
         +- Scan parquet  (1)

The built-in equivalent:

orders.select(F.upper(F.col("status")).alias("s"))
* Project (3)
+- * ColumnarToRow (2)
   +- Scan parquet  (1)

Two things changed. BatchEvalPython appeared, which is the operator that serializes rows, ships them to a Python worker process, and reads results back. And notice the asterisks: * marks an operator fused into a whole-stage generated function. BatchEvalPython has no asterisk, so it splits the generated stage in two, and the operators either side of it can no longer be fused together.

The cost is therefore not only the Python interpreter. It is the serialization round trip per batch plus the loss of code generation across the boundary.

The rule. Reach for a built-in first; the pyspark.sql.functions catalogue covers more than most people assume. Where a UDF is genuinely necessary, a pandas UDF amortizes the boundary crossing over a whole batch instead of paying it per row, and a Scala implementation avoids the process boundary altogether.

6. Choose between repartition and coalesce deliberately

These two look interchangeable and are not.

print("source partitions            =", orders.rdd.getNumPartitions())
print("after repartition(50)        =", orders.repartition(50).rdd.getNumPartitions())
print("after coalesce(2)            =", orders.coalesce(2).rdd.getNumPartitions())
print("after repartition(customer)  =", orders.repartition(50, "customer_id").rdd.getNumPartitions())
source partitions            = 5
after repartition(50)        = 50
after coalesce(2)            = 2
after repartition(customer)  = 50

Both reach a partition count, by different means. repartition shuffles and can raise or lower the count, distributing rows evenly, and with a column argument it co-locates rows sharing a key. coalesce does not shuffle; it merges existing partitions into fewer, which is nearly free and is the right tool immediately before a write to reduce output file count.

The trap is that coalesce reaches backwards. Because it introduces no shuffle, the reduced parallelism applies to the stage that produces the data, not just the write. coalesce(2) before a heavy aggregation does not aggregate with full parallelism and then narrow; it runs the aggregation itself on two tasks. When you want fewer output files without throttling the computation that makes them, repartition is the one that costs a shuffle and keeps the upstream parallelism.

7. Know your skew before you tune for it

Skew is the failure mode most often misdiagnosed as a memory problem, and it is measurable in one query.

top = orders.groupBy("customer_id").count().orderBy(F.desc("count")).limit(5).collect()
customer_id=      5 rows=   280751 share=7.02%
customer_id=      2 rows=   280675 share=7.02%
customer_id=      8 rows=   280332 share=7.01%
customer_id=      0 rows=   280332 share=7.01%
customer_id=      9 rows=   280304 share=7.01%
total rows = 4000000

Five keys holding 7% of the table each is a partition that will be roughly fourteen thousand times the size of an average one when there are 200,000 distinct keys. No memory setting fixes that; one task simply has more work than the others and the stage waits for it.

Measure this before reaching for a fix. If the distribution is flat, your slow stage is not skew and salting it will only add a shuffle. The fixes themselves are a subject of their own, in data skew in Apache Spark.

In what order should you try these?

Techniques are worth ranking, because they are not equally likely to matter.

  1. Read less. Partition pruning and column pruning change how many bytes enter the job at all, and they are usually a one-line change.
  2. Shuffle less. A broadcast that replaces a sort-merge join removes two exchanges. Nothing you tune inside a shuffle beats not having it.
  3. Let AQE size the shuffles you keep. It is on by default and it is better at choosing partition counts after the fact than you are before it.
  4. Remove boundary crossings. Replace Python UDFs with built-ins where they exist.
  5. Only then tune memory and parallelism. Spill and executor sizing are real, but they are the fourth thing to look at, not the first.

The ordering is not arbitrary: each step reduces the amount of data the later steps have to be clever about.

Common misconceptions

“Caching makes things faster.” Caching makes re-reads faster, and costs memory that execution then cannot use. A DataFrame read once and never reused is slower cached than not, because you paid to materialize it.

“More partitions is more parallelism, so it is better.” More partitions means smaller tasks and more scheduling and more shuffle blocks. Past the point where every core is busy, additional partitions are pure overhead, which is what AQEShuffleRead coalesced exists to undo.

repartition(1) and coalesce(1) are the same before a write.” coalesce(1) pushes single-threaded execution back up into the stage that produces the data. repartition(1) shuffles into one partition and leaves the upstream stage parallel.

“The physical plan tells you what will happen.” With adaptive execution on, the pre-execution plan is a proposal. Only isFinalPlan=true describes what ran.

“A Python UDF is a bit slower than a built-in.” It also splits whole-stage code generation and adds a serialization boundary, so the cost is structural rather than proportional.

A model worth keeping

Optimization in Spark is mostly the art of not doing work: not opening partitions, not decoding columns, not shuffling rows, not leaving the JVM. Every technique here is one of those four, and each one leaves a specific line in the query plan when it applies.

So the habit worth building is not a list of settings. It is opening explain("formatted") before and after a change and reading the five attributes in the table above. A change that does not alter the plan did not do anything, however much the wall clock moved.

References

Trademarks

Apache Spark, Apache Parquet, Apache and the Apache feather logo are either registered trademarks or trademarks of The Apache Software Foundation in the United States and other countries.

Found this useful?

These posts and tools are free. If one saved you an afternoon, you can buy me a coffee.

Buy me a coffee