DEV Community

Cover image for PySpark Coding Interview: 15 DataFrame Problems With Full Solutions
Gowtham Potureddi
Gowtham Potureddi

Posted on

PySpark Coding Interview: 15 DataFrame Problems With Full Solutions

Most pyspark interview questions are not trivia about which method returns a DataFrame and which returns a list — they are a laptop, a small sample table, and a sentence that starts with "transform this so that…". The interviewer is watching one thing: whether you think in the DataFrame API the way a Spark engineer does, chaining transformations that describe what you want and letting a single action trigger the how. Reach for a Python loop, collect the whole frame to the driver, or write a UDF where a built-in exists, and you have answered the real question — "does this person understand distributed computation" — with a quiet no, no matter how correct the output looks.

This guide is fifteen real DataFrame problems, the kind that actually show up in a pyspark coding interview, grouped into the five domains an interviewer walks through: the DataFrame mental model, selection and joins, groupBy and aggregations, window functions, and the dedup/UDF/performance questions that separate a junior answer from a senior one. Every problem gives you a sample pyspark dataframe, runnable PySpark you could paste into a spark-submit, a line-by-line trace of what Spark does, the exact output table, and a one-line rule of thumb. The headline problem in each domain also gets the full treatment — the solution code, a step-by-step trace, the output, and a concept-by-concept breakdown of why it is the idiomatic Spark answer and not just a working one.

PipeCode blog header for the PySpark coding interview — bold white headline 'PySpark Coding Interview' over a hero composition of four DataFrame-domain glyph medallions (filter/join, groupBy, window, performance) arranged on a wheel around a central purple Spark-DataFrame seal, on a dark gradient.

When you want hands-on reps alongside the reading, drill the DataFrame transformation library →, rehearse joins on the join practice library →, and sharpen ranking with the window-function practice library →.


On this page


1. The PySpark coding interview format + the DataFrame mental model

The interview measures whether you think in transformations and actions, not whether you memorised the API

The one-sentence framing that changes how you prepare: a PySpark coding interview is a small sample DataFrame plus a transformation request, and the interviewer is scoring whether you build a lazy chain of transformations that reads like the intent and let exactly one action materialise the result — not whether you can recite method names. Almost every weak answer fails on the same axis: it drags data to the driver (collect() then a Python loop), or it rebuilds in Python what a Spark built-in already does in the engine. Get the mental model right and the syntax follows; get it wrong and no amount of API recall saves you.

The mental model you must be able to state out loud.

  • Transformations are lazy. select, filter, withColumn, join, groupBy().agg(), orderBy — none of these run anything. They append a node to a logical plan (a DAG) and return a new DataFrame. Chaining a hundred of them costs nothing until you ask for a result.
  • Actions trigger a job. count, collect, show, take, write, toPandas — these are the only calls that submit work to the cluster. When someone asks "what actually runs the code?", the answer is "an action."
  • The Catalyst optimizer rewrites your plan. Because transformations are lazy, Spark sees the whole chain before executing and can push filters down, prune columns, and reorder joins. That is why declarative DataFrame code often beats hand-tuned RDD code — Catalyst optimises what you described.
  • DataFrames are immutable. Every transformation returns a new DataFrame; you never mutate in place. df.withColumn(...) does nothing unless you assign it.

How to actually talk through an answer. The candidates who pass narrate a small, repeatable routine instead of typing silently:

  • Restate the transformation and the grain of the output — "so you want one row per customer, keeping only customers with three or more orders" — which pins down groupBy-vs-window before you write a line.
  • Name the pattern out loud — "this is a top-N-per-group, so a row_number window" — so the interviewer sees the mapping even if you fumble a method name.
  • Write the lazy chain, then point at the action — "these are all transformations; show() is the only thing that runs."
  • Pre-empt the follow-up — mention the shuffle, the null edge case, or the broadcast opportunity before you are asked; it converts a coding screen into a design conversation.

How PySpark interview questions are written — the tell. Learn to hear the pattern behind the prompt because it maps straight to an API.

  • "Keep only rows where…" → filter / where with a column expression.
  • "Add a column that…" → withColumn + when/otherwise or a built-in.
  • "One row per group with…" → groupBy().agg() (aggregate) or a Window + row_number (pick a row).
  • "The Nth / latest / rank within each group" → a window function, never a self-join.
  • "Rows in A but not in B" → a left_anti join.
  • "It's slow / it's skewed / a UDF is the bottleneck" → broadcast, salt, or replace the UDF with a native function.

The vocabulary you should use unprompted. Interviewers relax the moment you speak the execution model fluently, because it signals you have actually run Spark at scale rather than just read about it.

  • SparkSession is the entry point — spark.read... gives you a DataFrame; spark.createDataFrame(...) builds one from local data for a quick demo.
  • DataFrame over RDD over Dataset (in PySpark). The DataFrame API is the default; RDDs are the low-level escape hatch for non-tabular or fine-grained control; the typed Dataset API is a Scala/Java concept, so in Python you say "DataFrame."
  • Jobs → stages → tasks. An action launches a job; a job splits into stages at each shuffle boundary; a stage runs as parallel tasks, one per partition. When you say "this shuffle creates a new stage," you sound like someone who has read a Spark UI.
  • Narrow vs wide transformations. select/filter/withColumn are narrow (no data movement); join/groupBy/distinct/orderBy are wide (a shuffle/Exchange). Wide transformations are where cost lives.
  • Adaptive Query Execution (AQE). Modern Spark re-optimises the plan at runtime — coalescing shuffle partitions, switching join strategies, and splitting skewed partitions. Mentioning AQE on a performance question is a strong senior signal.

What separates a strong answer.

  • Can you say which line is the action in your own solution? — the single fastest credibility signal.
  • Do you default to built-in pyspark.sql.functions before writing a UDF? — required.
  • When output has one row per key, can you explain why you used groupBy vs a Window? — the senior distinction.
  • Can you name the shuffle boundaries in your chain and roughly how many stages the job has? — the reliability signal.

Warm-up 1 — select, filter, and the action that runs it

Detailed explanation. The gentlest opener asks you to narrow a DataFrame to a subset of columns and rows and report a number. The trap is answering "how many active users?" by collecting the frame and counting in Python. The idiomatic path is a select + filter chain (both lazy) terminated by a single count (the action). Being able to point at count() and say "this is the only line that runs a job" is the whole point of the warm-up.

  • select narrows columns — lazy, returns a new DataFrame.
  • filter/where narrows rows with a boolean column expression — lazy.
  • count is the action that submits the job and returns an int.

Question. Given a users DataFrame, how many users are active and aged 18 or older, using only DataFrame operations?

Input.

user_id age status
1 25 active
2 17 active
3 40 inactive
4 33 active

Code.

from pyspark.sql import functions as F

adults_active = (
    users
    .select("user_id", "age", "status")     # transformation (lazy)
    .filter((F.col("age") >= 18) & (F.col("status") == "active"))  # transformation (lazy)
)

result = adults_active.count()               # ACTION — this is what runs the job
print(result)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. select appends a projection node to the logical plan; nothing executes.
  2. filter appends a predicate node; Catalyst will later push this filter down toward the scan.
  3. Up to here adults_active is just a plan — no data has moved.
  4. count() is the action: Spark builds a physical plan, reads the rows, applies the pushed-down filter, and returns the count.

Output:

result
2

Rule of thumb. Build the chain with select/filter (lazy), then name the one action that runs it — if your answer has no clear action, you are not thinking in Spark yet.

Warm-up 2 — derive a categorical column with when/otherwise

Detailed explanation. The next warm-up asks you to add a derived column — a spend tier, a fraud flag, a bucket. The instinct to avoid is a UDF or a .collect() loop; the idiomatic tool is withColumn plus the vectorised F.when(...).otherwise(...) ladder, which stays inside Catalyst and runs distributed. Chained when clauses read top-to-bottom like an if/elif/else, and the trailing otherwise is the default bucket.

  • withColumn adds or replaces a column and returns a new DataFrame.
  • F.when(cond, value) is the vectorised conditional; chain .when(...) for more branches.
  • .otherwise(default) supplies the fallback; without it unmatched rows get null.

Question. Add a tier column: spend ≥ 1000 → "gold", ≥ 500 → "silver", else "bronze".

Input.

customer_id spend
c1 1200
c2 700
c3 150

Code.

from pyspark.sql import functions as F

tiered = customers.withColumn(
    "tier",
    F.when(F.col("spend") >= 1000, "gold")
     .when(F.col("spend") >= 500, "silver")
     .otherwise("bronze"),
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. withColumn("tier", ...) appends a projection that computes tier per row — still lazy.
  2. The when ladder is evaluated top-down per row: the first matching predicate wins, so ordering the thresholds high-to-low matters.
  3. c1 (1200) matches the first branch → "gold"; c2 (700) skips branch one, matches branch two → "silver"; c3 (150) falls through to otherwise → "bronze".
  4. Nothing executes until a later action (show, write) materialises the frame.

Output:

customer_id spend tier
c1 1200 gold
c2 700 silver
c3 150 bronze

Rule of thumb. Derive columns with withColumn + when/otherwise, ordering the branches most-specific first; never reach for a UDF for logic a when ladder expresses.

Warm-up 3 — distinct, dropDuplicates, and reading the plan with explain()

Detailed explanation. The third warm-up quietly probes whether you know the difference between distinct() (dedupe on all columns) and dropDuplicates(subset) (dedupe on a subset, keeping an arbitrary surviving row), and whether you can read a query plan. Interviewers love asking "what does explain() show?" because it exposes whether you understand that your lazy chain became an optimised physical plan. You are not expected to memorise every operator — just to recognise the scan, the exchange (shuffle), and the aggregate that a distinct implies.

  • distinct() removes fully-duplicate rows (all columns must match).
  • dropDuplicates(["country"]) keeps one row per distinct country (which row is arbitrary without ordering).
  • explain() prints the physical plan — a distinct shows up as a hash aggregate with an Exchange (shuffle) because deduping needs data grouped by key across the cluster.

Question. How many distinct countries appear in users, and what does the plan reveal about the cost?

Input.

user_id country
1 US
2 IN
3 US
4 DE

Code.

distinct_countries = users.select("country").distinct()
distinct_countries.explain()          # inspect the physical plan
n = distinct_countries.count()        # ACTION
print(n)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. select("country") prunes to one column; Catalyst pushes this projection to the scan so only country is read.
  2. distinct() compiles to a hash aggregate keyed by country, which requires an Exchange — a shuffle that moves matching keys to the same executor.
  3. explain() prints roughly HashAggregate <- Exchange hashpartitioning(country) <- HashAggregate <- Scan, revealing the shuffle is the real cost.
  4. count() fires the job; US collapses to one row, leaving US, IN, DE → 3.

Output:

n
3

Rule of thumb. Use distinct() for whole-row dedupe and dropDuplicates(subset) for per-key dedupe; when asked "is this expensive?", point at the Exchange/shuffle in explain(), because that is where the money goes.


2. Selection, filtering & joins

Combining DataFrames is where interviews get real: pick the join type from the question, and broadcast the small side

Iconographic PySpark join diagram — a fact DataFrame and a dimension DataFrame flowing through a filter-funnel into a join node showing inner/left/anti join glyphs, a small broadcast badge on the dimension side, and the joined result grid on the right.

The invariant to burn in: filter as early as possible so less data reaches the join, choose the join type directly from the wording of the question (inner = matches only, left = keep all of the left, left-anti = left rows with no match), and broadcast any dimension small enough to fit in memory so Spark skips the expensive shuffle. Every join scenario is a variation on picking the right type and deciding whether the small side can be broadcast.

The join types and when each wins.

  • inner — keep only rows with a match on both sides. The default; use it when unmatched rows are irrelevant.
  • left (left outer) — keep every left row; unmatched right columns become null. Use it when the fact table must be preserved and you will coalesce the nulls.
  • right / full (outer) — keep the right side, or both sides. Rarer; full surfaces rows missing from either side.
  • left_semi — keep left rows that have a match, but return only left columns (a filtered existence check, no column widening).
  • left_anti — keep left rows that have no match. The idiomatic "in A but not in B" answer — no NOT IN subquery, no null-checking after an outer join.

Broadcast joins — the performance answer. A normal (shuffle/sort-merge) join moves both sides across the network to co-locate keys. If one side is small (a dimension, a lookup), wrap it in F.broadcast(dim) and Spark ships a copy to every executor, turning the join into a local map-side lookup with no shuffle. "Join a huge fact to a small dimension with minimal shuffle" → broadcast.

The three join strategies Spark can pick — and when. Interviewers love a follow-up here, so know the physical menu behind your logical join:

  • Sort-merge join — the default for two large tables: both sides are shuffled by the key and sorted, then merged. Robust but pays for a full shuffle on both sides.
  • Broadcast hash join — when one side is under the broadcast threshold (spark.sql.autoBroadcastJoinThreshold, ~10 MB by default), Spark broadcasts it and does a hash lookup on the big side — no shuffle. F.broadcast(df) forces this even above the threshold.
  • Shuffle hash join — a middle option Spark occasionally chooses when one side is smallish but not broadcastable; it builds a hash table per partition after a shuffle.
  • AQE can flip the strategy at runtime — if a stage's actual output turns out small, Adaptive Query Execution may convert a planned sort-merge join into a broadcast join automatically.

A note on join keys and correctness. Spark null-comparisons follow SQL: null == null is not true, so rows with a null key never match in an inner/left join. When a key can legitimately be null and you want nulls to match, use the null-safe equality eqNullSafe (<=>). And after any join, prefer selecting the exact columns you want so a repeated column name (like id from both sides) never becomes ambiguous downstream.

Column expressions — the grammar under every filter and join. Before the join types, be fluent in referencing columns, because interviewers notice sloppy references:

  • F.col("x") is the portable way to reference a column in an expression; df["x"] and df.x also work but tie you to one frame, which bites you after a join where both sides have id.
  • Arithmetic and comparison operators are overloadedF.col("price") * 1.1, F.col("age") >= 18 — and return new column expressions, not values.
  • Aliasing.alias("net") names a derived column; do it so the output schema reads cleanly.
  • Chained expressions stay lazy — a filter expression is itself just a plan fragment until an action fires.

Filtering that keeps the plan cheap.

  • F.col("x").isin(a, b, c) — membership test instead of chained ORs.
  • F.col("x").between(lo, hi) — inclusive range.
  • Null-safe logicisNull()/isNotNull(), and <=> (the null-safe equality) when a key can be null.
  • Filter before join — push predicates ahead of the join so fewer rows shuffle (Catalyst often does this for you, but writing it that way makes intent clear).

Common trap answers to pre-empt.

  • NOT IN / outer-join-then-null-check when the question is "rows with no match" — the clean answer is left_anti.
  • Forgetting to broadcast a tiny dimension and eating a full shuffle join.
  • Ambiguous column references after a join — alias the frames or select explicit columns so id is unambiguous.
  • Filtering after the join when the predicate only touches one side — filter first.

Multi-column filtering with isin, between & null handling — a worked teaching example

Detailed explanation. Before any join, interviewers check that you can express a compound predicate cleanly. Consider filtering orders to a set of regions, a price band, and non-null coupons in one readable expression. The idiomatic answer combines isin, between, and isNotNull with &/| — remembering that each comparison must be parenthesised because Python's operator precedence binds & tighter than >=.

  • isin(...) replaces region == 'US' | region == 'EU'.
  • between(lo, hi) is inclusive on both ends.
  • isNotNull() filters out missing values explicitly.
  • Parenthesise every clause(a) & (b), never a & b around comparisons.

Question. Keep orders where region is US or EU, price is between 50 and 500 inclusive, and coupon is not null.

Input.

order_id region price coupon
o1 US 120 SAVE10
o2 APAC 200 SAVE10
o3 EU 600 SAVE10
o4 US 75 null
o5 EU 300 SUMMER

Code.

from pyspark.sql import functions as F

filtered = orders.filter(
    (F.col("region").isin("US", "EU"))
    & (F.col("price").between(50, 500))
    & (F.col("coupon").isNotNull())
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. region.isin("US","EU") drops o2 (APAC).
  2. price.between(50, 500) drops o3 (600, above the band).
  3. coupon.isNotNull() drops o4 (null coupon).
  4. o1 and o5 satisfy all three predicates; the & chain is one pushed-down filter node.

Output:

order_id region price coupon
o1 US 120 SAVE10
o5 EU 300 SUMMER

Rule of thumb. Compose filters with isin/between/isNotNull and parenthesise every clause joined by &/|; it reads like the requirement and compiles to a single filter.

Left join a fact to a dimension with coalesce — a worked teaching example

Detailed explanation. The bread-and-butter join question keeps every fact row and enriches it from a dimension, defaulting the misses. You left join orders to customers on customer_id, then coalesce the possibly-null dimension column to a sensible default. The interviewer is watching that you (a) keep the left side with left, (b) resolve the ambiguous key by joining on a column name (not two aliased columns), and (c) fill nulls declaratively rather than after collect().

  • how="left" preserves all order rows.
  • Joining on the string "customer_id" collapses the duplicate key into one column (no ambiguity).
  • F.coalesce(col, default) returns the first non-null — the idiomatic null fill.

Question. Attach each order's customer name, defaulting unmatched customers to "UNKNOWN".

Input.

orders

order_id customer_id
o1 c1
o2 c2
o3 c9

customers

customer_id name
c1 Ana
c2 Ben

Code.

from pyspark.sql import functions as F

enriched = (
    orders.join(customers, on="customer_id", how="left")
          .withColumn("name", F.coalesce(F.col("name"), F.lit("UNKNOWN")))
          .select("order_id", "customer_id", "name")
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. The left join keeps o1, o2, o3; o3 (customer c9) has no dimension match, so name is null.
  2. Joining on the string "customer_id" means the result has a single customer_id column — no orders.customer_id vs customers.customer_id ambiguity.
  3. coalesce(name, "UNKNOWN") replaces the null for o3 with "UNKNOWN" and leaves matched names intact.
  4. select pins the exact output columns.

Output:

order_id customer_id name
o1 c1 Ana
o2 c2 Ben
o3 c9 UNKNOWN

Rule of thumb. Use how="left" to preserve the fact table, join on the shared column name to avoid ambiguity, and fill unmatched dimension columns with coalesce — not a post-collect Python patch.

Interview scenario on joins — customers with no orders

You are given a large customers DataFrame and an orders DataFrame. Return every customer who has never placed an order — a churn / re-engagement list. The interviewer explicitly says the customer table is large and the orders table is also large, and asks for the cleanest, most efficient DataFrame expression.

Solution Using a left-anti join

Answer choices (as the interviewer would probe them).

  • A. customers.join(orders, "customer_id", "left").filter(F.col("order_id").isNull()) then drop order columns.
  • B. Collect all order customer ids to the driver and filter customers with a Python not in.
  • C. customers.join(orders, "customer_id", "left_anti").
  • D. customers.subtract(orders.select("customer_id")) on mismatched schemas.

Code.

from pyspark.sql import functions as F

# Idiomatic: left-anti keeps left rows with NO match, returns only left columns
no_orders = customers.join(orders, on="customer_id", how="left_anti")
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. Keywords: "customers who have never placed an order" = left rows with no match in orders → the definition of left_anti.
  2. A works but is wasteful — it widens every row with order columns, then filters on a null, then drops the extra columns; more shuffle and more code.
  3. B collects order ids to the driver — it does not scale, defeats distribution, and OOMs on a large orders table — reject outright.
  4. D relies on subtract with identical schemas and whole-row equality, which is fragile and not what the question means.
  5. C expresses the intent exactly: left_anti returns only the customer columns for customers absent from orders, in one shuffle, no null gymnastics.

Output:

approach correct? why
left + null filter (A) works extra columns + extra work
driver not in (B) no doesn't scale, OOM risk
left_anti (C) yes intent-exact, one shuffle, left columns only
subtract (D) fragile needs identical schemas

Why this works — concept by concept:

  • left_anti semantics — an anti-join is defined as "left rows with no join partner," which is precisely "customers with no orders"; matching the API to the sentence is the whole skill.
  • No column widening — unlike a left join, anti-join returns only the left side's columns, so there is nothing to drop and no ambiguous order_id to null-check.
  • Stays distributed — the driver never sees the data; the join runs across executors, which is why B's collect is the anti-pattern the question is fishing for.
  • Cost — one shuffle join versus A's join-plus-filter-plus-project; on two large tables the anti-join is both the cleanest and the cheapest, and if orders' distinct keys were small you could even broadcast them.

PySpark
Topic — joins
DataFrame join and anti-join problems

Practice →

ETL Topic — data-transformation DataFrame selection and transformation problems

Practice →


3. GroupBy & aggregations

One row per group is a groupBy question: run every aggregate in a single pass and filter the aggregates for HAVING

Iconographic PySpark groupBy diagram — an input DataFrame partitioning into colour-grouped buckets that collapse into aggregate rows (sum, avg, countDistinct), with a pivot glyph turning rows into columns and a filter chip for the HAVING pattern.

The invariant: when the answer is "one row per key with some totals," it is a groupBy(keys).agg(...) — compute every aggregate you need in the single agg call, alias each one, then apply a post-agg filter for the HAVING-style condition and pivot when the question wants categories spread across columns. Rebuilding this with a loop, or aggregating one metric at a time and joining, is the answer to avoid.

The aggregation toolkit the exam expects.

  • groupBy(*keys).agg(...) — collapses rows sharing the keys into one row; put all aggregates in the one agg.
  • The aggregate functionsF.sum, F.avg, F.min, F.max, F.count, and F.countDistinct (distinct count of a column), each .alias("name")d so the output columns are readable.
  • count("*") vs countDistinct(col) — total rows per group vs distinct values per group; interviewers probe that you know the difference.
  • Post-agg filter — the DataFrame equivalent of SQL HAVING: aggregate first, then filter on the aggregated column.
  • pivotgroupBy(a).pivot("b").agg(...) turns distinct values of b into columns; supply the value list (pivot("month", months)) to avoid an extra scan.

Multiple aggregates in one pass — the efficient pattern. A common mistake is computing sum, then average, then distinct count as three separate groupBys and joining them. One agg with several expressions does it in a single shuffle:

df.groupBy("k").agg(F.sum("x").alias("sx"), F.avg("y").alias("ay"))
Enter fullscreen mode Exit fullscreen mode

The aggregate functions worth knowing beyond sum/avg. Interviewers reward a wider vocabulary because real reporting needs more than a total:

  • countDistinct vs approx_count_distinct — the exact distinct count shuffles all values; approx_count_distinct uses HyperLogLog for a fast, memory-cheap estimate. On a "distinct users over billions of rows, exact count not required" question, the approximate version is the senior answer.
  • collect_list / collect_set — gather a group's values into an array (collect_set de-duplicates). Useful for "list every product each customer bought," and they pair naturally with a later explode.
  • first / last (with ignorenulls=True) — grab a representative value per group; combine with an ordered window when "first" must be deterministic.
  • Conditional aggregationF.sum(F.when(cond, 1).otherwise(0)) counts rows matching a predicate per group, the DataFrame idiom for "count of paid orders and count of refunded orders in one pass."

Why groupBy is efficient. Spark computes a partial aggregate on each partition first (a map-side combine), then shuffles only the partial results and merges them. That is why groupBy().agg() scales: for a sum, each task sends one partial sum per key instead of every raw row. countDistinct is the exception — it must move the distinct values themselves — which is exactly why approx_count_distinct exists.

Common trap answers.

  • collect() then a Python defaultdict — never; that abandons the cluster.
  • One aggregate per groupBy, then join them back — one agg call computes them together in a single pass.
  • Filtering before aggregating when the condition is on the aggregate — "customers with ≥ 3 orders" must filter after the count, i.e. HAVING.
  • pivot without a value list on a high-cardinality column — it triggers a scan to discover values and can explode the column count.

Multi-aggregate groupBy in a single pass — a worked teaching example

Detailed explanation. The core groupBy question asks for several metrics per key at once — total spend, average order value, and number of distinct products. The idiomatic answer computes all three inside one agg so Spark shuffles the data once and emits one row per customer. Aliasing each aggregate keeps the output self-describing, which interviewers notice.

  • sum totals a numeric column per group.
  • avg averages it.
  • countDistinct counts unique values (here, distinct products).
  • One agg, one shuffle — all metrics computed together.

Question. Per customer_id, compute total revenue, average order value, and the number of distinct products purchased.

Input.

customer_id order_id product amount
c1 o1 A 100
c1 o2 B 60
c1 o3 A 40
c2 o4 C 200

Code.

from pyspark.sql import functions as F

summary = customers_orders.groupBy("customer_id").agg(
    F.sum("amount").alias("total_revenue"),
    F.avg("amount").alias("avg_order_value"),
    F.countDistinct("product").alias("distinct_products"),
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. groupBy("customer_id") marks customer_id as the shuffle key; rows for a customer land on the same executor.
  2. Inside one agg, Spark computes sum, avg, and countDistinct for each group in the same pass — one shuffle, not three.
  3. For c1: total = 100+60+40 = 200; avg = 200/3 ≈ 66.67; distinct products = {A, B} = 2.
  4. For c2: total = 200; avg = 200; distinct products = {C} = 1.

Output:

customer_id total_revenue avg_order_value distinct_products
c1 200 66.67 2
c2 200 200.0 1

Rule of thumb. Put every metric in one groupBy().agg() with aliases; computing metrics in separate passes and joining them is more shuffle for the same answer.

Pivot revenue by category across months — a worked teaching example

Detailed explanation. When the question wants a matrix — categories down the side, months across the top — it is a pivot. You groupBy the row key, pivot the column key (passing the explicit value list to skip a discovery scan), and agg the measure. The result is a wide DataFrame with one column per month, which is exactly how a "monthly revenue by category" report is asked for.

  • groupBy(row_key) sets the rows.
  • .pivot(col_key, [values]) sets the columns; the value list avoids an extra pass.
  • .agg(F.sum(measure)) fills each cell.

Question. Build a category-by-month revenue matrix for months ["Jan", "Feb"].

Input.

category month revenue
Books Jan 300
Books Feb 500
Toys Jan 200
Toys Feb 100

Code.

from pyspark.sql import functions as F

pivoted = (
    sales.groupBy("category")
         .pivot("month", ["Jan", "Feb"])   # explicit values -> no discovery scan
         .agg(F.sum("revenue"))
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. groupBy("category") makes one output row per category.
  2. pivot("month", ["Jan","Feb"]) reserves two output columns, Jan and Feb; passing the list means Spark does not scan the data first to find the month values.
  3. agg(F.sum("revenue")) fills each cell with the summed revenue for that (category, month).
  4. Books → Jan 300, Feb 500; Toys → Jan 200, Feb 100.

Output:

category Jan Feb
Books 300 500
Toys 200 100

Rule of thumb. Reach for pivot when the report wants values-as-columns, and always pass the value list for a known, bounded set of columns to avoid the discovery scan.

A frequent follow-up is "how would you undo a pivot?" — the inverse is an unpivot, which in PySpark you express with F.stack inside selectExpr (or by explodeing an array of structs), turning the month columns back into (category, month, revenue) rows. Knowing both directions signals you understand that pivot is just a reshape, not a computation, and that the cell values come from the aggregate you supplied.

Interview scenario on aggregation — high-value repeat customers

Given an orders DataFrame (customer_id, order_id, amount), return each customer's total spend and order count, but only for customers who placed at least 3 orders. This is the classic aggregate-then-filter (HAVING) question, and the interviewer wants to see you filter on the aggregate, not on raw rows.

Solution Using groupBy + agg + a post-aggregation filter

Answer choices.

  • A. Filter raw rows somehow, then group — but there is no row-level predicate for "≥ 3 orders."
  • B. groupByagg(count, sum)filter(count >= 3).
  • C. Compute counts, collect them, filter in Python, then re-join.
  • D. A Window count over the whole partition, then filter and dedupe.

Code.

from pyspark.sql import functions as F

repeat_customers = (
    orders.groupBy("customer_id")
          .agg(
              F.count("order_id").alias("order_count"),
              F.sum("amount").alias("total_spend"),
          )
          .filter(F.col("order_count") >= 3)     # HAVING: filter the aggregate
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. The condition "≥ 3 orders" is on an aggregate (a per-customer count), so it cannot be a row-level filter — that eliminates A.
  2. groupBy("customer_id").agg(...) collapses to one row per customer with order_count and total_spend in a single pass.
  3. .filter(order_count >= 3) is the HAVING step — it runs after aggregation on the computed column.
  4. C's collect-and-rejoin abandons distribution for no benefit; D's window-then-dedupe computes the count without collapsing rows, so it needs an extra dedupe — more work than a groupBy for a pure aggregate. B is the direct answer.

Output:

customer_id order_count total_spend
c1 5 820
c7 3 410

Why this works — concept by concept:

  • Aggregate then filter is HAVING — SQL's HAVING is exactly a filter applied after groupBy().agg(); recognising "≥ 3 orders" as a condition on the count is the crux.
  • One pass for all metrics — count and sum are computed together in the same shuffle, so the "also return total spend" requirement costs nothing extra.
  • groupBy beats a window here — because the output is genuinely one row per key, groupBy collapses directly; a window would keep all rows and force a dedupe.
  • Cost — a single shuffle by customer_id plus a cheap post-filter; no driver-side collection, so it scales to any number of customers.

PySpark
Topic — aggregation
GroupBy and aggregation problems

Practice →

Course Course — PySpark fundamentals PySpark fundamentals for data engineering interviews

Practice →


4. Window functions & ranking

The moment a question says "per group" and "keep all rows," it is a window — partition, order, then rank or offset

Iconographic PySpark window-function diagram — a DataFrame partitioned by category and ordered by sales, a sliding window frame highlighting rows, row_number/rank/dense_rank chips assigning ranks, and a lag arrow computing day-over-day change with a running-total band.

The invariant: a window function computes a value per row relative to a group of related rows without collapsing them — you define a Window.partitionBy(group).orderBy(sort) spec once, then apply row_number/rank/dense_rank for ranking, lag/lead for row-to-row deltas, or a sum over a rowsBetween frame for running totals. The tell for a window (versus a groupBy) is "keep every row but add a ranked/relative column," or "the top-N per group."

The window vocabulary.

  • Window.partitionBy(cols) — the grouping; each partition is ranked/scanned independently.
  • .orderBy(cols) — the order within a partition (add .desc() for descending).
  • row_number() — a strict 1,2,3 with no ties (ideal for top-N-per-group and dedupe).
  • rank() — ties share a rank and skip the next (1,1,3); dense_rank() — ties share and do not skip (1,1,2).
  • lag(col, n) / lead(col, n) — the value n rows back / ahead within the partition (day-over-day change, gaps).
  • rowsBetween(start, end) — the frame for running aggregates, e.g. Window.unboundedPrecedingcurrentRow for a running total.

Frame types — rowsBetween vs rangeBetween. The frame decides which rows the window function sees, and mixing the two up is a classic bug:

  • rowsBetween(start, end) counts physical rows — "the 3 rows before this one" — regardless of their values. Use it for a fixed-size sliding window and for running totals (unboundedPrecedingcurrentRow).
  • rangeBetween(start, end) counts by the value of the orderBy column — "all rows within 7 days of this row" — so ties and gaps in the ordering key change the frame. Use it for value-based windows like a trailing-7-day sum on a date column.
  • Default frames differ — a window with an orderBy but no explicit frame defaults to rangeBetween(unboundedPreceding, currentRow) for aggregates, which occasionally surprises people whose ordering column has duplicates.

Ranking vs aggregation — the distinction interviewers test. A groupBy().agg() returns one row per group; a window returns every original row plus a new column. "Give me the top 3 products per category" needs every product row scored, then filtered to rank ≤ 3 — a window, not a groupBy.

Windows shuffle too. A window with partitionBy(k) shuffles the data by k just like a groupBy, then sorts within each partition for the orderBy. Two performance cautions worth voicing: a window with no partitionBy funnels all rows into a single partition (a serious bottleneck — avoid unbounded global windows on big data), and stacking several windows with the same spec lets Spark reuse one shuffle, so define the Window object once and reuse it.

Common trap answers.

  • A self-join to find "the latest per key" — a row_number window is cleaner and cheaper.
  • groupBy when the question says "keep all rows and add a rank" — that collapses rows you were told to keep.
  • rank vs dense_rank vs row_number confusion — pick row_number for top-N/dedupe (no ties), dense_rank when tied items should share a place without gaps.
  • Forgetting .orderBy in a running total — an unordered frame gives meaningless partial sums.

Rank products by sales within each category — a worked teaching example

Detailed explanation. The canonical window question ranks items inside a group. You build a spec partitioned by category and ordered by sales descending, then add a rank column. The interviewer may probe the difference between rank, dense_rank, and row_number on ties, so state it: row_number breaks ties arbitrarily, rank leaves gaps after ties, dense_rank does not.

  • SpecWindow.partitionBy("category").orderBy(F.col("sales").desc()).
  • rank() — assigns 1,1,3 on a tie (gap).
  • Apply with withColumn — the rank is a new column, every row preserved.

Question. Add a rank of each product by sales within its category (highest sales = rank 1).

Input.

category product sales
Books A 500
Books B 500
Books C 200
Toys D 300

Code.

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

w = Window.partitionBy("category").orderBy(F.col("sales").desc())

ranked = df.withColumn("rank", F.rank().over(w))
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. The spec partitions rows by category, so Books and Toys rank independently.
  2. Within Books, ordering by sales desc puts A (500) and B (500) tied at the top, then C (200).
  3. rank() gives A and B rank 1 each, then skips to rank 3 for C (the gap that distinguishes rank from dense_rank).
  4. Toys has one product D → rank 1. Every input row survives, now with a rank column.

Output:

category product sales rank
Books A 500 1
Books B 500 1
Books C 200 3
Toys D 300 1

Rule of thumb. Partition by the group, order by the metric, and pick the ranking function by tie behaviour: row_number (no ties), rank (ties + gaps), dense_rank (ties, no gaps).

Day-over-day change with lag and a running total — a worked teaching example

Detailed explanation. Time-series questions want each row compared to the previous one (a delta) or a cumulative sum. lag(col, 1) fetches the prior row's value within the ordered partition; subtracting gives day-over-day change. A running total uses a sum over a frame from unboundedPreceding to currentRow. Both require an orderBy — without it the "previous" row and the cumulative frame are undefined.

  • lag(amount, 1).over(w) — yesterday's value; null for the first row.
  • amount - lag(...) — the delta.
  • F.sum(amount).over(w.rowsBetween(Window.unboundedPreceding, Window.currentRow)) — running total.

Question. For a single account's daily revenue, add day-over-day change and a running total.

Input.

day revenue
2026-08-01 100
2026-08-02 150
2026-08-03 120

Code.

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

w = Window.orderBy("day")
w_run = w.rowsBetween(Window.unboundedPreceding, Window.currentRow)

out = (
    daily
    .withColumn("dod_change", F.col("revenue") - F.lag("revenue", 1).over(w))
    .withColumn("running_total", F.sum("revenue").over(w_run))
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. Window.orderBy("day") sequences the rows chronologically.
  2. lag("revenue", 1) returns the prior day's revenue: null, 100, 150; the delta is null, 50, −30.
  3. w_run frames from the first row to the current, so sum accumulates: 100, 250, 370.
  4. Both columns are added per row — the frame is preserved, nothing collapses.

Output:

day revenue dod_change running_total
2026-08-01 100 null 100
2026-08-02 150 50 250
2026-08-03 120 -30 370

Rule of thumb. Always orderBy before lag/lead or a running sum; the first lag row is null by design, so handle it (or coalesce) if downstream math needs a number.

Interview scenario on windows — top 3 products per category

Given a sales DataFrame (category, product, sales), return the top 3 products by sales within each category. This is the most-asked window question in a PySpark interview, and the interviewer is checking that you use row_number (not a self-join, not a groupBy) and filter on the rank.

Solution Using Window row_number filtered to rank ≤ 3

Answer choices.

  • A. groupBy("category").agg(F.max("sales")) — only returns the single top, not top 3, and drops product identity.
  • B. Self-join sales to itself counting how many products score higher, keep count < 3.
  • C. Window.partitionBy("category").orderBy(sales.desc()) + row_number() + filter(rn <= 3).
  • D. Sort globally and limit(3) — ignores per-category grouping entirely.

Code.

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

w = Window.partitionBy("category").orderBy(F.col("sales").desc())

top3 = (
    sales
    .withColumn("rn", F.row_number().over(w))
    .filter(F.col("rn") <= 3)
    .drop("rn")
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. Keywords: "top 3" + "per category" + implicitly "keep the product rows" → a partitioned ranking window, not an aggregate.
  2. A collapses each category to one max value and loses which product it was — eliminate.
  3. B (self-join counting greater rows) works but is O(n²)-ish and needlessly heavy versus a single window pass — reject on cost.
  4. D takes the global top 3 across all categories, ignoring the per-category requirement — eliminate.
  5. C partitions by category, orders by sales desc, numbers rows 1..n with row_number (no ties to muddy the cut), keeps rn <= 3, and drops the helper column.

Output:

category product sales
Books A 500
Books B 480
Books C 300
Toys D 300
Toys E 250
Toys F 100

Why this works — concept by concept:

  • row_number for top-Nrow_number gives a strict 1..n with no ties, so rn <= 3 cleanly yields exactly three rows per group; rank/dense_rank could return more than three on ties.
  • partitionBy scopes the ranking — ranking resets per category, which is what "per category" demands; a global sort (D) cannot express that.
  • Window beats a self-join — one partitioned pass replaces an O(n²) self-join, the classic senior-vs-junior tell on this question.
  • Cost — a single shuffle to partition by category plus a sort within each partition; drop the helper column and you have the exact requested schema, no extra join.

PySpark
Topic — window-functions
Window function and ranking problems

Practice →

ETL Topic — data-transformation Running totals and time-series transformation problems

Practice →


5. Dedup, UDFs vs built-ins & performance/skew

The senior signal: dedupe with a window, prefer built-ins over UDFs, and know how to break a skewed join

Iconographic PySpark performance diagram — a dedup window keeping the latest row per key, a UDF box crossed out in favour of a native-function box, and a skewed join shown before and after salting with a broadcast badge on the small side.

The invariant: keep the latest record per key with a row_number window (not a blind dropDuplicates), always reach for a native pyspark.sql.functions before a Python UDF because UDFs are opaque to Catalyst and force per-row serialization, and when one key dominates a join, fix the skew with a broadcast join or by salting the hot key. These are the questions that decide the senior loop.

Deduplication done right.

  • dropDuplicates(subset) keeps an arbitrary surviving row per key — fine when rows are truly identical, wrong when you need "the latest."
  • row_number window — partition by the key, order by a timestamp descending, keep rn == 1 to deterministically keep the newest.
  • State which you need — "keep any one" (dropDuplicates) vs "keep the latest/best" (window) is the distinction interviewers probe.

UDFs vs built-ins — why it matters.

  • Native functions are vectorised and Catalyst-visible — Spark can optimise, push down, and codegen them.
  • Python UDFs are a black box — Spark serializes each row to Python, runs your function, serializes back; no pushdown, no codegen, big slowdown.
  • The reflex — express logic with when, regexp_replace, split, F.expr, coalesce, etc.; only write a UDF when no built-in (and no pandas UDF) can express it.
  • If you must — a pandas (vectorised) UDF is far faster than a plain Python UDF because it operates on Arrow batches, not row-by-row.

Why a Python UDF is so costly, concretely. For a plain F.udf, Spark cannot see inside the function, so it (1) serializes each row from the JVM to a Python worker process, (2) runs your Python one row at a time, and (3) serializes the result back — per row, across a process boundary. That kills three things at once: codegen (the whole-stage Java the DataFrame path compiles to), predicate pushdown (Catalyst will not push a filter through an opaque UDF), and CPU efficiency (row-at-a-time Python versus vectorised JVM). A pandas UDF recovers most of the loss by handing your function an Arrow-backed pandas.Series (a whole batch), so the round-trip amortises over thousands of rows — use it for genuinely custom logic like a bespoke parsing or a model inference call that no built-in covers. The decision ladder is simple: built-in first, pandas UDF if the logic is truly custom, plain Python UDF only as a last resort.

Performance & skew — the reliability questions.

  • Broadcast joinF.broadcast(small_df) avoids the shuffle when one side fits in memory.
  • Data skew — one key (a null, a mega-customer) sends most rows to one task; symptoms are one straggler task while the rest finish. Fix by salting (append a random suffix to the hot key on both sides so it spreads across tasks) or by isolating and broadcasting.
  • repartition vs coalescerepartition(n) shuffles to n even partitions (use before a heavy stage); coalesce(n) merges without a full shuffle (use to reduce output files).
  • Cache deliberatelycache()/persist() only a DataFrame reused across multiple actions; caching a once-used frame just wastes memory.
  • Let AQE help — Adaptive Query Execution (on by default in Spark 3.x) coalesces small shuffle partitions, converts eligible joins to broadcast at runtime, and has a skew-join feature that splits an oversized partition automatically; know it exists, but also know how to salt by hand for the cases AQE cannot cover.

The read-path optimisations interviewers expect you to mention. Half of "make it faster" is reading less in the first place:

  • Column pruningselect only the columns you need so a columnar format (Parquet/ORC) never reads the rest.
  • Predicate pushdown — a filter on a partitioned or Parquet-statistics column is pushed to the scan, skipping whole files/row-groups before any data is deserialized.
  • Partitioned storage — writing data partitioned by a date column lets a date filter prune whole directories, the file-layout twin of BigQuery partition pruning.
  • Watch for spill — when a task's data does not fit in memory it spills to disk (visible in the Spark UI); the fix is usually more/again-partitioned parallelism or fixing skew, not just more memory.

Common trap answers.

  • dropDuplicates for "keep the latest" — non-deterministic; use a row_number window ordered by timestamp.
  • A Python UDF for string cleanup that regexp_replace/when already does — the classic performance red flag.
  • Ignoring skew and blaming "Spark is slow" when one task is the straggler.
  • repartition when you meant coalesce (or vice-versa) — one shuffles, one does not.

Keep the latest record per key with a window — a worked teaching example

Detailed explanation. A CDC or event stream often has several versions of the same entity; you want the newest per key. dropDuplicates(["id"]) keeps an arbitrary row — wrong. The deterministic answer partitions by the key, orders by the update timestamp descending, numbers rows with row_number, and keeps rn == 1. This is the single most common dedup pattern in production PySpark.

  • Partition by the entity key (id).
  • Order by updated_at desc so the newest is first.
  • row_number() then filter(rn == 1) — deterministically the latest.

Question. From an events table with multiple updates per id, keep only the most recent row per id.

Input.

id updated_at status
u1 2026-08-01 new
u1 2026-08-03 paid
u2 2026-08-02 new

Code.

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

w = Window.partitionBy("id").orderBy(F.col("updated_at").desc())

latest = (
    events
    .withColumn("rn", F.row_number().over(w))
    .filter(F.col("rn") == 1)
    .drop("rn")
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. Partition by id so all versions of u1 land together.
  2. Order by updated_at desc puts u1's 2026-08-03 "paid" row first, its 2026-08-01 "new" row second.
  3. row_number assigns 1 to the newest row in each partition; u1 → the paid row gets rn = 1, u2 → its single row gets rn = 1.
  4. filter(rn == 1) keeps exactly one deterministic row per id; drop the helper column.

Output:

id updated_at status
u1 2026-08-03 paid
u2 2026-08-02 new

Rule of thumb. "Keep the latest per key" = row_number window ordered by timestamp desc, filter to rn == 1; dropDuplicates is only correct when any surviving duplicate is acceptable.

Replace a Python UDF with native functions — a worked teaching example

Detailed explanation. Interviewers plant a slow Python UDF and ask you to make it fast. The fix is almost always to express the same logic with native functions so it stays inside Catalyst and codegen. Here a UDF normalises a phone string; regexp_replace (strip non-digits) plus when (flag invalid lengths) does the same work vectorised, with no per-row Python round-trip.

  • The UDF serializes each row to Python — slow, opaque.
  • regexp_replace(col, pattern, repl) does the cleanup in-engine.
  • when(length(...)==10, ...).otherwise(...) flags validity without a UDF.

Question. Strip all non-digits from phone and flag whether the result has 10 digits — without a Python UDF.

Input.

id phone
1 (415) 555-2671
2 555-99

Code.

from pyspark.sql import functions as F

# BAD: a Python UDF — opaque to Catalyst, per-row serialization
# clean_udf = F.udf(lambda s: re.sub(r"\D", "", s))
# df.withColumn("digits", clean_udf("phone"))

# GOOD: native, vectorised, Catalyst-visible
cleaned = (
    df
    .withColumn("digits", F.regexp_replace("phone", r"\D", ""))
    .withColumn(
        "valid",
        F.when(F.length("digits") == 10, F.lit(True)).otherwise(F.lit(False)),
    )
)
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. regexp_replace("phone", r"\D", "") removes every non-digit in-engine: (415) 555-26714155552671, 555-9955599.
  2. length("digits") is a native column function; no Python round-trip per row.
  3. when(length == 10, True).otherwise(False) flags row 1 valid (10 digits) and row 2 invalid (5 digits).
  4. The whole chain codegens into a single stage — dramatically faster than the UDF at scale.

Output:

id digits valid
1 4155552671 true
2 55599 false

Rule of thumb. Before writing a UDF, ask "does a pyspark.sql.functions built-in already do this?" — for string, date, and conditional logic the answer is almost always yes, and native keeps Catalyst in play.

Interview scenario on performance — a skewed join that runs forever

You join a huge events DataFrame to a users dimension on user_id. The job crawls: 199 tasks finish fast and one runs for an hour. Investigation shows a sentinel user_id = "0" (unattributed events) accounts for 60% of the rows. Make the join fast. The interviewer wants the skew diagnosis and the fix.

Solution Using broadcast for the small side + salting the hot key

Answer choices.

  • A. repartition(1000) the events frame and hope the extra partitions help.
  • B. Increase executor memory so the one hot task doesn't spill.
  • C. Broadcast the small users dim to avoid the shuffle; if the hot key still skews (e.g. self-ish join), salt the hot user_id on both sides so it spreads across tasks.
  • D. Filter out user_id = "0" entirely and ignore those events.

Code.

from pyspark.sql import functions as F

# Case 1: users is small -> broadcast avoids the shuffle join entirely
joined = events.join(F.broadcast(users), on="user_id", how="left")

# Case 2: if both sides are large and one key is hot -> salt the hot key
SALT = 16
events_salted = events.withColumn(
    "salt", F.when(F.col("user_id") == "0", (F.rand() * SALT).cast("int")).otherwise(F.lit(0))
)
# explode the dim's hot key across the same salt range so keys still match
users_salted = users.withColumn(
    "salt",
    F.explode(F.when(F.col("user_id") == "0", F.array([F.lit(i) for i in range(SALT)]))
               .otherwise(F.array(F.lit(0)))),
)
joined_salted = events_salted.join(users_salted, on=["user_id", "salt"], how="left")
Enter fullscreen mode Exit fullscreen mode

Step-by-step trace.

  1. Diagnosis: 199 fast tasks + 1 hour-long task + one key at 60% of rows = classic data skew on user_id = "0".
  2. A (repartition) redistributes rows but the hot key still hashes to one partition — it does not fix skew — reject.
  3. B (more memory) treats the symptom (spill), not the cause (one task doing 60% of the work) — reject.
  4. D drops real data to dodge the problem — a correctness change the question did not authorise — reject.
  5. C fixes the cause: if users fits in memory, broadcast removes the shuffle so there is no skewed shuffle stage at all; if both sides are genuinely large, salting appends a 0..15 suffix to the hot key on the events side and replicates the matching dim rows across the same salts, so the hot key's rows spread across 16 tasks instead of one.

Output:

approach fixes skew? why
repartition (A) no hot key still hashes to one task
more memory (B) no treats spill, not the imbalance
drop key (D) no changes results / loses data
broadcast + salt (C) yes no shuffle, or hot key spread across tasks

Why this works — concept by concept:

  • Diagnose skew from the symptom — one straggler task while the rest finish is the fingerprint of a hot key; naming it is half the answer the interviewer wants.
  • Broadcast removes the shuffle — when the dimension fits in memory, shipping it to every executor turns the join into a local lookup, so there is no skewed shuffle stage to straggle.
  • Salting spreads the hot key — adding a random suffix on the big side and replicating the small side across the same suffixes splits the 60%-key's work across many tasks, eliminating the single straggler.
  • Cost — broadcast trades a little executor memory for zero shuffle; salting adds a small replication of the hot dim rows but converts an hour-long straggler into balanced parallel tasks — the right trade when the data is genuinely skewed.

PySpark
Topic — optimization
Spark performance, skew, and optimization problems

Practice →

Course
Course — Spark internals
Apache Spark internals for data engineering interviews

Practice →


Cheat sheet — PySpark DataFrame interview recipes

Transformation vs action (know which line runs the job).

Kind Examples Runs a job?
Transformation (lazy) select, filter, withColumn, join, groupBy().agg(), orderBy, distinct No — builds the DAG
Action (eager) count, collect, show, take, write, toPandas, first Yes — submits to the cluster

Join-type decision table.

The question says… Join
"only rows that match both" inner
"keep all of the left, null the misses" left
"rows in A but not in B" left_anti
"rows in A that have a match (A's columns only)" left_semi
"everything from both sides" full (outer)
"huge fact × small dim, minimal shuffle" F.broadcast(dim)

Window-spec recipes.

  • Top-N per group → partitionBy(group).orderBy(metric.desc()) + row_number() + filter(rn <= N).
  • Latest per key (dedup) → partitionBy(key).orderBy(ts.desc()) + row_number() + filter(rn == 1).
  • Day-over-day → orderBy(day) + col - lag(col, 1).
  • Running total → orderBy(day).rowsBetween(unboundedPreceding, currentRow) + sum.
  • Ties matter → row_number (no ties) vs rank (ties + gaps) vs dense_rank (ties, no gaps).

GroupBy vs window. Output is one row per keygroupBy().agg(). Output keeps every row and adds a ranked/relative column → Window. "≥ N per group" on a count → groupBy().agg(count).filter(count >= N) (HAVING).

Performance / skew checklist.

  • Small dimension → F.broadcast() to skip the shuffle.
  • One straggler task + a dominant key → salt the hot key (and replicate the small side across salts).
  • Prefer native pyspark.sql.functions over Python UDFs; if unavoidable, use a pandas UDF (Arrow-batched).
  • repartition(n) shuffles to even partitions before a heavy stage; coalesce(n) merges output files without a full shuffle.
  • cache()/persist() only DataFrames reused across multiple actions.
  • Never collect() a large frame to the driver — it defeats the point of Spark.

Pattern → API map (say the API out loud). "per group + keep rows" → window · "one row per key" → groupBy · "in A not in B" → left_anti · "latest per key" → row_number==1 · "slow UDF" → native/pandas UDF · "skewed join" → broadcast/salt.

Built-in functions worth having on the tip of your tongue. Reaching for the right pyspark.sql.functions (imported as F) is what keeps a UDF off the whiteboard.

Need Function
Conditional column F.when(cond, a).otherwise(b)
Null fallback F.coalesce(col, F.lit(default))
Clean a string F.regexp_replace, F.trim, F.lower, F.split
Extract / test a pattern F.regexp_extract, col.rlike(pattern)
Membership / range col.isin(...), col.between(lo, hi)
Flatten an array column F.explode(col) (or explode_outer to keep empties)
Group values into an array F.collect_list, F.collect_set
Distinct count (exact / approx) F.countDistinct, F.approx_count_distinct
Dates F.to_date, F.datediff, F.date_add, F.months_between
Raw SQL expression F.expr("...")

Common mistakes that quietly fail the round.

  • collect() on a big DataFrame to loop in Python — the single biggest anti-pattern; it pulls the whole frame to the driver and defeats Spark.
  • Un-parenthesised boolean predicatesF.col("a") == 1 & F.col("b") == 2 binds wrong; wrap each comparison: (F.col("a") == 1) & (F.col("b") == 2).
  • Assuming order without orderBy — Spark does not preserve input order across a shuffle; a "first row" is meaningless without an explicit order.
  • Forgetting DataFrames are immutabledf.withColumn(...) returns a new frame; if you do not assign it, nothing changed.
  • A UDF where a built-in exists — reread the built-ins table above before writing one.

Frequently asked questions

What kind of questions are asked in a PySpark coding interview?

Almost all of them hand you a small sample DataFrame and a transformation request — filter and join, aggregate per group, rank the top-N per group, deduplicate to the latest record, or fix a slow/skewed job. The pyspark interview questions that matter test whether you build a lazy chain of transformations and finish with a single action, and whether you reach for the right built-in (a window for top-N, a left_anti for "not in B") instead of collecting to the driver or writing a UDF. Expect at least one performance question about broadcast joins, UDFs, or data skew in a senior loop.

Is PySpark DataFrame or RDD tested more in interviews?

The pyspark dataframe API is what nearly every modern interview tests, because it is what teams actually write and because Catalyst optimises it. You should still be able to explain RDDs and when the lower-level API is justified (fine-grained control, non-tabular data), but you will almost never be asked to solve a problem in raw RDDs. Frame your answers in DataFrame/pyspark.sql.functions code and mention RDDs only if the interviewer steers there.

How do I answer "transformations vs actions"?

Say it in one breath: transformations (select, filter, join, groupBy().agg(), withColumn) are lazy — they only append to the logical plan and return a new DataFrame — while actions (count, collect, show, write) are the only calls that trigger a Spark job. Then point at your own code and name the action line. This is one of the highest-signal spark interview questions because it instantly reveals whether you understand the execution model.

Do I need to write UDFs in a PySpark interview?

Usually the opposite — interviewers plant a slow Python UDF and want you to replace it with native pyspark.sql.functions (regexp_replace, when, split, date functions) so the logic stays inside Catalyst and codegen. Writing a UDF where a built-in exists is a red flag, because a plain Python UDF serializes every row to Python and kills performance. If a UDF is genuinely unavoidable, reach for a pandas (vectorised) UDF, which processes Arrow batches instead of row-by-row.

How do window functions come up in Spark interviews?

Constantly — any question that says "per group" while keeping every row is a window. The most common is "top N products/users per category," solved with Window.partitionBy(group).orderBy(metric.desc()) and row_number() <= N; close behind are "latest record per key" (dedup with row_number == 1), day-over-day change with lag, and running totals with a rowsBetween frame. Being able to pick row_number vs rank vs dense_rank by their tie behaviour is a frequent follow-up.

How should I prepare for a PySpark coding interview in 2026?

Drill real DataFrame problems until the pattern-to-API mapping is automatic — "in A not in B" instantly becomes left_anti, "top-N per group" becomes a row_number window, "≥ N orders" becomes aggregate-then-filter. These pyspark interview questions for data engineer roles reward reps on the same five domains this guide covers, plus the ability to reason about shuffles, broadcast joins, and skew out loud. Build one small end-to-end job so you have seen a real shuffle and a real straggler task, because those experiences are what the performance pyspark coding questions are written around. Rehearse explaining each solution out loud — naming the action, the shuffle, and the trade-off — because the verbal walkthrough is scored as heavily as the code you type.


Practice on PipeCode

Turn PySpark patterns into muscle memory

Reading solutions is not the same as writing them under a clock. PipeCode drills build the reflex a PySpark interview actually tests — reading a prompt, mapping it to the right DataFrame API, and defending the join type, window, or skew fix out loud. Pipecode.ai is Leetcode for Data Engineering — scenario-first practice on transformations, joins, aggregations, and window functions tuned to the trade-offs real Spark interviews reward.

Practice DataFrame problems →
Practice window-function problems →

Top comments (0)