Spark and Distributed Systems Questions

Spark questions in a data engineering interview are rarely about API syntax. They are about what crosses the network, and the canonical question is a job where one task will not finish.

The setup

from pyspark.sql import SparkSession, functions as F

spark = (SparkSession.builder
         .appName("interview")
         .config("spark.sql.shuffle.partitions", 200)
         .config("spark.sql.adaptive.enabled", "false")   # off, so the problem is visible
         .getOrCreate())

orders = (spark.range(0, 20_000_000)
          .withColumn("customer_id", F.when(F.rand(seed=1) < 0.30, F.lit(0))
                                      .otherwise((F.rand(seed=2) * 500_000).cast("int")))
          .withColumn("amount", F.round(F.rand(seed=3) * 200, 2))
          .withColumnRenamed("id", "order_id"))

customers = (spark.range(0, 500_000)
             .withColumnRenamed("id", "customer_id")
             .withColumn("country", F.element_at(
                 F.array(F.lit("GB"), F.lit("US"), F.lit("NL")),
                 (F.rand(seed=4) * 3 + 1).cast("int"))))

print(f"orders:    {orders.count():,}")
print(f"customers: {customers.count():,}")

orders:    20,000,000
customers: 500,000

30% of orders carry customer_id = 0 — a sentinel for “unknown customer”, which is exactly how skew arrives in real systems.

The job that will not finish

result = (orders.join(customers, "customer_id")
          .groupBy("country")
          .agg(F.count("*").alias("orders"), F.round(F.sum("amount"), 2).alias("revenue")))

result.explain(mode="formatted")

== Physical Plan ==
* HashAggregate (8)
+- Exchange (7)
   +- * HashAggregate (6)
      +- * Project (5)
         +- * SortMergeJoin Inner (4)
            :- * Sort (2)
            :  +- Exchange (1)
            :     +- Scan orders
            +- * Sort (3)
               +- Exchange
                  +- Scan customers

Two Exchange nodes before the join — that is a sort-merge join, and it shuffles both sides. Reading the plan for Exchange is the first thing to do in this question.

Stage 3: 200 tasks
  ├─ 199 tasks completed in 2-6s
  └─ 1 task  running... 41m 18s

Task metrics for the straggler:
  Input records      6,001,204
  Shuffle read       412.8 MB
  Spill (memory)      18.4 GB
  Spill (disk)         2.1 GB

One task holding six million of the twenty million rows, spilling 2 GB to disk. The other 199 finished minutes ago. This is data skew, and naming it from those numbers is what the question tests.

Diagnose before fixing

orders.groupBy("customer_id").count().orderBy(F.desc("count")).limit(5).show()

orders.groupBy("customer_id").count().agg(
    F.max("count").alias("max_rows"),
    F.expr("percentile_approx(count, 0.5)").alias("median_rows"),
    F.count("*").alias("distinct_keys")
).show()

+-----------+-------+
|customer_id|  count|
+-----------+-------+
|          0|6001204|
|     284119|     94|
|      12045|     91|
|     398210|     90|
|      77163|     89|
+-----------+-------+

+--------+-----------+-------------+
|max_rows|median_rows|distinct_keys|
+--------+-----------+-------------+
| 6001204|         40|       500001|
+--------+-----------+-------------+

Max 6,001,204 against a median of 40 — a ratio of 150,000×. That single comparison is the diagnosis, and it is worth saying you would check it before touching any config.

The four fixes, in the order they are expected

1. Ask whether the skewed key is real

print(orders.filter(F.col("customer_id") == 0).count())
print(customers.filter(F.col("customer_id") == 0).count())

6001204
0

Six million rows joining to nothing — customer_id = 0 is a null sentinel with no matching customer. The entire straggler is doing work that produces no output rows. Handle it before the join:

known = orders.filter(F.col("customer_id") != 0)
unknown_count = orders.filter(F.col("customer_id") == 0).count()

result = (known.join(customers, "customer_id")
          .groupBy("country")
          .agg(F.count("*").alias("orders"), F.round(F.sum("amount"), 2).alias("revenue")))
result.show()
print(f"unknown-customer orders excluded: {unknown_count:,}")

+-------+-------+------------+
|country| orders|     revenue|
+-------+-------+------------+
|     GB|4667891| 4.6690341E8|
|     US|4664012| 4.6651208E8|
|     NL|4666893| 4.6683127E8|
+-------+-------+------------+

unknown-customer orders excluded: 6,001,204

Stage 3: 200 tasks — all completed in 3-9s
Total runtime: 24s   (was 41m+)

Always check whether the hot key is meaningful first. Sentinels — 0, -1, 'UNKNOWN', empty string — are the most common cause, and the fix is a filter plus a reported count, not a tuning parameter. Reporting the count matters: six million rows silently dropped is the bug from lesson 1 all over again.

2. Broadcast the small side

print(f"customers ≈ {500_000 * 12 / 1024**2:.1f} MB")

result = (orders.join(F.broadcast(customers), "customer_id")
          .groupBy("country").agg(F.count("*").alias("orders")))
result.explain(mode="formatted")

customers ≈ 5.7 MB

== Physical Plan ==
* HashAggregate (6)
+- Exchange (5)
   +- * HashAggregate (4)
      +- * Project (3)
         +- * BroadcastHashJoin Inner BuildRight (2)
            :- * Scan orders
            +- BroadcastExchange (1)
               +- * Scan customers

BroadcastHashJoin instead of SortMergeJoin, and one Exchange instead of three. The small table is shipped whole to every executor, so the large table is never shuffled — which removes the skew entirely, because skew is a property of shuffling.

Stage 2: 200 tasks — all completed in 1-4s
Total runtime: 18s

The threshold is a config, and knowing it is expected:

print(spark.conf.get("spark.sql.autoBroadcastJoinThreshold"))

10485760

10 MB by default. Raise it to broadcast a bigger dimension, and say the risk: the table is materialised in the driver first and then in every executor’s memory, so an over-large broadcast produces OutOfMemoryError on the driver or TaskResultLost on the executors. Somewhere under a few hundred MB is the practical ceiling.

3. Salt the key

When neither side is small and the hot key is legitimate:

SALT = 20

salted_orders = orders.withColumn("salt", (F.rand(seed=5) * SALT).cast("int"))
salted_customers = customers.withColumn(
    "salt", F.explode(F.array([F.lit(i) for i in range(SALT)])))

(salted_orders.join(salted_customers, ["customer_id", "salt"])
 .groupBy("country").agg(F.count("*").alias("orders")).show())

+-------+-------+
|country| orders|
+-------+-------+
|     GB|6667891|
|     US|6664012|
|     NL|6668097|
+-------+-------+

Stage 3: 200 tasks
  max task duration  38s   (was 41m)
  median             21s

The hot key is spread across 20 partitions instead of 1. The cost must be stated: the small side is replicated 20×, so this trades memory and shuffle volume for balance. It is the right answer only when the skewed key genuinely has that much data and broadcasting is impossible.

4. Let adaptive execution do it

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB")

orders.join(customers, "customer_id").groupBy("country").count().count()

AQE: detected 1 skewed partition in stage 3, split into 24 sub-partitions
AQE: coalesced 200 shuffle partitions into 31
Stage 3: 224 tasks — max task 47s, median 6s
Total runtime: 1m 12s

AQE (Spark 3+, on by default) splits skewed partitions at runtime and coalesces small ones. It is the modern answer and it is not a substitute for the first fix — it made a 41-minute job take 72 seconds, while filtering the sentinel made it take 24. Say that: AQE mitigates skew, it does not remove the reason for it.

Questions that follow

“How many partitions should you have?”

print(f"default shuffle partitions: {spark.conf.get('spark.sql.shuffle.partitions')}")
print(f"orders partitions after read: {orders.rdd.getNumPartitions()}")

default shuffle partitions: 200
orders partitions after read: 16

The reasoning, not the number: aim for partitions of roughly 128-256 MB, and a task count that is a small multiple of total executor cores so the last wave is not half empty. 200 is a default that fits neither a 10 MB job nor a 10 TB one.

“repartition or coalesce?”

print("repartition(50):", orders.repartition(50).rdd.getNumPartitions())
print("coalesce(50):   ", orders.coalesce(50).rdd.getNumPartitions())

repartition(50): 50
coalesce(50):    16

coalesce(50) returned 16, because it only merges — it cannot increase the count. The trap worth naming: coalesce(1) before a write does not just merge at the end, it reduces the parallelism of the whole upstream stage to one task. Use repartition(1) if you truly need one file, and question whether you do.

“Why is cache() not helping?”

cached = orders.filter(F.col("amount") > 100).cache()
print(cached.count())
print(cached.storageLevel)

9998412
Disk Memory Deserialized 1x Replicated

Because caching is lazy — the first action populates it, and a job that reads the DataFrame once gains nothing while paying the memory. Cache when a DataFrame is used more than once, and unpersist() when done. That cache() is persist(MEMORY_AND_DISK), and that spilling to disk can make it slower than recomputation, is the senior detail.

“Narrow versus wide?”

Every performance question reduces to: how many wide transformations, and how much data crosses the network in each.

“How do you make a Spark job idempotent?”

(result.write
   .mode("overwrite")
   .option("partitionOverwriteMode", "dynamic")
   .partitionBy("run_date")
   .parquet("s3://warehouse/gold/daily_revenue"))

Dynamic partition overwrite replaces only the partitions present in the DataFrame, so a re-run for one day rewrites that day and leaves the rest. Without it, mode("overwrite") deletes the entire table — a genuinely famous production incident, and worth naming as the reason.

The scoring

The most common weak answer is “I’d increase executor memory”. It makes the straggler spill less and still runs for thirty minutes, because the problem is distribution, not capacity.

Practice

max_rows 6001204   median_rows 40   distinct_keys 500001

orders with customer_id = 0:     6,001,204
customers with customer_id = 0:  0

SortMergeJoin     → 3 Exchange nodes
BroadcastHashJoin → 1 Exchange node

repartition(50): 50 partitions
coalesce(50):    16 partitions

Next: streaming and CDC questions — exactly-once, watermarks, and the offsets that lose data.