PySpark & SQL Drills

from pyspark.sql import functions as F, Window as W

w = (W.partitionBy("account_id")
      .orderBy(F.col("updated_at").desc(), F.col("source_seq").desc()))

latest = (df
    .withColumn("rn", F.row_number().over(w))
    .filter("rn = 1")
    .drop("rn"))

What the interviewer is probing: do you know that dropDuplicates(["account_id"]) keeps an arbitrary row, while row_number with an explicit ordering is deterministic? Always add a tie-breaker column (source_seq) because timestamps collide.

Shuffle/complexity: one hash shuffle on the partition key, then a per-partition sort. Roughly O(n log n) within partitions. max_by / groupBy().agg() avoids the sort if you only need a couple of columns.

The trick: a changed key needs both an UPDATE (close) and an INSERT (open), but MERGE matches each source row once. So stage the source twice — once with the real key (drives the close) and once with a NULL merge key (never matches, so it inserts).

from delta.tables import DeltaTable
from pyspark.sql import functions as F

target = DeltaTable.forName(spark, "silver.dim_account")
current = target.toDF().filter("is_current")

changed = (src.alias("s")
    .join(current.alias("t"), "account_id")
    .where("s.cost_center <> t.cost_center OR s.owner <> t.owner")
    .select("s.*"))

staged = (src.withColumn("merge_key", F.col("account_id"))
    .unionByName(changed.withColumn("merge_key", F.lit(None))))

(target.alias("t")
 .merge(staged.alias("s"),
        "t.account_id = s.merge_key AND t.is_current = true")
 .whenMatchedUpdate(
     condition="s.cost_center <> t.cost_center OR s.owner <> t.owner",
     set={"is_current": "false", "valid_to": "s.effective_date"})
 .whenNotMatchedInsert(values={
     "account_id": "s.account_id", "cost_center": "s.cost_center",
     "owner": "s.owner", "valid_from": "s.effective_date",
     "valid_to": "NULL", "is_current": "true"})
 .execute())

Probing: do you understand MERGE semantics, idempotency, and why brand-new keys also flow through the NOT MATCHED branch? Mention hashing tracked attributes (xxhash64) instead of a long OR-chain for change detection.

Shuffle: the join and the merge each shuffle on the key; the merge also rewrites matched files, so partitioning/clustering the dimension by a stable column reduces write amplification.

First question to ask out loud: how big is the dimension? If it fits in executor memory, broadcast and skew is irrelevant:

joined = fact.join(F.broadcast(dim), "customer_id")

If both sides are large, enable AQE skew handling first (on by default in recent runtimes: spark.sql.adaptive.skewJoin.enabled), which splits oversized partitions at runtime. Manual salting is the fallback:

SALT = 16
fact_s = fact.withColumn("salt", (F.rand() * SALT).cast("int"))
dim_s  = dim.withColumn("salt",
            F.explode(F.array(*[F.lit(i) for i in range(SALT)])))
joined = fact_s.join(dim_s, ["customer_id", "salt"]).drop("salt")

Probing: do you reach for AQE before hand-rolling salting, and do you know the cost of salting (the small side is replicated SALT times)? Salting only the known-hot keys is the refined answer.

from pyspark.sql import types as T

schema = T.StructType([T.StructField("line_items", T.ArrayType(
    T.StructType([T.StructField("qty", T.IntegerType()),
                  T.StructField("unit_price", T.DecimalType(18, 2))])))])

totals = (raw
    .select("order_id", F.from_json("payload", schema).alias("p"))
    .select("order_id", F.explode_outer("p.line_items").alias("li"))
    .groupBy("order_id")
    .agg(F.sum(F.col("li.qty") * F.col("li.unit_price")).alias("order_total")))

Probing: explicit schema vs schema_of_json inference (explicit wins in production — inference is a separate job and drifts); explode_outer so orders with empty arrays survive; and the sharper follow-up — you don't need explode at all here: F.aggregate("p.line_items", ...) sums the array in place with zero row blow-up. Saying that unprompted is a strong signal.

Shuffle: explode multiplies rows before the groupBy shuffle; the higher-order-function version aggregates per row and shuffles nothing extra.

hwm = (spark.table("etl.watermarks")
       .filter("table_name = 'gl_postings'")
       .agg(F.max("hwm_value")).first()[0]) or "1900-01-01"

inc = (spark.read.format("jdbc").options(**jdbc_opts)
       .option("query",
               f"SELECT * FROM gl_postings WHERE updated_at > '{hwm}'")
       .load())

(DeltaTable.forName(spark, "bronze.gl_postings").alias("t")
 .merge(inc.alias("s"), "t.doc_id = s.doc_id AND t.line_id = s.line_id")
 .whenMatchedUpdateAll().whenNotMatchedInsertAll().execute())

new_hwm = inc.agg(F.max("updated_at")).first()[0]
# persist new_hwm to etl.watermarks ONLY after the merge succeeds

Probing: ordering of operations (advance the watermark after, never before, a successful write); using MERGE instead of append so reruns and overlap windows are idempotent; pushing the filter into the JDBC query rather than reading the whole table and filtering in Spark. Mention the classic failure: updated_at set by app servers with clock drift — subtract a safety overlap (e.g. reload the last 15 minutes) and let MERGE absorb the duplicates. If asked about file sources, note that Auto Loader / COPY INTO solve the same problem with file-level state — know them as concepts even if your production stack used watermark tables and ADF.

w = W.partitionBy("cost_center").orderBy(F.col("amount").desc())

top3 = (df.withColumn("rnk", F.dense_rank().over(w))
          .filter("rnk <= 3"))

Probing: the rank-family distinction. row_number breaks ties arbitrarily (exactly 3 rows), rank leaves gaps, dense_rank includes ties without gaps. Ask which behavior the "business" wants — that question is the point of the drill. Same single-shuffle profile as Drill 1.

months = ["2026-01", "2026-02", "2026-03"]
pivoted = (actuals.groupBy("cost_center")
           .pivot("month", months)      # pass values explicitly!
           .agg(F.sum("amount")))

unpivoted = pivoted.selectExpr(
    "cost_center",
    "stack(3, '2026-01', `2026-01`, '2026-02', `2026-02`, "
    "'2026-03', `2026-03`) AS (month, amount)")

Probing: passing the pivot value list explicitly skips the extra job Spark runs to discover distinct values and keeps the schema stable. For unpivot, stack() (or the SQL UNPIVOT clause in recent runtimes) — interviewers like hearing that tall/narrow is the modeling-friendly shape and wide pivots belong at the BI layer.

w = W.partitionBy("user_id").orderBy("event_ts")

sessions = (events
  .withColumn("prev_ts", F.lag("event_ts").over(w))
  .withColumn("is_new",
      F.when(F.col("prev_ts").isNull() |
             ((F.col("event_ts").cast("long") -
               F.col("prev_ts").cast("long")) > 1800), 1).otherwise(0))
  .withColumn("session_seq", F.sum("is_new").over(w))
  .withColumn("session_id",
      F.concat_ws("-", "user_id", "session_seq")))

Probing: the lag → flag → running-sum pattern (it reappears in gaps-and-islands below); handling the first event's NULL lag; and whether you mention that in streaming this becomes session_window() with a watermark. Two window passes over the same partitioning = one shuffle, two sorts.

SELECT
  cost_center,
  fiscal_period,
  SUM(amount) AS actual,
  SUM(SUM(amount)) OVER (
      PARTITION BY cost_center ORDER BY fiscal_period
      ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS ytd_actual,
  SUM(amount) - LAG(SUM(amount)) OVER (
      PARTITION BY cost_center ORDER BY fiscal_period) AS mom_variance
FROM gold.gl_actuals
GROUP BY cost_center, fiscal_period

Probing: windows evaluate after GROUP BY, so SUM(SUM(amount)) is legal — saying that sentence cleanly is half the points. Specify the frame: the default with ORDER BY is RANGE ... CURRENT ROW, which lumps peer rows together; ROWS is what you almost always mean. YTD needs a partition reset on fiscal year — add fiscal_year to the PARTITION BY and say so.

WITH flagged AS (
  SELECT account_id, activity_date,
         DATE_SUB(activity_date,
           ROW_NUMBER() OVER (PARTITION BY account_id
                              ORDER BY activity_date)) AS grp
  FROM daily_activity
)
SELECT account_id,
       MIN(activity_date) AS streak_start,
       MAX(activity_date) AS streak_end,
       COUNT(*)           AS streak_days
FROM flagged
GROUP BY account_id, grp

Probing: the insight that date - row_number is constant within a consecutive run. Walk a 4-row example out loud to prove you understand it rather than memorized it. Assumes deduped input — state that. One shuffle for the window, one for the group.

SELECT segment, business_unit, cost_center, expense,
  expense / SUM(expense) OVER ()                                   AS pct_company,
  expense / SUM(expense) OVER (PARTITION BY segment)               AS pct_segment,
  expense / SUM(expense) OVER (PARTITION BY segment, business_unit) AS pct_bu
FROM gold.expense_summary

Probing: multiple windows at different grains in one pass — no self-joins. Note the empty OVER () computes a global total (single-partition exchange for that one aggregate — fine for an aggregate, dangerous for an ordered window, per the warning above). Guard division by zero with NULLIF in production.

SELECT *
FROM bronze.jde_gl_extract
QUALIFY ROW_NUMBER() OVER (
    PARTITION BY doc_id, line_id
    ORDER BY extracted_at DESC) = 1

Probing: QUALIFY filters on window results the way HAVING filters on aggregates — it removes the wrapper subquery. Supported in Databricks SQL (and Snowflake, BigQuery, Teradata) but it is not universal ANSI; say "where QUALIFY isn't available I fall back to a CTE with row_number" to show portability awareness.

WITH spine AS (
  SELECT explode(sequence(DATE '2026-01-01', DATE '2026-12-01',
                          INTERVAL 1 MONTH)) AS period_start
),
cc AS (SELECT DISTINCT cost_center FROM gold.gl_actuals)
SELECT s.period_start, c.cost_center,
       COALESCE(a.amount, 0) AS amount
FROM spine s
CROSS JOIN cc c
LEFT JOIN gold.gl_actuals a
       ON a.fiscal_period = s.period_start
      AND a.cost_center  = c.cost_center

Probing: recognizing that LAG/variance silently skips missing periods unless you densify first — a wrong variance number is worse than a missing one in finance. The CROSS JOIN is deliberate and tiny (12 × cost centers); saying "this cartesian is intentional and bounded" preempts the obvious follow-up.

collect() pulls the whole dataset to the driver — OOM at scale, and the aggregation runs single-threaded in Python. Fix: df.agg(F.sum("amount")).first()[0] — the cluster aggregates, one scalar returns. The probe: do you keep computation on executors and reserve collect/toPandas for provably small results?

A Python UDF for logic that built-ins cover: F.upper(F.trim("name")). The UDF forces JVM↔Python serialization per row, blocks codegen and Photon, and is opaque to Catalyst. Escalation ladder to recite: built-ins → higher-order functions → pandas_udf (vectorized via Arrow) → row-wise Python UDF as last resort.

repartition(1) funnels the entire write through one task — one core does all the work and one giant file lands. Let the engine handle file sizing (optimized writes / AQE partition coalescing, plus OPTIMIZE after) instead of forcing single-file output. Bonus distinction: repartition always shuffles; coalesce narrows without a shuffle but can starve upstream parallelism because it propagates back through the plan.

The typo compares a column to itself — true for every non-null key, so the join degenerates into a near-cartesian product (rows × rows). Spark may surface it as CartesianProduct or BroadcastNestedLoopJoin in the plan. Fix the condition (fact.cust_id == dim.cust_id); the senior habit is running .explain() and checking the join node before launching anything expensive. If a cross join is ever intended (Drill 5's date spine), write crossJoin() explicitly so reviewers see intent.

# dedupe preserving order (dicts are insertion-ordered)
unique = list(dict.fromkeys(items))

# schema diff between two DataFrames
missing = set(expected_cols) - set(df.columns)
extra   = set(df.columns) - set(expected_cols)

# frequency count
from collections import Counter
top = Counter(codes).most_common(5)

Know the complexity story: set/dict membership is O(1) average vs O(n) for a list — the one-line justification interviewers want.

def read_batches(path, size=10_000):
    batch = []
    for line in open(path):          # lazy: one line in memory at a time
        batch.append(parse(line))
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch

total = sum(r.amount for r in rows)  # generator expr: no list built

The point to articulate: generators trade memory for laziness — the same idea as Spark's lazy evaluation, which makes a tidy bridge back to distributed thinking.