Apache Spark Query Optimization on Databricks: Catalyst, AQE, and Photon Engine

A deep dive into how Spark transforms your SQL into a physical execution plan — and how Databricks layers Adaptive Query Execution and the Photon vectorized engine on top to squeeze out maximum performance.

Why Query Optimization Matters
The Catalyst Optimizer Pipeline
Stage 1: Parsing — From SQL to Unresolved Logical Plan
Stage 2: Analysis — Binding to the Catalog
Stage 3: Logical Optimization — Rule-Based Rewrites
Stage 4: Physical Planning — Strategies and Cost Models
Adaptive Query Execution (AQE)
The Photon Engine
Reading Explain Plans
Tuning Reference Table
References

Why Query Optimization Matters

A Spark query written by a human and a Spark query executed by the engine are often very different things. The gap between them — the optimization — is what separates a job that runs in 3 minutes from one that runs in 3 hours on identical hardware.

Databricks compounds Spark's native Catalyst optimizer with two additional layers:

Adaptive Query Execution (AQE) — re-optimizes the query at runtime using actual statistics collected mid-job
Photon — a C++ vectorized execution engine that replaces the JVM-based Spark executor for eligible operators

Understanding all three lets you write queries that cooperate with the engine rather than fight it.

The Catalyst Optimizer Pipeline

Catalyst is Spark's rule-based and cost-based query optimizer. Every query — whether written in SQL, DataFrame API, or Dataset API — passes through the same four-stage pipeline before a single byte of data is read.

Stage 1: Parsing — From SQL to Unresolved Logical Plan

# ── Catalyst Stage 1: Parsing ─────────────────────────────────────────────────
# Spark uses ANTLR4 to parse SQL into an Abstract Syntax Tree (AST).
# At this point column names are NOT validated — the plan is "unresolved".

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("catalyst-demo").getOrCreate()

# Both of these produce identical internal representations
df_api = (
    spark.table("prod.silver.events_clean")
        .filter("event_type = 'purchase'")
        .groupBy("platform")
        .agg({"revenue": "sum"})
)

sql_api = spark.sql("""
    SELECT platform, SUM(revenue) AS total_revenue
    FROM prod.silver.events_clean
    WHERE event_type = 'purchase'
    GROUP BY platform
""")

# Inspect the unresolved logical plan (before analysis)
df_api.explain(mode="formatted")
# Output includes:
# == Parsed Logical Plan ==
# 'Aggregate ['platform], ['platform, unresolvedAlias('sum('revenue), None)]
#   +- 'Filter ('event_type = 'purchase)
#      +- 'UnresolvedRelation [prod, silver, events_clean]

The key insight here: UnresolvedRelation and unresolvedAlias mean Spark hasn't touched the catalog yet. Column names could be typos at this point and Catalyst doesn't know.

Stage 2: Analysis — Binding to the Catalog

The Analyzer walks the unresolved AST and looks up every relation and attribute against the Catalog (in Databricks, this is Unity Catalog). It resolves column names, infers data types, validates references, and binds functions.

# ── Catalyst Stage 2: Analysis ────────────────────────────────────────────────
# After analysis, every column is resolved to a specific attribute with a type.
# AnalysisException is thrown HERE if a column doesn't exist.

from pyspark.sql import functions as F
from pyspark.sql.utils import AnalysisException

# Example of what Analysis catches:
try:
    spark.table("prod.silver.events_clean") \
        .select("nonexistent_column") \
        .show()
except AnalysisException as e:
    print(f"Analysis failed: {e}")
    # → AnalysisException: [UNRESOLVED_COLUMN.WITH_SUGGESTION]
    #   A column or function parameter with name `nonexistent_column` cannot be resolved.

# After successful analysis, inspect the resolved plan
df = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")
        .select("platform", "revenue", "user_id")
)

# The analyzed plan shows fully qualified attribute IDs like:
# == Analyzed Logical Plan ==
# platform: string, revenue: double, user_id: string
# Project [platform#42, revenue#67, user_id#31]
#   +- Filter (event_type#39 = purchase)
#      +- Relation prod.silver.events_clean[...] parquet
print(df._jdf.queryExecution().analyzed())

Stage 3: Logical Optimization — Rule-Based Rewrites

This is where Catalyst applies its ~100+ built-in rules to produce an equivalent but cheaper logical plan. Rules fire repeatedly in fixed-point iteration until the plan stabilises.

# ── Catalyst Stage 3: Key Optimization Rules ──────────────────────────────────

# RULE 1: Predicate Pushdown
# Catalyst moves filters as close to the data source as possible,
# so Spark reads fewer rows from Parquet.

df_before = (
    spark.table("prod.silver.events_clean")
        .join(
            spark.table("prod.silver.users_clean"),
            on="user_id"
        )
        .filter(F.col("event_type") == "purchase")  # ← filter AFTER join
)

# Catalyst rewrites this internally as if you wrote:
df_after_equivalent = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")  # ← filter BEFORE join
        .join(
            spark.table("prod.silver.users_clean"),
            on="user_id"
        )
)
# Result: potentially millions fewer rows shuffled during the join

# RULE 2: Column Pruning
# Catalyst removes columns not needed by downstream operators.
# Even if you select(*), Spark will only read the columns it needs.
df_pruned = (
    spark.table("prod.silver.events_clean")
        .select("*")
        .filter(F.col("event_type") == "purchase")
        .groupBy("platform")
        .agg(F.sum("revenue").alias("total_revenue"))
)
# Internally, Catalyst prunes all columns except: event_type, platform, revenue

# RULE 3: Constant Folding
# Expressions with only literals are evaluated at plan time, not per-row.
df_constants = spark.range(1000).select(
    F.lit(2 + 3 * 4).alias("always_14"),      # folded to Literal(14) at plan time
    F.col("id") * F.lit(1).alias("same_id"),   # simplified to just col("id")
)

# RULE 4: Boolean Simplification
# AND/OR chains with tautologies or contradictions are collapsed
df_simplified = spark.range(100).filter(
    (F.col("id") > 10) & F.lit(True)   # simplified to just (col("id") > 10)
)

# See all optimizations applied:
print(df_pruned._jdf.queryExecution().optimizedPlan())

Stage 4: Physical Planning — Strategies and Cost Models

The physical planner maps each logical operator to one or more physical implementations and selects the best one using a cost model. The most impactful decision here is join strategy selection.

# ── Catalyst Stage 4: Physical Planning & Join Strategies ────────────────────

# JOIN STRATEGY 1: Broadcast Hash Join (BHJ)
# Best when one side is small enough to fit in executor memory.
# No shuffle — the small table is broadcast to all workers.

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "10mb")  # default

large_df = spark.table("prod.silver.events_clean")      # 500GB
small_df = spark.table("prod.gold.product_catalog")     # 8MB ← will be broadcast

result_bhj = large_df.join(small_df, on="product_id")  # BHJ auto-selected

# Force BHJ with a broadcast hint (overrides threshold check):
from pyspark.sql.functions import broadcast
result_forced = large_df.join(broadcast(small_df), on="product_id")


# JOIN STRATEGY 2: Sort Merge Join (SMJ)
# Default for large-large joins. Both sides are sorted and merged.
# Requires a full shuffle — expensive but handles any size.

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")  # disable BHJ
large_df2 = spark.table("prod.silver.transactions_clean")   # 200GB

result_smj = large_df.join(large_df2, on="user_id")  # SMJ selected


# JOIN STRATEGY 3: Shuffle Hash Join (SHJ)
# Hash-based, no sort. Chosen by AQE when one side is much smaller
# than the other but still above the broadcast threshold.
spark.conf.set("spark.sql.join.preferSortMergeJoin", "false")


# WHOLE-STAGE CODEGEN: Spark fuses multiple operators into a single
# Java function to avoid virtual dispatch overhead and intermediate objects.
# Verify it's active in your plan:
spark.conf.set("spark.sql.codegen.wholeStage", "true")  # default

result_bhj.explain(mode="formatted")
# Look for: *(1) BroadcastHashJoin — the *(N) prefix = WholeStageCodegen stage N

Adaptive Query Execution (AQE)

AQE is Databricks' most impactful runtime optimization layer. It materializes shuffle map output statistics at shuffle boundaries and uses them to make three key decisions after data has been partially processed.

# ── AQE Configuration ─────────────────────────────────────────────────────────

# AQE is ON by default in Databricks Runtime 7.3+
spark.conf.set("spark.sql.adaptive.enabled", "true")

# 1. Dynamic Partition Coalescing
# Merges small post-shuffle partitions to avoid thousands of tiny tasks
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled",       "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes",     "128mb")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionNum", "1")

# 2. Dynamic Join Strategy Switching
# Allows AQE to downgrade SMJ → BHJ at runtime if a side turns out small
spark.conf.set("spark.sql.adaptive.localShuffleReader.enabled",       "true")

# AQE broadcast threshold (can be higher than static threshold since
# we now KNOW the actual size)
spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold",       "30mb")

# 3. Skew Join Optimization
# Splits oversized partitions and replicates the non-skewed side
spark.conf.set("spark.sql.adaptive.skewJoin.enabled",                 "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor",   "5")    # 5x median
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256mb")

# Verify AQE decisions in the query plan:
df = (
    spark.table("prod.silver.events_clean")
        .join(spark.table("prod.silver.users_clean"), on="user_id")
        .groupBy("platform")
        .agg(F.sum("revenue").alias("total"))
)
df.explain(mode="formatted")
# Look for: AdaptiveSparkPlan isFinalPlan=true
# and: == Final Physical Plan == (shows post-AQE decisions)

The Photon Engine

Photon is Databricks' native vectorized query engine written in C++. It replaces the JVM-based Spark executor for eligible operations, processing data in column-oriented batches (vectors) rather than row-by-row.

# ── Photon Configuration & Verification ───────────────────────────────────────
# Photon is available on Databricks Runtime 9.1+ with Photon-enabled clusters.
# Enable it at the cluster level (UI: Cluster > Configuration > Enable Photon)
# or via config:

spark.conf.set("spark.databricks.photon.enabled", "true")

# Photon-accelerated operators (as of DBR 13.x):
# ✅ Scan (Parquet, Delta)         ✅ Filter / Project
# ✅ Hash Aggregate                ✅ Sort
# ✅ Broadcast Hash Join           ✅ Sort Merge Join
# ✅ Window functions              ✅ Union / Expand
# ✅ String functions              ✅ Math functions
# ❌ UDFs (Python/Scala)          ❌ Some complex types
# ❌ Streaming (partial)          ❌ RDD-based operations

# Verify Photon is executing your query:
df = spark.sql("""
    SELECT
        platform,
        DATE_TRUNC('month', event_ts)   AS month,
        SUM(revenue)                    AS total_revenue,
        COUNT(DISTINCT user_id)         AS unique_buyers,
        AVG(revenue)                    AS avg_order_value
    FROM prod.silver.events_clean
    WHERE event_type = 'purchase'
      AND event_ts >= '2024-01-01'
    GROUP BY platform, DATE_TRUNC('month', event_ts)
    ORDER BY month DESC, total_revenue DESC
""")

df.explain(mode="formatted")
# Look for operators prefixed with "Photon" in the physical plan:
# == Physical Plan ==
# PhotonResultStage
# +- PhotonSort [month DESC NULLS LAST, total_revenue DESC NULLS LAST]
#    +- PhotonShuffleExchangeSink hashpartitioning(platform, month)
#       +- PhotonGroupingAgg [platform, month], [sum(revenue), count(user_id), avg(revenue)]
#          +- PhotonFilter (event_type = purchase AND event_ts >= 2024-01-01)
#             +- PhotonScan parquet prod.silver.events_clean

# Photon performance metrics appear in Spark UI under "Photon Metrics":
# - Photon scan time
# - Photon total compute time
# - Rows processed by Photon vs fallback JVM

Reading Explain Plans

The explain(mode="formatted") output is your primary debugging tool. Here's how to read it efficiently:

# ── Explain Plan Modes ────────────────────────────────────────────────────────

df = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")
        .join(broadcast(spark.table("prod.gold.product_catalog")), on="product_id")
        .groupBy("platform", "category")
        .agg(
            F.sum("revenue").alias("total_revenue"),
            F.count("*").alias("transaction_count")
        )
)

# Mode 1: simple (default) — compact tree
df.explain()

# Mode 2: extended — all 4 plan stages side by side
df.explain(mode="extended")

# Mode 3: formatted — human-readable with operator details (RECOMMENDED)
df.explain(mode="formatted")

# Mode 4: cost — includes estimated row counts and sizes (requires ANALYZE TABLE)
df.explain(mode="cost")

# Mode 5: codegen — shows generated Java code for WholeStageCodegen
df.explain(mode="codegen")


# ── Key Signals to Look For ───────────────────────────────────────────────────

# ✅ GOOD signs:
# *(N) prefix          → WholeStageCodegen active (operators fused)
# BroadcastHashJoin    → small table correctly broadcast, no shuffle
# PhotonXxx            → Photon accelerating this operator
# AdaptiveSparkPlan    → AQE is engaged
# PartitionFilters     → Delta/Parquet file skipping active
# PushedFilters        → filters pushed to Parquet reader

# ❌ WARNING signs:
# Exchange (shuffle)   → unexpected shuffle (missing broadcast hint?)
# SortMergeJoin        → large-large join (may need Z-ORDER or AQE tuning)
# HashAggregate x2     → partial + final agg = shuffle involved
# CartesianProduct     → missing join condition! Will OOM on large tables
# ObjectHashAggregate  → non-codegen path, JVM overhead
# GenerateXxx          → explode() or similar, can't be fused


# ── ANALYZE TABLE: feed statistics to CBO ─────────────────────────────────────
# Without stats, Catalyst uses default estimates (1M rows, 8 bytes/col).
# Run ANALYZE to give the Cost-Based Optimizer real numbers.

spark.sql("ANALYZE TABLE prod.silver.events_clean COMPUTE STATISTICS")
spark.sql("""
    ANALYZE TABLE prod.silver.events_clean
    COMPUTE STATISTICS FOR COLUMNS user_id, event_type, platform, revenue
""")
# Now explain(mode="cost") shows real row counts and sizes

Tuning Reference Table

A quick-reference guide for the most impactful Spark/Databricks configs, what they control, and when to change them:

Config Key	Default	What It Controls	When to Tune
`spark.sql.adaptive.enabled`	`true`	Master AQE switch	Keep on; only disable for debugging
`spark.sql.adaptive.advisoryPartitionSizeInBytes`	`64mb`	Target post-coalesce partition size	Increase to `128mb`–`256mb` for large shuffles
`spark.sql.adaptive.skewJoin.enabled`	`true`	AQE skew split	Keep on; tune `skewedPartitionFactor` if needed
`spark.sql.autoBroadcastJoinThreshold`	`10mb`	Static BHJ threshold	Increase to `50mb`–`100mb` if executor memory allows
`spark.sql.adaptive.autoBroadcastJoinThreshold`	`30mb`	AQE runtime BHJ threshold	Increase if AQE isn't catching small tables
`spark.sql.shuffle.partitions`	`200`	Default shuffle partition count	Set to `8 × num_cores` for your cluster
`spark.sql.files.maxPartitionBytes`	`128mb`	Max bytes per Parquet read partition	Reduce for high-parallelism scans
`spark.databricks.photon.enabled`	`true`	Photon vectorized engine	Keep on; disable only for UDF-heavy jobs
`spark.sql.codegen.wholeStage`	`true`	Whole-Stage CodeGen fusion	Keep on; disable only for debugging
`spark.sql.statistics.histogram.enabled`	`false`	Column histograms for CBO	Enable after running ANALYZE TABLE
`spark.sql.cbo.enabled`	`true`	Cost-Based Optimizer	Keep on; requires ANALYZE TABLE to be useful
`spark.databricks.delta.optimizeWrite.enabled`	`true`	Auto bin-pack write files	Keep on for all Delta writes

Key Takeaways

Catalyst has four stages: Parse → Analyze → Optimize → Plan. Each stage has a distinct job, and understanding them tells you exactly where to look when a query misbehaves.
Predicate pushdown and column pruning are the two most impactful automatic optimizations — they reduce the data volume Spark has to move before any aggregation or join.
AQE is not a set-and-forget feature: tune advisoryPartitionSizeInBytes to your actual data sizes, and verify its decisions with explain(mode="formatted") — look for AdaptiveSparkPlan isFinalPlan=true.
Photon drops in transparently for most SQL and DataFrame operations. The exceptions are Python UDFs, RDD operations, and some complex types — refactor these away from hot paths.
Run ANALYZE TABLE ... COMPUTE STATISTICS FOR COLUMNS on your most-joined tables. The CBO's join ordering and strategy decisions improve dramatically with real statistics vs. default estimates.
explain(mode="formatted") is your most important debugging tool — learn to read it before reaching for cluster config changes.