PySpark Practice in Your Browser

df.filter(F.col("amount") > 100)
  .withColumn("is_weekend",
              F.dayofweek("date").isin(1, 7))
  .select("user_id", "is_weekend", "amount")
  .orderBy(F.desc("amount"))

select, filter, withColumn, drop, alias, when/otherwise. Warm-up tier; fluency here keeps the rest of the interview moving.

df.groupBy("category", "month")
  .agg(
    F.sum("revenue").alias("rev"),
    F.countDistinct("customer_id").alias("uniq"),
    F.sum(F.when(F.col("is_new"), F.col("revenue"))
           .otherwise(0)).alias("new_cust_rev"),
  )

agg with named columns, conditional aggregation, distinct counts. The DataFrame equivalent of GROUP BY + HAVING + CASE WHEN.

w = (Window
      .partitionBy("user_id")
      .orderBy(F.desc("event_at"), F.desc("event_id")))
df.withColumn("rn", F.row_number().over(w))
  .filter("rn = 1")
  .drop("rn")

Window.partitionBy, orderBy, row_number, lag, lead, running totals. PySpark window syntax differs from SQL in subtle ways the evaluator checks.

events.join(
    F.broadcast(users), "user_id", "left"
).join(
    products, ["product_id"], "inner"
)

Inner, left, full, semi, anti joins. broadcast() hint for small dimensions. Default threshold is 10MB; the evaluator expects you to know that or look it up before guessing.

# salt the hot key
events_salted = events.withColumn(
    "salt", F.rand() * 20).cast("int")
users_salted = users.crossJoin(
    F.broadcast(spark.range(20).select(
        F.col("id").alias("salt"))))
events_salted.join(users_salted,
    ["user_id", "salt"], "left")

Salting hot keys, repartition vs coalesce, AQE settings, bucketed tables. The senior-level differentiator.

# read physical plan
df.explain(mode="formatted")
# inspect via Spark UI metrics:
#   task durations, shuffle read/write,
#   GC time, executor heap

Read EXPLAIN output, identify shuffles, diagnose OOM. Problems provide Spark UI screenshots; you write the fix.

from pyspark.storagelevel import StorageLevel
df.persist(StorageLevel.MEMORY_AND_DISK)
df.count()  # trigger materialization
# downstream uses df 3 times
df.unpersist()

When to cache, what storage level, when to checkpoint instead, when to refactor and avoid the reuse.