Spark

Spark

Purpose

Write Spark jobs whose cost is understood. Almost all Spark performance problems are one of three things: too much shuffle, skewed partitions, or reading far more data than the query needs.

When to Use

• Building or reviewing a Spark pipeline.
• A job that is slow, failing with out-of-memory errors, or has one straggling task.
• Tuning partitioning and join strategy.
• Reading the Spark UI to diagnose a stage.

Capabilities

• Partitioning strategy and repartitioning.
• Shuffle minimization and broadcast joins.
• Skew detection and mitigation.
• Caching and persistence levels.
• File-format and predicate-pushdown optimization.
• Spark UI interpretation.

Inputs

• The job, its input data volume, and its physical plan.
• The Spark UI: stage timings, task distribution, shuffle read/write.
• Cluster resources.

Outputs

• A plan with fewer or smaller shuffles.
• Balanced partitions with no straggling tasks.
• Measured improvement in wall-clock time and cost.

Workflow

1. Read the plan first — df.explain(True). Every Exchange is a shuffle, and a shuffle writes to disk and crosses the network. It is the dominant cost.
2. Prune early — Select the columns and filter the rows you need before joining, not after. With Parquet, this pushes down to the file reader and never reads the data at all.
3. Broadcast the small side — A join where one side fits in memory (roughly under 100 MB) should be a broadcast join. That eliminates the shuffle entirely.
4. Find the skew — In the Spark UI, look at the task duration distribution within a stage. If the max is 50x the median, one partition holds most of the data. That single task is your job's runtime.
5. Mitigate the skew — Salting the key, or enabling adaptive query execution's skew join handling.
6. Cache only what is reused — Caching a DataFrame used once costs memory and gains nothing.

Best Practices

• Enable adaptive query execution (spark.sql.adaptive.enabled=true). It coalesces partitions, converts joins to broadcasts at runtime, and handles skew automatically. It is on by default from Spark 3.2 and is the single largest free improvement available.
• A shuffle is the most expensive operation in Spark. groupBy, join, distinct, and repartition all shuffle. Count them in the plan.
• collect() brings the entire dataset to the driver. On anything non-trivial this is an out-of-memory error waiting for a larger input.
• Use Parquet or Delta, never CSV or JSON, for anything that will be read more than once. Columnar formats support predicate and projection pushdown; row formats do not.
• Too many small partitions means task-scheduling overhead dominates; too few means poor parallelism and memory pressure. Aim for partitions of roughly 128 MB.
• cache() is lazy. It does nothing until an action runs, and it can silently fall back to recomputation if it does not fit in memory.

Examples

Prune, broadcast, and avoid the shuffle:

from pyspark.sql import functions as F

# Costly: joins the full orders table, then filters. The shuffle moves
# everything, including the rows about to be discarded.
result = (
    orders.join(customers, "customer_id")
          .filter(F.col("created_at") >= "2026-01-01")
          .select("order_id", "segment", "total_cents")
)

# Better: filter and project first (pushed into the Parquet reader), and
# broadcast the small dimension so the large side is never shuffled.
result = (
    orders
      .filter(F.col("created_at") >= "2026-01-01")     # predicate pushdown
      .select("order_id", "customer_id", "total_cents") # projection pushdown
      .join(F.broadcast(customers.select("customer_id", "segment")), "customer_id")
)

Diagnosing and fixing skew:

# Symptom in the Spark UI: stage 7 has 200 tasks; 199 finish in ~4s and one
# runs for 22 minutes. That one task is the entire job's runtime.

# Confirm: one key dominates.
orders.groupBy("customer_id").count().orderBy(F.desc("count")).show(5)
# +------------+---------+
# | customer_id|    count|
# +------------+---------+
# | cus_bulk_01| 41200000|   <-- 68% of all rows: a single bulk-import account
# | cus_9f2a3b |    18400|

# Fix: salt the hot key so it spreads across partitions.
SALTS = 64
orders_salted = orders.withColumn(
    "salt",
    F.when(F.col("customer_id") == "cus_bulk_01",
           (F.rand() * SALTS).cast("int")).otherwise(F.lit(0)),
)
customers_exploded = customers.withColumn(
    "salt",
    F.explode(F.when(F.col("customer_id") == "cus_bulk_01",
                     F.array([F.lit(i) for i in range(SALTS)])).otherwise(F.array(F.lit(0)))),
)
result = orders_salted.join(customers_exploded, ["customer_id", "salt"])

Notes

• Adaptive query execution handles most skew automatically from Spark 3.2 onward (spark.sql.adaptive.skewJoin.enabled). Manual salting is still needed for extreme skew and for non-join operations.
• spark.sql.shuffle.partitions defaults to 200 regardless of your data size, which is wrong for both small and large jobs. AQE's coalescing largely fixes this, but the default is still worth setting deliberately.
• If the job fits comfortably on one machine, do not use Spark. A single-node DuckDB or Polars job will be faster and dramatically simpler than a distributed one.

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: nimadorostkar
• Source: nimadorostkar/Claude-Skills-collection
• License: MIT
• Homepage: https://github.com/nimadorostkar/Claude-Skills-collection

Install and usage instructions live in the source repository linked above.