Distributed processing with Spark

The execution model

A Spark application has a driver (your program: builds the plan, coordinates) and executors (worker processes that hold data partitions in memory and run tasks). Nothing runs until an action (write, count, collect) triggers a job. The job is split into stages at every point where data must be redistributed across executors — a shuffle — and each stage runs one task per partition of its input. Partitions are the unit of parallelism: two hundred partitions means up to two hundred tasks running concurrently, one task per core.

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.appName("daily_revenue").getOrCreate()

orders = spark.read.parquet("s3://acme-lake/silver/orders/")          # lazy: nothing read yet
customers = spark.read.parquet("s3://acme-lake/silver/customers/")

result = (
    orders
    .filter(F.col("ordered_at") >= "2026-03-01")                        # partition-local; pushed to the scan
    .filter(F.col("status") == "paid")
    .join(customers.select("customer_id", "country"), "customer_id")     # shuffle 1 (or broadcast, see below)
    .groupBy(F.to_date("ordered_at").alias("day"), "country")            # shuffle 2: group by key
    .agg(F.sum("amount").alias("revenue"), F.countDistinct("order_id").alias("orders"))
    .repartition(1)                                                      # shuffle 3: single output file (small result)
)

result.write.mode("overwrite").partitionBy("day").parquet("s3://acme-lake/gold/daily_revenue/")  # action: runs the job

The optimiser (Catalyst) rewrites the plan before execution: it pushes filters down to the scan so only the March partitions are read, prunes unused columns, and picks join strategies. Adaptive query execution (on by default in current versions) re-plans between stages using observed sizes: coalescing tiny partitions, switching join strategies, splitting skewed partitions. Your job is to write transformations the optimiser can see through — DataFrame operations and SQL — and to avoid Python UDFs and collect() that turn the plan opaque or pull data to the driver.

Shuffles: the cost that matters

A shuffle writes every partition's data to local disk grouped by key, then every executor fetches its share over the network. It is the most expensive thing Spark does and it happens on groupBy, join, distinct, repartition, window functions with partitionBy, and orderBy. Reducing shuffles, or the data that flows through them, is most of Spark tuning. Filter and project before the shuffle so less data moves; aggregate before joining when possible; and use a broadcast join when one side is small (under a few hundred MB): the small table is copied to every executor and the join runs partition-local with no shuffle of the big side.

from pyspark.sql import functions as F

countries = spark.read.parquet("s3://acme-lake/silver/countries/")          # a few thousand rows

joined = orders.join(F.broadcast(countries), "country_code")
joined.explain(mode="formatted")
# look for: BroadcastHashJoin ... and no Exchange (shuffle) on the orders side

# the threshold below which Spark broadcasts automatically (default 10 MB); raise for medium dimensions
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", str(256 * 1024 * 1024))

The number of shuffle partitions (spark.sql.shuffle.partitions, default 200) sets how many tasks the post-shuffle stage has. Too few means huge tasks and spills; too many means thousands of tiny tasks and files. Adaptive execution coalesces small partitions automatically, so the modern advice is to set the value high enough for the largest shuffle and let AQE shrink it. Output file counts follow from partition counts; a final coalesce(n) (no shuffle) or repartition(n) (shuffle, balanced) controls how many files land.

Skew

Data skew is when one key holds far more rows than the others: a NULL customer id, a default account, one giant tenant. After a shuffle on that key, one task gets most of the data and the whole stage waits for it; the Spark UI shows a stage where the median task took seconds and the maximum took an hour. AQE's skew join optimisation splits oversized partitions for sort-merge joins, which handles many cases automatically. When it does not, the fixes are manual: handle the hot key separately (filter it out and union its result), salt the key (append a random suffix on the big side and explode the small side to match, so the hot key spreads over N tasks), or broadcast the small side so no shuffle happens at all.

from pyspark.sql import functions as F

N = 16
big_salted = events.withColumn("salt", (F.rand() * N).cast("int"))
small_replicated = users.crossJoin(spark.range(N).withColumnRenamed("id", "salt"))

joined = big_salted.join(
    small_replicated,
    on=["user_id", "salt"],
).drop("salt")

# confirm AQE settings that help before resorting to salting
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

Find skew before it finds you: df.groupBy(key).count().orderBy(F.desc('count')).show(10) on a sample tells you whether one key dominates. Skew on groupBy keys is handled by two-phase aggregation (partial aggregates before the shuffle, which Spark does automatically for sums and counts but not for collect_list or countDistinct over a hot key).

Reading the Spark UI and the plan

The Spark UI (port 4040 while running, the history server afterwards) is where every performance question is answered. Jobs tab: which action took the time. Stages tab: for the slow job, which stage; within it, the task duration distribution (skew shows as a long tail), shuffle read and write sizes (how much data moved), and spill (memory pressure). SQL tab: the physical plan with per-operator rows and time, where you see the scan's pushed filters, the join strategy and every Exchange. Executors tab: memory use, GC time, failed tasks — an executor spending half its time in garbage collection needs more memory or fewer cached DataFrames.

== Physical Plan ==
AdaptiveSparkPlan
+- HashAggregate(keys=[day, country], functions=[sum(amount), count(distinct order_id)])
   +- Exchange hashpartitioning(day, country, 200)          <- the one shuffle; check its size in the SQL tab
      +- HashAggregate(keys=[day, country], functions=[partial_sum(amount), ...])   <- partial agg before shuffle
         +- Project [day, country, amount, order_id]
            +- BroadcastHashJoin [customer_id], [customer_id], Inner    <- small side broadcast: no shuffle of orders
               :- Filter (status = paid)
               :  +- FileScan parquet orders[...]
               :       PartitionFilters: [ordered_at >= 2026-03-01]      <- partition pruning is working
               :       PushedFilters: [EqualTo(status,paid)]             <- predicate pushed to the reader
               +- BroadcastExchange
                  +- FileScan parquet customers[customer_id, country]    <- column pruning: 2 of 40 columns

• No PartitionFilters on a partitioned table: the filter is not prunable; rewrite it.
• SortMergeJoin with a tiny side: broadcast it.
• An Exchange right after another Exchange on the same key: a repartition you did not need.
• Stage with max task time far above median: skew.
• Spill (memory and disk) on a stage: more partitions, more executor memory, or less cached data.
• A BatchEvalPython operator: a Python UDF; replace with built-in functions or a pandas UDF.

When Spark is the wrong tool

Spark's overhead is real: cluster start-up, scheduling, shuffles, and the operational weight of managing it. For data that fits on one large machine — which today means hundreds of GB — single-node engines (DuckDB, polars) are often faster end to end and far simpler. For SQL-expressible transformations over data already in the warehouse, the warehouse's engine is usually cheaper and needs no cluster. Spark earns its place when data is genuinely large, when the logic needs code (complex parsing, ML feature pipelines, custom algorithms), when it must read and write many formats and systems, or when streaming and batch share one codebase.