Seven PySpark coding problems with worked DataFrame code, the follow-up questions about scale, and the traps hidden in each one.
Many PySpark coding questions are ordinary data problems written in Spark syntax: keep the latest record per key, rank rows inside groups, find rows with no match, count words, build running totals, compare this month with last month. A common follow-up is what happens to that code when the table has a billion rows. To prepare, learn the seven problems below in the DataFrame API, know where each one shuffles data across the cluster, and be able to read the output of explain(). That is what this page covers, and nothing else.
For how this fits into the whole loop, including the data modeling and pipeline design rounds, see our guide to passing the data engineer interview.
The snippets use the DataFrame API as documented for PySpark 4.2; run them yourself on small data to see the output. The tables show the expected result on small sample DataFrames, and the plan output in problem 7 is trimmed. The examples use a few small DataFrames: records (customer profile updates), sales, customers, orders and countries. Assume this import at the top:
from pyspark.sql import Window, functions as FBefore you write a line, ask four questions
What is one row? What is the key? Can the key or the sort column be null? What should happen on a tie? Several of the problems below have a wrong answer that runs without errors and returns plausible rows. Asking these questions out loud is how you avoid writing it, and it shows the interviewer you think about data before syntax.
The seven problems
1. Keep only the latest record for each customer
What it tests: whether you reach for a window function or for dropDuplicates().
w = Window.partitionBy("customer_id").orderBy(
F.col("updated_at").desc(), F.col("record_id").desc()
)
latest = (
records
.withColumn("rn", F.row_number().over(w))
.filter(F.col("rn") == 1)
.drop("rn")
)+---------+-----------+-------+-------------------+
|record_id|customer_id| email| updated_at|
+---------+-----------+-------+-------------------+
| 3| 101|c@x.com|2026-03-01 09:00:00|
| 4| 102|d@x.com|2026-02-01 09:00:00|
+---------+-----------+-------+-------------------+dropDuplicates(["customer_id"]) looks shorter, but the API reference only promises to remove duplicate rows, "optionally only considering certain columns". It says nothing about which row survives, so it cannot answer "keep the latest". The second sort key matters too: in the test data, records 2 and 3 share the same updated_at. Without record_id as a tie-breaker, which one wins is not defined. Say that before the interviewer asks.
2. Top three products by revenue in each category
What it tests: aggregating before ranking, and knowing that the three ranking functions treat ties differently.
revenue = sales.groupBy("category", "product_id").agg(F.sum("amount").alias("revenue"))
w = Window.partitionBy("category").orderBy(F.col("revenue").desc())
top3 = revenue.withColumn("rank", F.dense_rank().over(w)).filter(F.col("rank") <= 3)+--------+----------+-------+----+
|category|product_id|revenue|rank|
+--------+----------+-------+----+
| books| x| 30.0| 1|
| books| y| 20.0| 2|
| toys| a| 55.0| 1|
| toys| b| 40.0| 2|
| toys| c| 40.0| 2|
| toys| d| 10.0| 3|
+--------+----------+-------+----+"Top three" returned four toys, because b and c tie. row_number() would return exactly three rows (a, b and c) and number the tied b and c as 2 and 3 in an arbitrary order, so d is dropped and the tie is hidden. rank() would give 1, 2, 2 and skip 3, so d would disappear. dense_rank() leaves no gaps. None of these is wrong; the right one depends on what the business means by "top three", which is exactly the clarifying question to ask.
3. Customers who have never placed an order
What it tests: whether you know the anti join.
no_orders = customers.join(orders, on="customer_id", how="left_anti")+-----------+----+
|customer_id|name|
+-----------+----+
| 103| Cy|
+-----------+----+A left join followed by filter(orders.order_id.isNull()) also works, but it drags the right-hand columns along and hides the intent. left_anti keeps left rows with no match; left_semi keeps left rows that do have a match. Both return only left-side columns.
4. Word count on a large text file
What it tests: split and explode, and whether you clean the input.
lines = spark.read.text(path) # one row per line, in a column named "value"
counts = (
lines
.select(F.explode(F.split(F.lower(F.col("value")), r"[^a-z']+")).alias("word"))
.filter(F.col("word") != "")
.groupBy("word")
.count()
.orderBy(F.col("count").desc())
)On the input "The cat sat." and "the CAT, the hat!" this returns the: 3, cat: 2, sat: 1, hat: 1. The RDD version most people memorise (flatMap(lambda l: l.lower().split()), then map to pairs, then reduceByKey) returned cat and cat, as two different words on the same input, and sat. with its full stop. Normalise punctuation, or say why you are not. If the interviewer wants the RDD version, use reduceByKey rather than groupByKey: the RDD guide notes that for per-key aggregations reduceByKey or aggregateByKey "will yield much better performance".
5. Running total and a seven-day rolling sum per customer
What it tests: window frames, and the difference between counting rows and counting time.
running = (
Window.partitionBy("customer_id").orderBy("order_date", "order_id")
.rowsBetween(Window.unboundedPreceding, Window.currentRow)
)
last_7_days = (
Window.partitionBy("customer_id")
.orderBy(F.unix_date("order_date")) # days since 1970-01-01
.rangeBetween(-6, 0)
)
result = (
orders
.withColumn("running_total", F.sum("amount").over(running))
.withColumn("spend_7d", F.sum("amount").over(last_7_days))
)+--------+-----------+------+----------+-------------+--------+
|order_id|customer_id|amount|order_date|running_total|spend_7d|
+--------+-----------+------+----------+-------------+--------+
| 1| 101| 20.0|2026-01-01| 20.0| 20.0|
| 2| 101| 30.0|2026-01-03| 50.0| 50.0|
| 3| 101| 10.0|2026-01-09| 60.0| 40.0|
| 4| 102| 15.0|2026-02-14| 15.0| 15.0|
| 5| 102| 25.0|2026-03-02| 40.0| 25.0|
+--------+-----------+------+----------+-------------+--------+On 9 January the seven-day sum is 40, not 60, because the 1 January order falls outside the frame. rowsBetween(-6, 0) would have meant "this order and the six before it", which is a different question when customers do not order every day. rangeBetween needs a numeric sort key, which is why the window orders by unix_date rather than the date itself.
6. Month-over-month revenue change
What it tests: lag, and a Spark 4 behaviour change that breaks old answers.
monthly = (
orders
.withColumn("month", F.trunc("order_date", "month"))
.groupBy("customer_id", "month")
.agg(F.sum("amount").alias("revenue"))
)
w = Window.partitionBy("customer_id").orderBy("month")
mom = (
monthly
.withColumn("prev_revenue", F.lag("revenue").over(w))
.withColumn("pct_change", F.round(
F.try_divide(F.col("revenue") - F.col("prev_revenue"), F.col("prev_revenue")) * 100, 1))
)Customer 102 goes from 15.0 to 25.0, a change of 66.7%. Each customer's first month has no previous value, so prev_revenue and pct_change are null. The trap is a previous month with zero revenue. Since Spark 4.0, ANSI mode is on by default, so a plain / with a zero divisor fails the whole job with a DIVIDE_BY_ZERO error. Older tutorials, written when the same division quietly returned null, do not warn you. try_divide returns null when the divisor is 0.
7. Enrich a large table with a small lookup table
What it tests: broadcast joins, and proving the plan does what you think.
enriched = orders.join(F.broadcast(countries), on="customer_id", how="left")
enriched.explain()+- BroadcastHashJoin [customer_id#75], [customer_id#159], LeftOuter, BuildRight, false, false
:- Scan ExistingRDD[order_id#74,customer_id#75,amount#76,order_date#77]
+- BroadcastExchange HashedRelationBroadcastMode(...)broadcast() marks a DataFrame as small enough to copy to every executor, so the large side does not have to be shuffled by key. Two details make the answer stronger. Spark already broadcasts automatically when a table's estimated size is under spark.sql.autoBroadcastJoinThreshold, 10 MB by default, and adaptive query execution can switch a join to a broadcast at runtime. And join hints are not guaranteed: the tuning guide says some strategies do not support every join type. Run explain() and point at BroadcastHashJoin instead of assuming.
The follow-up: "Now the data is a hundred times bigger"
This is a common place for interviewers to probe whether you understand the engine rather than only the syntax. Four answers cover most versions of the question.
Find the shuffles. Grouping, window partitionBy, non-broadcast joins and repartition() all move data between executors. The RDD programming guide calls the shuffle expensive because it "involves disk I/O, data serialization, and network I/O". In a plan, a shuffle appears as an Exchange node. A groupBy("k").count() shows a partial aggregation before the Exchange and a final one after it, so Spark pre-aggregates on each partition before moving anything, which is the DataFrame equivalent of preferring reduceByKey.
Look for skew. If one key holds a large share of the rows, one task does most of the work while the rest finish early. In the Spark UI, the stage detail page shows shuffle read size and summary metrics for task duration; one task running far longer than the median is the sign. Since Spark 3.2, adaptive query execution is on by default, and its skew-join handling splits a partition when it is more than five times the median partition size and larger than 256 MB. If the interviewer wants the manual fix, salt the key:
SALTS = 8
big_salted = big.withColumn("salt", (F.rand() * SALTS).cast("int"))
small_salted = small.crossJoin(spark.range(SALTS).withColumnRenamed("id", "salt"))
joined = big_salted.join(small_salted, on=["key", "salt"]).drop("salt")The cost to mention: the small side is copied eight times. The salted join should return the same rows as the plain join; check by comparing counts on your data.
Keep Python out of the hot path. The PySpark UDF guide says scalar Python UDFs rely on cloudpickle and "encounter performance bottlenecks" with large inputs. Prefer a built-in function; if you need custom logic, mention Arrow-optimised or pandas UDFs, which move data in columnar batches.
Do not collect. toPandas() loads the whole result into the driver's memory, and the API reference says to use it only when the result is expected to be small.
Concept questions that come with the code
- Transformation or action? Transformations (
select,filter,join) are lazy: Spark records them and computes nothing until an action (count,show,collect, a write) needs a result. That delay is what lets the optimiser plan the whole job. - `repartition()` or `coalesce()`?
repartition()does a full shuffle and can raise or lower the count; with column arguments it hash-partitions by those columns.coalesce()only lowers it, with a narrow dependency and no shuffle; the API reference warns that coalescing down to one partition can push the work onto a single node. - When should you cache? When the same DataFrame feeds more than one action.
cache()uses theMEMORY_AND_DISK_DESERstorage level by default. Caching something used once only costs memory. - Explicit schema or `inferSchema`? For CSV,
inferSchemais off by default and, when turned on, needs one extra pass over the data. An explicit schema skips that pass and keeps column types the same from run to run. - What does `checkpoint()` do? It saves the DataFrame to the checkpoint directory and truncates its logical plan, which helps iterative jobs whose plans keep growing.
A one-week practice plan
- Days 1-2: write problems 1 to 4 from a blank file without autocomplete, then run them and compare.
- Day 3: problems 5 and 6. Then break your own test data: add a tie, a null key, a zero-revenue month.
- Day 4: run
explain()on every solution and find eachExchange. - Day 5: talk through two problems aloud with a timer: the four questions, the code, then "a hundred times bigger".
- Day 6: answer each concept question in two sentences.
- Day 7: a full mock run with someone who will interrupt you.
For practice sessions, Verve AI's Coding Copilot captures a problem from your screen and returns a solution in Python or SQL, with quick actions such as Edge Case and Complexity that you can use to check your own version against. If your loop is at Databricks, the Databricks interview questions guide and the Databricks coding questions cover the company side, and the technical round questions cover the rest of a technical screen.
Sources
- dropDuplicates does not specify which row is kept: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.dropDuplicates.html
- dense_rank vs rank tie behaviour: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.dense_rank.html
- Join types including left_anti and left_semi: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html
- Lazy transformations, reduceByKey vs groupByKey note, shuffle cost: https://spark.apache.org/docs/latest/rdd-programming-guide.html
- unix_date returns days since 1970-01-01: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.unix_date.html
- try_divide returns null for a zero divisor: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.try_divide.html
- ANSI mode on by default since Spark 4.0: https://spark.apache.org/docs/latest/sql-migration-guide.html
- ANSI mode throws instead of returning null on invalid inputs: https://spark.apache.org/docs/latest/sql-ref-ansi-compliance.html
- broadcast() marks a DataFrame for broadcast joins: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.broadcast.html
- 10 MB broadcast threshold, AQE on by default since 3.2, skew-join thresholds, hints not guaranteed: https://spark.apache.org/docs/latest/sql-performance-tuning.html
- Spark UI stage detail (shuffle read size, task summary metrics): https://spark.apache.org/docs/latest/web-ui.html
- Scalar Python UDF performance bottlenecks, Arrow optimisation: https://spark.apache.org/docs/latest/api/python/user_guide/udfandudtf.html
- Pandas UDFs use Arrow for vectorised operations: https://spark.apache.org/docs/latest/api/python/tutorial/sql/arrow_pandas.html
- toPandas loads all data into driver memory: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.toPandas.html
- repartition does a full shuffle; hash-partitioned by the given columns: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.repartition.html
- coalesce is a narrow dependency; drastic coalesce runs on fewer nodes: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.coalesce.html
- cache default storage level MEMORY_AND_DISK_DESER: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.cache.html
- CSV inferSchema default false, one extra pass: https://spark.apache.org/docs/latest/sql-data-sources-csv.html
- checkpoint truncates the logical plan: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.checkpoint.html
- explain() modes: https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.explain.html
- Coding Copilot capture, languages and quick actions: https://docs.vervecopilot.com/features/coding-copilot
Ryan, CEO of Verve AI
Archive
