Write the Logic First

Logic comes first and execution second. An execution story about a query that returns the wrong rows is worth nothing, because the interviewer will correct the query and the story will no longer match it. The plan is also derived from the logic: the grouping column decides what the aggregate shuffles by, the join condition decides what the join shuffles by, and the ORDER BY decides whether a range shuffle appears at all. Change the query and you change the plan. So write the query, check it against a handful of rows, state the result out loud, and only then turn to the cluster.

The running example

Here is the table the rest of the lesson runs on: 7 orders, where each row's amount is tagged with a category and a status. The question is the one from the intro in a retail disguise: revenue per category, counting only completed orders, largest revenue first. Mapped back, category stands in for title_id and amount for seconds_watched. The filter on status replaces the filter on the day, so the plan built here is the answer to the Netflix tail. With 7 rows you can check the logic by hand and still follow each row as it moves between partitions in the next section.

order_idcategorystatusamount
1Bookscompleted40
2Gamescompleted25
3Booksreturned60
4Musiccompleted35
5Gamescompleted50
6Bookscompleted20
7Musiccompleted55
SELECT
    category,
    SUM(amount) AS revenue
FROM orders
WHERE status = 'completed'
GROUP BY category
ORDER BY revenue DESC

Check it by hand. Order 3 is a returned order for 60, so the WHERE drops it and Books falls from 120 to 60. Games collects 25 and 50 for 75. Music collects 35 and 55 for 90. The ORDER BY on the alias puts Music first, then Games, then Books. Saying those 3 numbers to the interviewer proves the filter sits before the aggregate, and it gives you a result to point at later, when you describe what the driver collects at the end of the job.

categoryrevenue
Music90
Games75
Books60

The same logic as a DataFrame

The interviewer may ask for the DataFrame form, or you may prefer to write it. Spark turns both into the same plan, so the choice changes nothing about the tail. The DataFrame version reads top to bottom in the order the operations run, which makes it useful for the second half: each method call in the chain is an operator you will later label as narrow or wide.

result = (orders
    .filter(F.col("status") == "completed")
    .groupBy("category")
    .agg(F.sum("amount").alias("revenue"))
    .orderBy(F.col("revenue").desc())
    .collect())

3 correctness traps the tail will expose

The traps below each produce a plausible query with a different plan, which is why they surface during the tail. A HAVING clause on status fails, because status is neither grouped nor aggregated. The usual patch is to add status to the GROUP BY. The query then runs and returns the right sums, but it groups by category and status, so the partial sums of returned orders cross the shuffle before HAVING discards them. A SUM over a column with NULL amounts returns the sum of the non-NULL values, which is right for revenue and wrong for a count of rows, so say which one you mean. And ORDER BY with ties returns those rows in whichever order the sort shuffle delivers them, so add a second sort key if the output feeds anything that compares runs.

Do
  • Filter on row attributes in WHERE, before the aggregate
  • State the result with real numbers before the tail
  • Name the grouping column out loud: it becomes the shuffle key
  • Add a tiebreaker to ORDER BY when order must be stable
Avoid
  • Patch a HAVING on status by grouping on status: returned rows cross the shuffle
  • Use COUNT(amount) when you mean COUNT(*)
  • Start on execution with a query you haven't checked
  • Rewrite the query during the tail without saying the plan changed

Each clause is already an operator

Once the query is right, the plan is a reading of it. Spark's optimizer turns each clause into an operator, and the physical plan it runs names them. The table below maps this example's clauses to those operators. Of its 5 rows, 2 carry an Exchange, and those 2 are what you'll point at.

ClausePhysical operatorsMoves rows between partitions?
FROM ordersFileScanNo
WHERE status = 'completed'FilterNo
GROUP BY category, SUM(amount)HashAggregate (partial), Exchange hashed by category, HashAggregate (final)Yes, once
ORDER BY revenue DESCExchange by revenue range, SortYes, once
collect()result returned to the driverRows leave the executors for the driver, outside the stages

Calling explain() on the DataFrame prints these names, and each Exchange hashpartitioning or Exchange rangepartitioning node marks a stage boundary.

About this lesson

This section is part of the SQL at Scale lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.