Your Query Gets Rewritten

Concepts covered: Catalyst Optimizer

Once an action needs a DataFrame's result, Catalyst takes the chain you built and produces a new plan that returns exactly the same rows with less work. You describe the result, and Catalyst picks the order of operations. It works through rules, each a rewrite that is safe to apply wherever its pattern appears in the plan. Nearly every job you write will hit 4 of them. 2 cut data early by moving filters toward the read and pruning unused columns; the other 2 fold constants and combine filters.

Filters move toward the read

The rewrite with the biggest payoff moves a filter as early as the plan allows. Take a report of Completed orders sorted by profit, written in the order its author thought of the steps:

ranked = (
    orders
    .withColumn("profit_k", F.col("profit") / 1000)
    .orderBy(F.col("profit_k").desc())
    .filter(F.col("status") == "Completed")
)

Read literally, this chain sorts every order and then throws most of them away. A global sort needs a shuffle: every row crosses the network so that each task owns a slice of the sorted range. Catalyst sees that the filter reads only status. That column exists before withColumn and the sort leaves it untouched, so testing it first can't change which rows come out. Catalyst moves the filter below the sort and below withColumn, right after the read. Say orders holds 8,000,000 rows and 2,000,000 of them are Completed. As written, the sort would shuffle all 8,000,000. After the rewrite it shuffles 2,000,000.

• 75%

The animation runs that chain on 8 orders read as 2 partitions. Watch where the Cancelled and Pending orders disappear. The read tasks drop them, so only the 4 Completed orders cross the shuffle to be sorted, and order 3, with a profit of 4500, comes out on top.

(
    orders
    .withColumn("profit_k", F.col("profit") / 1000)
    .orderBy(F.col("profit_k").desc())
    .filter(F.col("status") == "Completed")
    .collect()
)

When a filter cannot move

A filter moves below a step only when that step can't affect its answer. Put a filter on region after groupBy("region") and it moves below the grouping, because removing a region's rows before grouping removes exactly that region's group. A filter on total_profit > 5000 after the same grouping stays where it is. total_profit does not exist until the grouping has run, and a rule cannot test a value nobody has computed yet.

You can apply the rule yourself. A filter runs earlier only if it can be written in columns that exist at that earlier point. A derived column qualifies when its formula reads such columns, so Catalyst rewrites profit_k > 2 as profit / 1000 > 2 and moves it too. An aggregate's result never qualifies, because computing it takes every row of the group.

Which filter can Catalyst run before the shuffle that this groupBy needs?

totals = (
    orders
    .groupBy("region")
    .agg(F.sum("profit").alias("total_profit"))
    .filter(F.col("region") != "Test")
    .filter(F.col("total_profit") > 5000)
)

• Only the region filter: region exists before grouping, and total_profit comes from it
• Both filters, because Catalyst moves every filter down to the read
• Only the total_profit filter, because it removes more rows than the other
• Neither, because a filter written after groupBy always runs after it

Columns nobody reads are never carried

The second rewrite works on columns. Catalyst walks the plan from the result back to the source and records which columns each step needs. A column that no step needs is cut at the read, so it is never carried between steps and never shuffled. When the files store each column separately, Spark also skips reading that column from disk. Wide tables are where this pays: a query that touches 3 columns of a 120-column events table moves 3 columns' worth of data through its shuffles.

In the next animation the query averages price by category over a 6-column products table. Catalyst dropped product_id and product_name before the first task ran, along with supplier and weight_kg, so none of those 4 columns crosses the shuffle. Kitchen comes out at (45 + 60 + 35) / 3, about 46.67.

products.groupBy("category").agg(F.avg("price").alias("avg_price")).collect()

The guarantee and its limit

Every rewrite returns the rows your written chain would have returned. That's why Catalyst applies them without asking you. Its limit is that it can only rewrite what it can read. It cannot remove a shuffle your logic requires: a groupBy still groups and an orderBy still sorts. It also can't see inside a Python function. Write the query in whatever order reads clearest, and spend your effort on the steps Catalyst cannot change.

Run this version, where every filter comes last. Then move the status and region filters up to sit right after spark.table("orders") and run again: the rows match. Now try moving the profit_k filter up as well. Spark rejects it at once, because profit_k does not exist at that point in your code. Catalyst can make that move by rewriting the filter in terms of profit, but your code cannot.

About this lesson

This section is part of the Catalyst Optimizer: Beginner 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.