PRACTICE / LESSON 19

The PySpark coding round: six patterns

Six problems that recur in 30–45 minute data engineering coding rounds — each solved, run, and costed, with the wrong answer that looks right.

25 min read + practicePySpark3.5.7 reference

THE IDEA TO KEEP

Most coding-round problems are a window, an aggregation, or a generator applied to a precisely stated grain. Say the grain, the tie rule, the null rule, and the time boundary before you code, then name the operator that shuffles.

How to work the round

In a 30–45 minute round the code is the last third of the work. Before typing, pin down the contract out loud:

Grain and keys — one output row per what, and which columns identify a record?
Ordering and ties — “latest” by which column, and what breaks a tie when two rows share it?
Nulls — null keys, null timestamps, null or empty arrays: dropped, kept, or sorted last?
Time — event time or ingest time, which time zone, and are boundaries inclusive? Is a gap of exactly 30 minutes a new session?

Then build a five-to-eight-row input with spark.createDataFrame that contains each edge case, say the expected output before you run it, and only then write the transformation. Finish with the cost: which operator shuffles (an Exchange in explain()), what is sorted, and whether anything collapses into a single partition.

Every block below was run exactly as shown on local PySpark 3.5.7 (local[2], JDK 17) with an existing spark session, spark.sql.session.timeZone set to UTC, and spark.sql.shuffle.partitions set to 4; the SCD block adds Delta Lake 3.2.1, and one demonstration also ran on a two-executor local cluster. Outputs are copied from those runs. Plans quoted in prose come from the same runtime with adaptive query execution (AQE) at its default, enabled, unless stated.

Check the sources: Window functions reference

1. Deduplicate, keeping the latest row per key

Problem. A customer feed repeats customer_id. Keep one row per customer: the newest by updated_at; when two rows share a timestamp, the higher ingest_id wins; a null timestamp loses to any real one. The tiny input contains all three cases.

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

customers = spark.createDataFrame([
    (1, "a@old.example", "2026-09-01 09:00:00", 101),
    (1, "a@new.example", "2026-09-03 18:30:00", 102),
    (2, "b@one.example", "2026-09-02 12:00:00", 201),
    (2, "b@two.example", "2026-09-02 12:00:00", 202),  # same timestamp: a tie
    (3, "c@only.example", None, 301),                 # null timestamp
], "customer_id int, email string, updated_at string, ingest_id int"
).withColumn("updated_at", F.to_timestamp("updated_at"))

latest_first = Window.partitionBy("customer_id").orderBy(
    F.col("updated_at").desc_nulls_last(), F.col("ingest_id").desc())

latest = (customers
    .withColumn("rn", F.row_number().over(latest_first))
    .filter(F.col("rn") == 1)
    .drop("rn"))
latest.orderBy("customer_id").show(truncate=False)

# Same answer without a window: max over a struct compares fields left to right
by_struct = (customers
    .groupBy("customer_id")
    .agg(F.max(F.struct("updated_at", "ingest_id", "email")).alias("r"))
    .select("customer_id", "r.email", "r.updated_at", "r.ingest_id"))
by_struct.orderBy("customer_id").show(truncate=False)

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7

+-----------+--------------+-------------------+---------+
|customer_id|email         |updated_at         |ingest_id|
+-----------+--------------+-------------------+---------+
|1          |a@new.example |2026-09-03 18:30:00|102      |
|2          |b@two.example |2026-09-02 12:00:00|202      |
|3          |c@only.example|NULL               |301      |
+-----------+--------------+-------------------+---------+

+-----------+--------------+-------------------+---------+
|customer_id|email         |updated_at         |ingest_id|
+-----------+--------------+-------------------+---------+
|1          |a@new.example |2026-09-03 18:30:00|102      |
|2          |b@two.example |2026-09-02 12:00:00|202      |
|3          |c@only.example|NULL               |301      |
+-----------+--------------+-------------------+---------+

Cost. The window is one hash shuffle on customer_id and a sort inside each partition. In Spark 3.5.7 the optimizer recognizes the rn == 1 filter and inserts WindowGroupLimit twice — Partial below the Exchange and Final above it (observed in explain() on an equivalent query) — so each task forwards only its local winners. The struct version is an aggregation: max over a struct compares fields left to right, so the field order is the tie rule, and in a separate check a struct whose first field was null lost to one with a real timestamp. Its plan was a SortAggregate with a partial max before the shuffle. max_by(col, ord), available since 3.3, is the one-column shortcut; its documentation says nothing about ties on ord, so pass a struct as ord when ties can happen.

The common wrong answer is df.orderBy(F.col("ts").desc()).dropDuplicates(["key"]). For a batch DataFrame the API documentation says only that it drops duplicate rows, not which one survives. In the 3.5.7 optimizer, ReplaceDeduplicateWithAggregate turns deduplication into an aggregate that keeps First of every non-key column. The observed plan is the global sort, a partial_first in each task, a hash Exchange on the key, and a final first: the winner depends on the order in which the final aggregate reads shuffle output, and the earlier sort does not control that. The block below makes it visible. It disables AQE itself, and its timestamps are unique, so no tie excuses the result.

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

spark.conf.set("spark.sql.adaptive.enabled", "false")
events = (spark.range(0, 200000, 1, 8)
    .withColumn("key", F.col("id") % 40000)
    .withColumn("ts", (F.col("id") * 7919) % 1000003))  # timestamps are unique

expected = (events
    .withColumn("rn", F.row_number().over(
        Window.partitionBy("key").orderBy(F.col("ts").desc())))
    .filter("rn = 1")
    .select("key", F.col("id").alias("expected_id")))

risky = events.orderBy(F.col("ts").desc()).dropDuplicates(["key"])

wrong = risky.join(expected, "key").filter("id <> expected_id").count()
print("keys:", expected.count(), "| dropDuplicates kept a non-latest row for:", wrong)

local[2]: keys: 40000 | dropDuplicates kept a non-latest row for: 0
local-cluster[2,1,1024], run 1: keys: 40000 | dropDuplicates kept a non-latest row for: 29351
local-cluster[2,1,1024], run 2: keys: 40000 | dropDuplicates kept a non-latest row for: 19926

The single-process run matched on every key, which is how this bug survives a unit test. With AQE left on, two runs of a close variant on the same two-executor cluster also matched; that is an observation, not a guarantee the API makes.

Check the sources: Window functions · dropDuplicates · Optimizer source · InferWindowGroupLimit · max_by

2. Sessionize events with a 30-minute inactivity gap

Problem. Give each event a session ID: a session starts at a user's first event or after more than 30 minutes without one. Decide the boundary aloud. Here a gap of exactly 30 minutes stays in the session (events 2 and 3), and a 40-minute gap splits it (events 3 and 4).

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

events = spark.createDataFrame([
    ("u1", 1, "2026-09-01 10:00:00"),
    ("u1", 2, "2026-09-01 10:20:00"),
    ("u1", 3, "2026-09-01 10:50:00"),  # exactly 30 minutes later: same session
    ("u1", 4, "2026-09-01 11:30:00"),  # 40 minutes later: new session
    ("u2", 5, "2026-09-01 10:05:00"),
    ("u2", 6, "2026-09-01 12:00:00"),
], "user_id string, event_id int, ts string").withColumn("ts", F.to_timestamp("ts"))

by_user = Window.partitionBy("user_id").orderBy("ts", "event_id")
gap_seconds = F.col("ts").cast("long") - F.lag("ts").over(by_user).cast("long")

sessions = (events
    .withColumn("new_session", F.when(gap_seconds.isNull() | (gap_seconds > 30 * 60), 1).otherwise(0))
    .withColumn("session_no", F.sum("new_session").over(
        by_user.rowsBetween(Window.unboundedPreceding, Window.currentRow)))
    .withColumn("session_id", F.concat_ws("-", "user_id", F.col("session_no").cast("string"))))
sessions.select("user_id", "event_id", "ts", "new_session", "session_id").orderBy("event_id").show()

# Built-in alternative (3.2+): one row per session, end = last event + gap
(events
    .groupBy("user_id", F.session_window("ts", "30 minutes").alias("s"))
    .agg(F.count("*").alias("events"))
    .select("user_id", "s.start", "s.end", "events")
    .orderBy("user_id", "start")
    .show())

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7

+-------+--------+-------------------+-----------+----------+
|user_id|event_id|                 ts|new_session|session_id|
+-------+--------+-------------------+-----------+----------+
|     u1|       1|2026-09-01 10:00:00|          1|      u1-1|
|     u1|       2|2026-09-01 10:20:00|          0|      u1-1|
|     u1|       3|2026-09-01 10:50:00|          0|      u1-1|
|     u1|       4|2026-09-01 11:30:00|          1|      u1-2|
|     u2|       5|2026-09-01 10:05:00|          1|      u2-1|
|     u2|       6|2026-09-01 12:00:00|          1|      u2-2|
+-------+--------+-------------------+-----------+----------+

+-------+-------------------+-------------------+------+
|user_id|              start|                end|events|
+-------+-------------------+-------------------+------+
|     u1|2026-09-01 10:00:00|2026-09-01 11:20:00|     3|
|     u1|2026-09-01 11:30:00|2026-09-01 12:00:00|     1|
|     u2|2026-09-01 10:05:00|2026-09-01 10:35:00|     1|
|     u2|2026-09-01 12:00:00|2026-09-01 12:30:00|     1|
+-------+-------------------+-------------------+------+

How it works and what it costs. lag reads the previous timestamp within the user; the flag is 1 at a first event or after a gap over 1,800 seconds; a running sum of flags numbers the sessions. Both window functions share one specification, so the plan is a single hash Exchange on the user, one sort, and one Window operator. A very active user still lands in one task: sessionization inherits key skew. The frame is explicit on purpose. With an ordering and no frame, PySpark uses a range frame from unbounded preceding to the current row, which includes every row with an equal ordering value; event_id in the ordering and a rows frame keep the numbering deterministic when timestamps collide.

The built-in alternative. session_window(ts, "30 minutes") (3.2+) groups events into sessions that end at the latest event plus the gap: the observed first session ends at 11:20, which is 10:50 plus 30 minutes. On this input it agreed with the window version, including the exactly-30-minute gap. It returns one row per session rather than labeling events, and the gap can be a column expression evaluated per row. In Structured Streaming, the guide's example pairs it with a watermark; Update output mode is not supported, and the grouping key needs at least one column besides the session window (batch queries may group by the session window alone). The lag version cannot run on a stream: starting it failed with NON_TIME_WINDOW_NOT_SUPPORTED_IN_STREAMING.

The common wrong answer is a fixed time bucket. Buckets cut on the clock, not on inactivity: a 30-minute tumbling window splits clicks at 10:29 and 10:31 into two sessions, and an hourly groupBy merges visits at 10:00 and 10:50 that were 50 minutes apart. The other frequent slip is stopping at the gap flag, which marks where sessions start but gives the rows in between no ID until you take the running sum.

Check the sources: Window default frames · session_window · Structured Streaming: session windows

3. A running 7-day active user count

Problem. For every date, count the distinct users with any activity on that date or the six days before it. The input has gaps (nothing on 09-03, 09-04, or 09-06 to 09-08) and a user active on consecutive days, because those are exactly what break the easy answers.

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

activity = spark.createDataFrame([
    ("u1", "2026-09-01"), ("u2", "2026-09-01"), ("u1", "2026-09-02"),
    ("u3", "2026-09-05"), ("u1", "2026-09-09"),
    ("u1", "2026-09-10"), ("u2", "2026-09-10"), ("u4", "2026-09-10"),
], "user_id string, d string").withColumn("d", F.to_date("d"))

try:
    activity.withColumn("n", F.countDistinct("user_id").over(Window.orderBy("d").rowsBetween(-6, 0)))
except Exception as e:
    print(str(e).split(";")[0])

# Correct: credit each active user-day to the 7 as-of dates whose window contains it
user_days = activity.select("user_id", "d").distinct()
credited = user_days.withColumn("as_of", F.explode(F.sequence("d", F.date_add("d", 6))))
active = credited.groupBy("as_of").agg(F.countDistinct("user_id").alias("active_7d"))
spine = (activity.agg(F.min("d").alias("lo"), F.max("d").alias("hi"))
    .select(F.explode(F.sequence("lo", "hi")).alias("as_of")))
correct = spine.join(active, "as_of", "left").fillna(0, ["active_7d"])

# Wrong: sum daily distinct counts over 7 rows, or over a 7-day range
by_day = Window.orderBy("day_no")
daily = (activity.groupBy("d").agg(F.countDistinct("user_id").alias("dau"))
    .withColumn("day_no", F.datediff("d", F.lit("1970-01-01")))
    .withColumn("rows_sum", F.sum("dau").over(by_day.rowsBetween(-6, 0)))
    .withColumn("range_sum", F.sum("dau").over(by_day.rangeBetween(-6, 0))))

(correct
    .join(daily.select(F.col("d").alias("as_of"), "dau", "rows_sum", "range_sum"), "as_of", "left")
    .orderBy("as_of")
    .show())

# Also correct, on active days only: a set over a 7-day range frame (one global partition)
(user_days
    .withColumn("day_no", F.datediff("d", F.lit("1970-01-01")))
    .withColumn("active_7d", F.size(F.collect_set("user_id").over(by_day.rangeBetween(-6, 0))))
    .select("d", "active_7d").distinct()
    .orderBy("d")
    .show())

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7

[DISTINCT_WINDOW_FUNCTION_UNSUPPORTED] Distinct window functions are not supported: "count(DISTINCT user_id) OVER (ORDER BY d ASC NULLS FIRST ROWS BETWEEN -6 FOLLOWING AND CURRENT ROW)".
+----------+---------+----+--------+---------+
|     as_of|active_7d| dau|rows_sum|range_sum|
+----------+---------+----+--------+---------+
|2026-09-01|        2|   2|       2|        2|
|2026-09-02|        2|   1|       3|        3|
|2026-09-03|        2|NULL|    NULL|     NULL|
|2026-09-04|        2|NULL|    NULL|     NULL|
|2026-09-05|        3|   1|       4|        4|
|2026-09-06|        3|NULL|    NULL|     NULL|
|2026-09-07|        3|NULL|    NULL|     NULL|
|2026-09-08|        2|NULL|    NULL|     NULL|
|2026-09-09|        2|   1|       5|        2|
|2026-09-10|        4|   3|       8|        5|
+----------+---------+----+--------+---------+

+----------+---------+
|         d|active_7d|
+----------+---------+
|2026-09-01|        2|
|2026-09-02|        2|
|2026-09-05|        3|
|2026-09-09|        2|
|2026-09-10|        4|
+----------+---------+

Reading the result. The first line is the exact error Spark 3.5.7 raises for a distinct count over a window frame: the literal translation of the requirement does not analyze. active_7d is correct. rows_sum adds the last seven rows of daily counts; with gaps those rows reach back more than seven days, so 09-09 reports 5 against a true 2. range_sum uses rangeBetween(-6, 0) on a day number, which fixes the calendar but still adds per-day distinct counts: u1 on 09-09 and 09-10 is counted twice, giving 5 against 4 (and 3 against 2 on 09-02). A range frame with an offset needs a single numeric ordering expression, which is why the date becomes day_no.

Cost. The explode approach deduplicates user-days, multiplies them by seven, and groups by as_of. For the distinct count the plan has two aggregation shuffles — one on (as_of, user_id), then one on as_of — and no global sort, so it spreads across executors. The spine join adds dates without events, drops the six trailing dates the explode creates past the last day, and turns an empty week into 0. The collect_set version is correct on active days, but a window with no partitionBy moves every row into one partition: the run logged No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation. Each row also carries that week's set of users. Use it only for small inputs, or partition it by a column such as tenant.

The common wrong answer is either sum above, or a countDistinct over a window that Spark refuses. State the definition first: daily active users do not add up to weekly active users.

Check the sources: rangeBetween · Window frames · Error classes

4. SCD Type 2 with a Delta MERGE

Problem. Maintain dim_customer with valid_from, valid_to, and is_current. Each batch carries (customer_id, city, effective_date). An unchanged row does nothing; a changed city closes the current version and inserts a new one; a new key inserts a current row. The block assumes a session with Delta Lake 3.2.1's SQL extension and catalog configured.

The trick. A changed key needs two actions in one MERGE — an update and an insert — but each source row either matches or it does not. So the source is staged as a union: every batch row keyed by merge_key = customer_id, which matches the current version and closes it if the city changed, plus an extra copy of each changed row with a NULL merge_key. NULL never equals anything, so that copy falls through to whenNotMatchedInsert and becomes the new version. A new key has no current row, so its keyed copy is inserted.

PYTHON / PYSPARK

import tempfile
from datetime import date
from delta.tables import DeltaTable
from pyspark.sql import functions as F

path = tempfile.mkdtemp() + "/dim_customer"
(spark.createDataFrame([], "customer_id int, city string, valid_from date, valid_to date, is_current boolean")
    .write.format("delta").save(path))

def apply_scd2(updates):
    dim = DeltaTable.forPath(spark, path)
    current = dim.toDF().filter("is_current")
    # Changed keys need a second row: a NULL merge_key never matches, so it is inserted
    new_versions = (updates.alias("u")
        .join(current.alias("c"), F.col("u.customer_id") == F.col("c.customer_id"))
        .filter(~F.col("u.city").eqNullSafe(F.col("c.city")))
        .select(F.lit(None).cast("int").alias("merge_key"), "u.*"))
    staged = new_versions.unionByName(
        updates.select(F.col("customer_id").alias("merge_key"), "*"))
    (dim.alias("t")
        .merge(staged.alias("s"), "t.customer_id = s.merge_key AND t.is_current")
        .whenMatchedUpdate(
            condition="NOT (t.city <=> s.city)",
            set={"is_current": "false", "valid_to": "s.effective_date"})
        .whenNotMatchedInsert(values={
            "customer_id": "s.customer_id", "city": "s.city",
            "valid_from": "s.effective_date", "valid_to": "null", "is_current": "true"})
        .execute())

batch = "customer_id int, city string, effective_date date"
load_1 = spark.createDataFrame([(1, "Austin", date(2026, 9, 1)), (2, "Denver", date(2026, 9, 1))], batch)
load_2 = spark.createDataFrame([
    (1, "Austin", date(2026, 9, 10)),   # unchanged
    (2, "Boston", date(2026, 9, 10)),   # changed attribute
    (3, "Chicago", date(2026, 9, 10)),  # new key
], batch)

apply_scd2(load_1)
apply_scd2(load_2)
spark.read.format("delta").load(path).orderBy("customer_id", "valid_from").show()

apply_scd2(load_2)  # replaying the same batch
m = F.col("operationMetrics")
(DeltaTable.forPath(spark, path).history()
    .select("version", "operation",
            m["numTargetRowsInserted"].alias("inserted"), m["numTargetRowsUpdated"].alias("updated"))
    .orderBy("version")
    .show())

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7 · DELTA LAKE 3.2.1

+-----------+-------+----------+----------+----------+
|customer_id|   city|valid_from|  valid_to|is_current|
+-----------+-------+----------+----------+----------+
|          1| Austin|2026-09-01|      NULL|      true|
|          2| Denver|2026-09-01|2026-09-10|     false|
|          2| Boston|2026-09-10|      NULL|      true|
|          3|Chicago|2026-09-10|      NULL|      true|
+-----------+-------+----------+----------+----------+

+-------+---------+--------+-------+
|version|operation|inserted|updated|
+-------+---------+--------+-------+
|      0|    WRITE|    NULL|   NULL|
|      1|    MERGE|       2|      0|
|      2|    MERGE|       2|      1|
|      3|    MERGE|       0|      0|
+-------+---------+--------+-------+

Reading the result. After two loads, customer 2 has a closed Denver row ending 2026-09-10 and a current Boston row; customer 1 is untouched; customer 3 is new. The history shows the second MERGE inserted 2 rows and updated 1, and replaying the same batch inserted and updated nothing, because unchanged rows match but fail the update condition. Each load is one table version, so a reader never sees a closed row without its replacement.

Details that separate a senior answer. Compare with <=> (eqNullSafe), not <>: the Delta documentation's SCD example uses <>, which is never true when either side is null, so a change to or from null would be missed. Put t.is_current in the match condition so closed history rows never match. Deduplicate the batch to one row per key first (problem 1): a separate run with two source rows for one target row failed with DELTA_MULTIPLE_SOURCE_ROW_MATCHING_TARGET_ROW_IN_MERGE. Replay is safe only for the same batch; a late batch with an older effective_date would close the current row with an earlier date, so apply batches in order or add a date guard to the conditions. Cost: the staging join against current rows is a shuffle, or a broadcast when the batch is small, and by default MERGE searches the entire target table for matches; the Delta documentation's tuning advice is to add known constraints, such as partition predicates, to the match condition.

The common wrong answer is a MERGE on customer_id alone with an update clause and an insert clause: a changed key matches, is closed, and its new version is never inserted, because a matched source row cannot also take the not-matched branch. The other is an UPDATE followed by a separate INSERT — two commits, with a moment in which the key has no current row, and a failure between them leaves it that way.

Check the sources: Delta Lake 3.2.1: merge and SCD Type 2

5. Top N per group

Problem. Return the top two reps by amount in each region. East has a tie at 90, so state the tie rule before choosing a function.

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

sales = spark.createDataFrame([
    ("east", "ann", 90), ("east", "bo", 90), ("east", "cy", 80), ("east", "di", 70),
    ("west", "ed", 95), ("west", "fay", 60),
], "region string, rep string, amount int")

by_amount = Window.partitionBy("region").orderBy(F.col("amount").desc())
by_amount_then_rep = Window.partitionBy("region").orderBy(F.col("amount").desc(), "rep")

(sales.select("*",
        F.row_number().over(by_amount_then_rep).alias("row_number"),
        F.rank().over(by_amount).alias("rank"),
        F.dense_rank().over(by_amount).alias("dense_rank"))
    .orderBy("region", "row_number")
    .show())

top2 = (sales
    .withColumn("rn", F.row_number().over(by_amount_then_rep))
    .filter(F.col("rn") <= 2))
top2.orderBy("region", "rn").show()

print(sales.orderBy(F.col("amount").desc(), "rep").limit(2).collect())

top2.explain()

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7

+------+---+------+----------+----+----------+
|region|rep|amount|row_number|rank|dense_rank|
+------+---+------+----------+----+----------+
|  east|ann|    90|         1|   1|         1|
|  east| bo|    90|         2|   1|         1|
|  east| cy|    80|         3|   3|         2|
|  east| di|    70|         4|   4|         3|
|  west| ed|    95|         1|   1|         1|
|  west|fay|    60|         2|   2|         2|
+------+---+------+----------+----+----------+

+------+---+------+---+
|region|rep|amount| rn|
+------+---+------+---+
|  east|ann|    90|  1|
|  east| bo|    90|  2|
|  west| ed|    95|  1|
|  west|fay|    60|  2|
+------+---+------+---+

[Row(region='west', rep='ed', amount=95), Row(region='east', rep='ann', amount=90)]
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Filter (rn#58 <= 2)
   +- Window [row_number() windowspecdefinition(region#0, amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) AS rn#58], [region#0], [amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST]
      +- WindowGroupLimit [region#0], [amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST], row_number(), 2, Final
         +- Sort [region#0 ASC NULLS FIRST, amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST], false, 0
            +- Exchange hashpartitioning(region#0, 4), ENSURE_REQUIREMENTS, [plan_id=177]
               +- WindowGroupLimit [region#0], [amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST], row_number(), 2, Partial
                  +- Sort [region#0 ASC NULLS FIRST, amount#2 DESC NULLS LAST, rep#1 ASC NULLS FIRST], false, 0
                     +- Scan ExistingRDD[region#0,rep#1,amount#2]

Choosing the ranking function. row_number gives exactly N rows per group, but which tied row wins is arbitrary unless the ordering includes a tie-breaker (rep here). rank gives ties the same number and skips the next (1, 1, 3), so rank <= 2 would return ann and bo for east. dense_rank does not skip (1, 1, 2), so dense_rank <= 2 would return ann, bo, and cy — more than N rows. Pick by the business rule: exactly N, every row tied at the cutoff, or the top N distinct values.

Cost. The plan above is one hash Exchange on region with a sort on each side. Spark 3.5.7's InferWindowGroupLimit rule recognized rn <= 2 and inserted WindowGroupLimit as Partial before the shuffle and Final after it, so tasks forward at most two rows per region instead of every row. The rule covers rank-like functions filtered by limit-like predicates and is controlled by the internal setting spark.sql.optimizer.windowGroupLimitThreshold (default 1000).

The common wrong answer is sales.orderBy(F.col("amount").desc()).limit(2). It returns the top two rows of the whole DataFrame — the printed rows are west's ed and east's ann, and every other row in east is gone. Cost is not the problem: in a separate check, orderBy(...).limit(2) planned as TakeOrderedAndProject(limit=2, ...), a top-N selection rather than a full global sort. Looping over regions with a filter and limit for each is correct but runs one query per group.

Check the sources: Ranking functions · InferWindowGroupLimit

6. Flatten nested JSON and arrays of structs

Problem. Orders arrive as JSON with a nested customer struct and an items array of structs. Produce one row per item with the customer fields, keeping orders whose array is empty or missing.

PYTHON / PYSPARK

from pyspark.sql import functions as F

lines = [
    '{"order_id": 1, "customer": {"id": "c1", "address": {"city": "Austin"}},'
    ' "items": [{"sku": "A", "qty": 2}, {"sku": "B", "qty": 1}]}',
    '{"order_id": 2, "customer": {"id": "c2", "address": null}, "items": []}',
    '{"order_id": 3, "customer": {"id": "c3"}}',
]
schema = ("order_id INT, customer STRUCT<id: STRING, address: STRUCT<city: STRING>>, "
          "items ARRAY<STRUCT<sku: STRING, qty: INT>>")

# Schema first: no inference pass; in production this is .json("s3://.../orders/")
orders = spark.read.schema(schema).json(spark.sparkContext.parallelize(lines))

flat = (orders
    .select("order_id",
            F.col("customer.id").alias("customer_id"),
            F.col("customer.address.city").alias("city"),
            F.posexplode_outer("items").alias("pos", "item"))
    .select("order_id", "customer_id", "city", "pos", "item.*"))
flat.orderBy("order_id", "pos").show()

print("explode rows:", orders.select(F.explode("items")).count(),
      "| explode_outer rows:", orders.select(F.explode_outer("items")).count())
orders.select("order_id", F.inline_outer("items")).orderBy("order_id").show()
orders.select("order_id", "customer.*").orderBy("order_id").show()

OBSERVED OUTPUT · LOCAL PYSPARK 3.5.7

+--------+-----------+------+----+----+----+
|order_id|customer_id|  city| pos| sku| qty|
+--------+-----------+------+----+----+----+
|       1|         c1|Austin|   0|   A|   2|
|       1|         c1|Austin|   1|   B|   1|
|       2|         c2|  NULL|NULL|NULL|NULL|
|       3|         c3|  NULL|NULL|NULL|NULL|
+--------+-----------+------+----+----+----+

explode rows: 2 | explode_outer rows: 4
+--------+----+----+
|order_id| sku| qty|
+--------+----+----+
|       1|   A|   2|
|       1|   B|   1|
|       2|NULL|NULL|
|       3|NULL|NULL|
+--------+----+----+

+--------+---+--------+
|order_id| id| address|
+--------+---+--------+
|       1| c1|{Austin}|
|       2| c2|    NULL|
|       3| c3|    NULL|
+--------+---+--------+

How it works. The DDL string fixes the schema, so Spark does not run an inference pass over the input (by default inference uses every input object: samplingRatio is 1.0), and a missing field becomes null rather than changing the schema. Dotted paths reach into structs: customer.address.city is null when address is null (order 2) or absent (order 3), and customer.* expands one level. explode produced 2 rows because it drops parents whose array is empty or null; explode_outer produced 4. posexplode_outer adds the element position, useful for keeping item order or deduplicating line items, and inline_outer turns an array of structs straight into columns. For malformed input, the default PERMISSIVE mode sets unparsable fields to null and keeps the raw record only if the schema includes the string column named by columnNameOfCorruptRecord.

Cost. Generators are narrow: an explode_outer plan in a separate check was a single Generate operator with no Exchange. The cost is row multiplication. Exploding two independent arrays one after the other yields their cross product, so explode one array per output table, or zip related arrays with arrays_zip and explode once.

The common wrong answer is plain explode on an optional array, which silently removes orders 2 and 3, often combined with schema inference on every run, so that one unusual file can change column types downstream.

Check the sources: JSON data source · explode_outer · inline_outer

Try it on a small dataset

Run this in a PySpark notebook with an existing spark session. Create a local session if needed. The example was run for review; this website does not execute Spark.

PYTHON / PYSPARK

from pyspark.sql import functions as F, Window

customers = spark.createDataFrame([
    (1, "a@old.example", "2026-09-01 09:00:00", 101),
    (1, "a@new.example", "2026-09-03 18:30:00", 102),
    (2, "b@one.example", "2026-09-02 12:00:00", 201),
    (2, "b@two.example", "2026-09-02 12:00:00", 202),  # same timestamp: a tie
    (3, "c@only.example", None, 301),                 # null timestamp
], "customer_id int, email string, updated_at string, ingest_id int"
).withColumn("updated_at", F.to_timestamp("updated_at"))

latest_first = Window.partitionBy("customer_id").orderBy(
    F.col("updated_at").desc_nulls_last(), F.col("ingest_id").desc())

latest = (customers
    .withColumn("rn", F.row_number().over(latest_first))
    .filter(F.col("rn") == 1)
    .drop("rn"))
latest.orderBy("customer_id").show(truncate=False)

# Same answer without a window: max over a struct compares fields left to right
by_struct = (customers
    .groupBy("customer_id")
    .agg(F.max(F.struct("updated_at", "ingest_id", "email")).alias("r"))
    .select("customer_id", "r.email", "r.updated_at", "r.ingest_id"))
by_struct.orderBy("customer_id").show(truncate=False)

Observed result

Both tables printed the same three rows: customer 1 kept a@new.example (2026-09-03 18:30:00, ingest_id 102); customer 2 kept b@two.example, the 2026-09-02 12:00:00 tie broken by the higher ingest_id 202; customer 3 kept c@only.example with a NULL updated_at, its only row. The tie-breaker is what makes the result deterministic; without it, customer 2's survivor is arbitrary.

Ran on local PySpark 3.5.7 (local[2], JDK 17) on 2026-09-15; output above is observed. The lesson's other blocks ran on the same runtime; the SCD block also used Delta Lake 3.2.1, and the dropDuplicates demonstration also ran on a two-executor local-cluster master.

Watch the assumption

Trusting a test that passed in local mode. orderBy followed by dropDuplicates returned the latest row for every key in local[2], and a non-latest row for tens of thousands of keys on a two-executor local cluster with AQE disabled. Put the tie, the null, and the gap into the tiny input, and prefer operators whose result is defined by the query rather than by the order in which tasks and shuffle blocks happen to run.

EXPLAIN IT BACK

Your deduplication used orderBy(...).dropDuplicates([key]) and its unit test passed. Why can it still keep the wrong row, and what do you write instead?

Say your answer out loud before opening the notes.

Reveal the reasoning

dropDuplicates promises that duplicates are removed, not which row survives. Spark 3.5.7 rewrites it into an aggregate that takes first() of every non-key column; that aggregate runs partially in each task, shuffles by key, and combines partial results in the order the shuffle output is read, so the preceding sort does not pick the winner. A single-process test can pass by accident: in local[2] every key matched, while a two-executor local cluster with AQE disabled kept a non-latest row for 29,351 and then 19,926 of 40,000 keys. Write row_number over a window partitioned by the key and ordered by timestamp descending plus a unique tie-breaker, and keep row 1, or take max over a struct whose field order encodes the same rule.

Sources & verification

Core claims checked against Apache Spark 3.5.7 documentation and source references on 2026-09-15. Example output observed in a local run: Ran on local PySpark 3.5.7 (local[2], JDK 17) on 2026-09-15; output above is observed. The lesson's other blocks ran on the same runtime; the SCD block also used Delta Lake 3.2.1, and the dropDuplicates demonstration also ran on a two-executor local-cluster master. Explanations, scenarios, and diagrams are original teaching material. Other releases and managed runtimes can differ.