PySpark: one problem shape at a time

Look for

• Keep some rows and some columns
• Add a column calculated from other columns
• Rename, drop or cast columns

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [(1, "paid", 120.0), (2, "paid", 80.0), (3, "refunded", 200.0)],
    ["order_id", "status", "amount"],
)
result = (
    df.filter((F.col("status") == "paid") & (F.col("amount") > 100))
    .select("order_id", "amount")
    .withColumn("amount_with_tax", F.round(F.col("amount") * 1.18, 2))
)

The usual trap

Each condition needs its own parentheses. F.col('status') == 'paid' & F.col('amount') > 100 does not mean what it reads like, because & binds tighter than ==. It raises an error instead of filtering. Wrap each condition in its own parentheses.

Not this

If the answer has 1 row per group, it needs groupBy and agg. filter and select keep the rows as they are, only fewer or narrower.

Look for

• Label values as Low, Medium or High
• Build 1 column from other columns
• A rule that a built-in function can already express

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame([(1, 45.0), (2, 120.0), (3, None)], ["id", "amount"])
result = df.withColumn(
    "band",
    F.when(F.col("amount").isNull(), "Unknown")
    .when(F.col("amount") < 50, "Low")
    .otherwise("High"),
)

The usual trap

A NULL is not less than 50, so it falls through to otherwise. Without the isNull branch first, the row with a missing amount is labelled High: 3 | NULL | High. Handle the missing case first.

Not this

Reach for a UDF (a Python function you register to run on each row) only when no built-in can express the rule. A UDF runs outside Spark's optimiser, so try when and otherwise first.

Look for

• A column that holds a list
• 1 row per tag, item or word
• Rows to columns, or columns to rows

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame([(1, ["a", "b"]), (2, []), (3, None)], ["id", "tags"])
result = df.select("id", F.explode_outer("tags").alias("tag"))

The usual trap

explode drops rows whose list is empty or NULL. With ids 1, 2 and 3, explode returns only id 1 (2 rows). explode_outer keeps ids 2 and 3 with a NULL tag. Choose on purpose.

Not this

groupBy and collect_list does the reverse: many rows become 1 row with a list. explode turns 1 row with a list into many rows.

Look for

• 1 row per user, product or day
• Total, average, count, maximum
• Only the groups that pass a test

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [("u1", "A", 10.0), ("u1", "B", None), ("u2", "A", 5.0)],
    ["user", "product", "amount"],
)
result = (
    df.groupBy("user")
    .agg(F.count("*").alias("rows"), F.count("amount").alias("amounts"), F.sum("amount").alias("total"))
    .orderBy("user")
)

The usual trap

count('*') counts rows and count('amount') skips NULLs. For u1 the result is rows 2, amounts 1. If a prompt says number of purchases, decide which of the two it means.

Not this

A window (covered in the window step) keeps every input row and adds a value next to it. groupBy returns only 1 row per group.

Look for

• Add columns from another table
• Keep orders even when the customer is unknown
• One side is small and one is large

Smallest shape

from pyspark.sql import functions as F

orders = spark.createDataFrame([(1, "u1"), (2, "u2"), (3, "u9")], ["order_id", "user"])
users = spark.createDataFrame([("u1", "Asha"), ("u2", "Ben")], ["user", "name"])
# broadcast sends the small table to every worker
result = orders.join(F.broadcast(users), "user", "left").orderBy("order_id")

The usual trap

The default join is inner, and it drops rows with no match. Joining these orders to users without a 'left' keeps 2 rows, and order 3 silently disappears. With 'left' all 3 stay, and order 3 gets a NULL name.

Not this

If you only need the rows that have no match, a left_anti join returns just those rows and no columns from the other side.

Look for

• A running total per user
• Rank or top N inside each group
• Compare a row with the one before it

Smallest shape

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

df = spark.createDataFrame(
    [("u1", 1, 10.0), ("u1", 2, 5.0), ("u2", 1, 7.0)],
    ["user", "day", "amount"],
)
w = Window.partitionBy("user").orderBy("day")
result = df.withColumn("running", F.sum("amount").over(w)).orderBy("user", "day")

The usual trap

A window without orderBy adds up the whole partition on every row. With Window.partitionBy('user') alone, u1's 2 rows both show 15, instead of 10 and then 15. Add orderBy for a running total.

Not this

groupBy gives 1 row per group. A window keeps all rows and adds a column, so use it when you still need the detail.

Look for

• Start a new session after a pause
• Longest unbroken run
• Reset a counter when something changes

Smallest shape

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

df = spark.createDataFrame([("u1", 100), ("u1", 160), ("u1", 4000), ("u2", 50)], ["user", "ts"])
w = Window.partitionBy("user").orderBy("ts")
flagged = df.withColumn("prev", F.lag("ts").over(w)).withColumn(
    "new_session",
    F.when(F.col("prev").isNull() | (F.col("ts") - F.col("prev") > 1800), 1).otherwise(0),  # 1800 seconds = 30 minutes
)
result = (
    flagged.withColumn("session_no", F.sum("new_session").over(w))
    .select("user", "ts", "session_no")
    .orderBy("user", "ts")
)

The usual trap

This is the same idea as the SQL sessionisation card. The first event of a user has no previous event. Without the isNull check, that row is not flagged as a start, and the sessions are numbered from 0: u1's first 2 rows get session_no 0 and the third gets 1.

Not this

A plain lag compares 1 row with the previous one and stops there. Sessions and streaks turn those comparisons into group numbers with a running sum.

Look for

• The same key appears more than once
• Keep the newest or the earliest row
• Remove duplicates but not the right one by accident

Smallest shape

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

df = spark.createDataFrame(
    [("e1", 1, "old"), ("e1", 3, "new"), ("e2", 2, "only")],
    ["event_id", "ingest", "payload"],
)
w = Window.partitionBy("event_id").orderBy(F.col("ingest").desc())
result = df.withColumn("rn", F.row_number().over(w)).filter(F.col("rn") == 1).drop("rn").orderBy("event_id")

The usual trap

dropDuplicates(['event_id']) keeps 1 row per key but does not let you say which. In our check it kept the old row (ingest 1, payload old), and Spark does not promise which copy survives. Rank the rows with row_number when the choice matters.

Not this

dropDuplicates is fine when the duplicates are exact copies. When copies differ and one is better, you need an order.

Look for

• How many values are missing in each column
• Does the data match the expected types
• Mask or check columns before sharing

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [(1, "[email protected]", None), (2, None, "9990001111"), (3, "[email protected]", None)],
    ["id", "email", "phone"],
)
# one count of NULLs per column
result = df.select([F.sum(F.col(c).isNull().cast("int")).alias(c) for c in df.columns])

The usual trap

A check that finds nothing can mean the data is clean or that the check never looked. Here the profile shows id 0, email 1, phone 2. Run a check on data you know is bad at least once, to see it fail.

Not this

Fixing the data (fill, drop, cast) changes it. A data quality check only measures, and tells you whether to fix.

Look for

• Too many small output files
• The same data is reused by 2 queries
• Spark is slow after a join or a wide step

Smallest shape

df = spark.createDataFrame([(i,) for i in range(6)], ["id"])
result = df.coalesce(2).orderBy("id")

The usual trap

coalesce can only lower the number of partitions, and it avoids a full shuffle. repartition can raise it and shuffles all the data. This is how Spark behaves in a cluster; the editor here runs a small simulation, so it does not show the difference.

Not this

Caching helps when you reuse the same frame. It does not fix a bad partition count or a skewed key. Read the plan first, then decide.

Look for

• One task runs much longer than the rest
• A few keys hold most of the rows
• A join or groupBy on a popular value

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame([("hot", 1)] * 5 + [("cold", 1)], ["k", "v"])
result = df.groupBy("k").count().orderBy(F.col("count").desc())

The usual trap

Averages hide skew. 6 rows over 2 keys average 3 per key, which looks fine, but the counts show hot with 5 and cold with 1, so one key holds 5 of the 6 rows. Look at rows per key before you change anything.

Not this

A slow job is not always skew. If every key is similar in size, look at partitions, caching and the plan instead.

Look for

• The driver ran out of memory
• collect, toPandas or a loop over rows
• Fine on a sample, fails on the full data

Smallest shape

df = spark.createDataFrame([(1, "a"), (2, "a"), (3, "b")], ["id", "v"])
summary = df.groupBy("v").count()
small = summary.collect()
result = summary.orderBy("v")

The usual trap

collect() brings every row to one machine. Collecting a summary of a few rows is fine; collecting the whole table is the usual cause of a driver out-of-memory error. Do the grouping in Spark, then collect only the small result. The editor cannot reproduce a real memory limit.

Not this

An out-of-memory error on an executor (a worker that processes the data) often points to a skewed or very large partition. A driver problem (the driver is the program that coordinates the job) usually comes from pulling data back, not from processing it.

Look for

• Events arrive late or out of order
• How long to wait before a result is final
• Drop events that are too old

Smallest shape

from pyspark.sql import functions as F

df = spark.createDataFrame([("e1", 100), ("e2", 130), ("e3", 40)], ["id", "event_time"])
newest = df.agg(F.max("event_time").alias("m")).collect()[0]["m"]
result = df.filter(F.col("event_time") >= newest - 60).orderBy("id")

The usual trap

A watermark is the rule for how long to wait for late events, and this skeleton imitates it with a batch filter. The cut-off is measured from the newest event time, not from the clock. Here the newest event is 130, so the limit is 70 and e3 (40) is dropped. Measured from the wall clock, a replay of old data would be thrown away entirely.

Not this

Ordinary batch filtering has all the data at once. A watermark is for data that is still arriving, where you must decide when to stop waiting.

Look for

• Update existing keys and add new ones
• Keep the old version of a row for history (type 2)
• Apply a batch of changes to a target

Smallest shape

target = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "val"])
updates = spark.createDataFrame([(2, "B"), (3, "c")], ["id", "val"])
result = target.join(updates, "id", "left_anti").unionByName(updates).orderBy("id")

The usual trap

Union alone duplicates the changed key. Without the left_anti step, key 2 appears twice: b and B. Remove the old rows for the keys you are replacing, then add the new ones.

Not this

This skeleton is a type 1 update: it overwrites the old value. A type 2 change, which the anchor problem asks for, keeps the old row, closes it with an end date, and adds a new current row.

Look for

• Raw data to cleaned data to a summary
• Bronze (raw), silver (cleaned) and gold (summary) layers
• Add a new batch to what already exists

Smallest shape

from pyspark.sql import functions as F

raw = spark.createDataFrame([("r1", "12.5"), ("r2", "bad"), ("r1", "12.5")], ["id", "reading"])
result = (
    raw.withColumn("reading", F.col("reading").cast("double"))
    .filter(F.col("reading").isNotNull())
    .dropDuplicates(["id"])
)

The usual trap

This is one silver step. Cleaning that drops bad rows hides them. The row r2 ('bad') vanishes from the result with no trace. Keep the rejected rows in a separate table, so someone can count and fix them.

Not this

A single transformation is one step. A medallion pipeline is the same idea repeated in layers, where each layer has a clear job and can be rerun.