🍎 WholeStageCodegen: The Plan Was Already Written 📓

The plan is evidence. The generated Java is the confession.

You write a DataFrame transformation, Spark gets to work, and usually that is enough. Until the physical plan starts showing names like *(1), Exchange, or BatchEvalPython—and the simple API begins to look like a cover story.

Underneath, Catalyst rewrites the query, Spark chooses physical operators, and compatible row-based fragments can be fused into generated Java by WholeStageCodegen.

This article follows those names: what Spark can see, what it generates, where the pipeline breaks, and when codegen is not the problem at all. The runnable experiments live in small companion notebooks.

Reading the Spark physical plan like it killed someone.

“I won’t let anyone get in my way”
(from “Death Note 1st Opening Theme”)

⚡ The First Name In The Plan

WholeStageCodegen sounds more mysterious than it is.

From Operators to a Generated Pipeline

Suppose Spark chooses:

Scan ➔ Filter ➔ Project

If those physical operators and their expressions are codegen-eligible, Spark can fuse them into one WholeStageCodegen stage and emit Java specialized for that fragment. Instead of repeatedly crossing generic operator boundaries, the hot path becomes closer to one program: read values, test predicates, compute expressions, emit results.

The useful mental shift is:

generic execution machinery ➔ specialized program for this plan

Your DataFrame describes the query. Spark may write another program to execute part of it.

The First Clue: *(1)

A normal physical plan may show:

*(1) Project
+- *(1) Filter
+- *(1) Range

The repeated *(1) marks operators in the same codegen stage. explain(“formatted”) exposes the same relationship through codegen IDs in node details.

Later, when a boundary splits generated execution, you may see the IDs change:

*(2) HashAggregate
+- Exchange
+- *(1) HashAggregate

Two IDs, two generated regions. We’ll investigate why the pipeline breaks there later.

The Star is Evidence, not an Alibi

*(1) means generated execution exists here. It doesn’t mean this query is fast.

A perfectly generated stage can still wait on shuffle, storage, spill, skew, or eleven million tiny Parquet files.

🍎 L wouldn’t close a case because he found one fingerprint. Neither should you.

🧪 Lab 01 — Find the Stage

01__find_the_stage.ipynb uses a tiny Filter => Project query. Inspect the normal and formatted plans, compare the codegen markers, and identify which operators share a stage.

The conclusion is intentionally modest: a shared codegen ID tells you where Spark fused execution, not whether that stage is the bottleneck.

🩸 The Cost of Being Generic

Now that we can see codegen, the useful question is: Why did Spark build it?

Because tiny execution costs get very interesting when they repeat a billion times.

Tiny Costs Have Many Friends

A general engine needs machinery that works across different operators, expressions, types, and schemas. That can mean repeated dispatch, operator hand-offs, intermediate row handling, generic expression machinery, boxing/object allocation, and runtime checks that specialization can sometimes simplify or inline.

Modern Spark already removes plenty of older overhead through Tungsten-era execution ideas such as compact internal rows. WholeStageCodegen pushes further: if the path is known, specialize it.

A tiny avoidable cost is boring once. A billion times, it gets a motive.

Codegen Has to Earn its Cost

Generated Java isn’t free. Spark must generate source, compile it, load the resulting class, and execute it. In Spark 4.2, the code-generation path uses Janino for runtime Java compilation; the JVM JIT may later optimize hot bytecode again.

For a tiny query, compilation can be noticeable. For a sufficiently large, repetitive workload, that setup cost can be amortized across the work.

So the useful question is not: Is WholeStageCodegen faster?
It is: Is enough repeated local CPU work happening here for codegen to pay for itself?

🧪 Lab 02 — Make It Pay

02__make_codegen_pay.ipynb runs the same CPU-heavy calculation with WholeStageCodegen enabled and disabled, feeding the computed value into sum so Catalyst cannot prune it. We first confirm that the physical execution path changes, then compare repeated warm runs.

The experiment isolates WholeStageCodegen, not every form of Spark code generation: if the warm timings differ, we have evidence that stage-level fusion mattered for this workload. If they barely move, the dominant cost was somewhere else.

📓The Plan Becomes A Program

A DataFrame does not jump directly from Python into generated Java. Several different systems touch it first.

Catalyst writes the case

A simplified journey is:

DataFrame / SQL
↓
logical plan
↓
analysis + optimization
↓
physical plan

Catalyst resolves names and types, simplifies expressions, prunes columns, propagates constraints, and reshapes the query before physical execution begins.

At this point Spark has established what the query means and how it can be logically rewritten. Physical execution comes next.

WholeStageCodegen Comes Later

Once physical operators exist, CollapseCodegenStages looks for row-based regions that support code generation. In Spark 4.2, that rule explicitly rejects expressions carrying CodegenFallback and checks input/output schema width before admitting an operator to a WholeStageCodegen stage.

The result is often several codegen islands, not necessarily one monster Java class for the whole query. It’s a set of codegen islands built from planning decisions already made.

Four Different Things People Call “Compilation”

Keep these layers separate:

Catalyst optimization ➔ changes the plan Spark code generation ➔ emits Java Janino ➔ compiles Java to JVM bytecode JVM JIT ➔ may optimize hot bytecode further

They cooperate. They are not aliases.

🍎 Four suspects entered the room wearing the word compiler. L makes them sit separately.

🧪 Lab 03 — Follow the Plan

03__follow_the_plan.ipynb inspects one query with extended, formatted, and codegen explain modes. The query stays fixed; only the camera changes.

By the end, the logical plan should answer what Spark understood, the physical plan how Spark intends to execute it, and the codegen view which stages produced Java.

🚧 Where the Pipeline Breaks

WholeStageCodegen usually appears in fragments. The interesting clue is often where a fragment stops.

Exchange: Data Has to Move

*(2) HashAggregate
+- Exchange hashpartitioning(...)
+- *(1) HashAggregate

Exchange marks data exchange—commonly a shuffle/repartition, and in other shapes a broadcast exchange. Codegen can optimize local work around it; it cannot turn distributed data movement into one local Java loop.

In this shape, the stage ID changes because the shuffle separates two generated execution regions.

InputAdapter: the Bridge Back in

InputAdapter lets a code-generated parent consume rows from a child that’s not part of the same generated stage:

non-codegen child ➔ InputAdapter ➔ codegen stage

It quietly tells you: something below this point could not join the same program.

WholeStageCodegen itself is row-based, so columnar operators can also mark a change of execution regime.

Python is Another Jurisdiction

A Python UDF creates a different runtime boundary. Spark may generate JVM code around it, but arbitrary Python cannot be pasted into the same generated Java pipeline.

Spark 4.2 enables Arrow optimization for regular Python UDFs by default when the required dependencies are available. Arrow improves JVM↔Python transfer and serialization; the scalar Python function still executes as Python logic.

Cheaper border. Still a border.

A Boundary is not Automatically a Bug

An Exchange may be required. A Python UDF may be justified. Another operator may deliberately use a different strategy.

The boundary tells you where to investigate, not what to condemn. And with Adaptive Query Execution, prefer the final executed plan when performance matters — the first plan may not be the final story.

🧪 Lab 04 — Find the Boundaries

04__find_the_boundaries.ipynb introduces repartitioning and a small Python-UDF example. Trace where codegen IDs change and where execution crosses into another strategy or runtime.

The lesson: WholeStageCodegen is local to compatible plan fragments; shuffles and runtime boundaries can split those fragments into separate generated regions.

🔬The Confession Is In processNext()

The physical plan says code was generated. explain(“codegen”) lets us read it.

You may find GeneratedIteratorForCodegenStage1 and a method called processNext().

That’s where the diagram becomes a program.

Operators Disappear into the Loop

A physical Filter and Project do not survive as two nicely labeled Java objects. Their logic is woven together:

while (input.hasNext()) {
// read values
// test predicate
// compute projection
// emit/update state
}

The operator boundaries have largely disappeared from the hot path. That is the fusion.

Why the Java Looks Mildly Possessed

Expect null flags, UnsafeRowWriter, aggregation state, metrics, and names such as project_value_0 or filter_isNull_2.

Spark is exposing what the DataFrame API hid: SQL null semantics, internal row access, type-specific reads/writes, and mutable state.

Read it Like Evidence, not Literature

Do not read 900 lines top to bottom. Find:

GeneratedIteratorForCodegenStage...
processNext()
the predicate
the arithmetic

Then map them back to the physical plan.

🧪 Lab 05 — Read the Confession

05__read_the_generated_java.ipynb prints generated Java for one small Range => Filter => Project pipeline. Locate processNext() and match the original predicate and expression to the emitted Java.

Nothing is benchmarked. We are proving that the operators in the plan became one generated execution path.

👁️ What Catalyst Cannot See

There is another boundary that may not look dramatic in the physical plan: the boundary of what Catalyst understands.

Built-ins Are Structured Evidence

A native expression such as F.col(“amount”) * 1.21 is represented as a Catalyst expression tree with known children, types, null behavior, and semantics.

That visibility gives Spark opportunities to simplify expressions, prune work, push eligible predicates toward data sources, reuse common subexpressions, and generate code.

The important word is visible.

A UDF is a Sealed Envelope

A UDF call is represented in Catalyst, but the function body is not exposed as the same kind of expression tree. Spark can optimize around the call; it generally cannot open the function and rewrite its internal algorithm as native Catalyst expressions.

Same answer. Different evidence.

JVM UDFs Solve a Different Problem

A Scala/Java UDF avoids the Python worker boundary, but its implementation is still largely opaque to Catalyst’s optimizer.

These are separate questions:

Where does the function execute?
How much of its logic can Catalyst understand?

Being in the JVM is not the same as being optimizer-visible.

🧪 Lab 06 — Hide the Evidence

06__hide_the_evidence.ipynb writes a small Parquet dataset and applies the same filter once natively and once with the logic hidden inside a UDF. Inspect what reaches the scan and what does not.

The point is planning, not stopwatch theater: optimization opportunities can be lost before the first row executes.

🧨 When Generated Code Becomes The Problem

Eventually the optimization becomes something Spark needs protection from.

You Can Generate too much Java

Huge projections, deep expressions, and giant CASE trees can produce giant methods and classes.

The JVM requires a method’s bytecode code_length to be less than 65,536 bytes. Spark therefore contains code-splitting machinery and other safeguards for large generated code.

A 20,000-branch CASE WHEN is not a business rule. It’s a boss fight with procurement approval.

Wide Plans Have Gravity

WholeStageCodegen admission also checks input and output schema width. Hundreds of unnecessary columns mean more state, null handling, writers, and generated variables.

select(“*”) is convenient. It’s not a query-planning philosophy.

Fallback Can Be Quiet

Spark can fall back if WholeStageCodegen compilation fails. It can also disable WholeStageCodegen for a plan when a compiled method exceeds its configured huge-method threshold, because JIT optimization may suffer.

yesterday ➔ generated path
today     ➔ fallback / disabled WSC
result    ➔ correct
runtime   ➔ haunted

No dramatic exception. Just a regression.

Don’t immediately raise the limit

If Spark is generating a Java cathedral, first ask why. Prune columns, simplify generated expressions, replace giant literal branching with data where appropriate, or split unhealthy transformations.

A larger threshold may tolerate a larger monster. It doesn’t make the monster good.

🧪 Lab 07 — Make the Compiler Flinch

07__make_the_compiler_flinch.ipynb grows a CASE WHEN expression gradually, tracks generated-code growth and reported method sizes, and checks whether the expression remains inside WholeStageCodegen.

The goal is not a universal “failure at N branches.”
It’s to see the limits before Ryuk asks for 50,000.

🕳️The Expression With Two Alibis

GraphFrames 0.12.1 gives us a beautifully cursed case: FiniteAXPlusB.

The Suspicious Combination

Its source declares extends TernaryExpression with CodegenFallback and the same class overrides doGenCode() with a real finite-field arithmetic loop.

So we have CodegenFallback and custom generated code in one expression.

doGenCode() is not the verdict

It’s tempting to see doGenCode() and conclude that the surrounding operator joins WholeStageCodegen. Not so fast.

Spark 4.2’s CollapseCodegenStages explicitly treats expressions carrying CodegenFallback as unsupported for WholeStageCodegen admission.

So both facts can coexist:

FiniteAXPlusB contains custom expression-codegen logic.
FiniteAXPlusB prevents its operator from joining a WSC stage.

Expression code generation and WholeStageCodegen eligibility are related layers, not synonyms.

GraphFrames Really Uses it

GraphFrames’ randomized-contraction connected-components implementation registers FiniteAXPlusB as _axpb and uses it in DataFrame expressions.

That gives us a useful hierarchy of evidence:

source code    ➔ what the expression implements
physical plan  ➔ what Spark admitted
codegen output ➔ what Spark emitted

🍎 Light reads doGenCode() and closes the case. L keeps reading.

🧪 Lab 08 — Put FiniteAXPlusB on Trial

08__finite_ax_plus_b_on_trial.ipynb verifies the GraphFrames 0.12.1 source and, when the matching runtime is available, runs randomized contraction and inspects its plan and codegen output. Otherwise, the runtime verdict remains unmeasured.

The verdict: having a doGenCode() implementation is not proof of WholeStageCodegen participation.

📌 Case update: After reading a draft of this article, Sem opened GraphFrames PR #888 to remove CodegenFallback from FiniteAXPlusB. Apparently, the investigation may have just changed the suspect’s future.

🧭Cause Of Death Matters

WholeStageCodegen is a local CPU optimization inside a distributed system. That sentence is more useful than any star in the plan.

It can help repetitive work such as filtering, arithmetic, hashing, aggregation updates, and join probes. It cannot erase shuffle, skew, spill, storage latency, bad partitioning, tiny files, or a poor join strategy.

You optimized the murder weapon. The victim drowned.

Amdahl’s law applies mercilessly: if only a small fraction of runtime is codegen-sensitive, even a dramatic improvement there barely moves the whole query.

🧪 Lab 09 — Find the Cause

09__find_the_cause.ipynb compares a CPU-heavy local workload with one dominated more by redistribution. Inspect plans and stage metrics, then decide whether local computation or redistribution is the stronger suspect.

The lesson is not: Codegen wins.
It’s: The optimization must match the dominant cost.

🖋️Write The Name Yourself

There is one level deeper: teach Catalyst a new expression and let Spark generate evaluation code for it.

A UDF gives Spark a function; an expression gives Spark semantics

A custom Catalyst expression can declare children, input/output types, null behavior, interpreted evaluation, and generated evaluation.

That moves from:

Spark ➔ call my function

toward:

Spark ➔ represent my expression structurally ➔ generate its evaluation code

That’s a much deeper contract.

Two Paths Must Tell the Same Story

A custom expression often has an interpreted path such as nullSafeEval(…) and a generated path such as doGenCode(…).

They must agree for ordinary values, NULL, edge cases, overflow behavior, and randomized inputs.

A fast wrong expression is still wrong. It just reaches the incident report sooner.

Registration Brings the Maintenance Bill

One integration route is SparkSessionExtensions.injectFunction. Spark 4.2 marks SparkSessionExtensions experimental/unstable and explicitly gives no binary or source compatibility guarantee.

So your tiny expression can acquire a Scala project, JAR packaging, registration, classpath management, version compatibility, codegen tests, and deployment.

Congratulations. You optimized fourteen nanoseconds and accidentally founded a platform team.

Test Three Layers

Keep these separate:

Interpreted correctness

generated-expression correctness

WholeStageCodegen integration

Turning spark.sql.codegen.wholeStage off addresses the third layer. It’s not a universal switch for every form of expression code generation.

🧪 Lab 10 — Write the Name

10__write_the_name_yourself.ipynb uses one intentionally boring custom expression. Compare a reference implementation with an equivalent generated Spark expression across edge cases, then inspect the plan and codegen output. The actual custom Catalyst integration still requires the Scala/JAR extension.

The lab makes the trade-off visible: custom Catalyst work is for hot, stable paths worth owning — not every slow UDF that annoyed you on Tuesday.

🕵The Performance Crime Scene

We now have enough clues. What we need is a procedure.

Start with what Actually Ran

Catalyst may rewrite or prune work; AQE may change the physical plan during execution. So begin with the executed plan, not the DataFrame chain you remember writing.

The source code is the witness statement. The executed plan is closer to CCTV.

Then Ask where the Time Went

Use the Spark UI and stage metrics to look for disproportion: huge shuffle read/write, spill, long GC, pathological tasks, excessive scan input, or high CPU with little I/O.

plan    ➔ what Spark built
metrics ➔ where it hurt

Only then decide whether generated Java deserves your attention.

If 85% of runtime is shuffle read, studying processNext() is mostly performance-themed procrastination.

Make Sure the Benchmark Executes the Thing

This can be misleading:

df.select(expensive_expression).count()

If the projection cannot affect row cardinality, Catalyst may prune it. You benchmarked Spark not doing the work.

Force the value into a required result and inspect the optimized plan.

Never benchmark your intention. Benchmark the executed plan.

Compact Checklist

□ What physical plan actually ran? □ Did AQE change it? □ Which stage dominates runtime? □ CPU, shuffle, I/O, spill, GC, or skew? □ Is the join strategy sensible? □ Are unnecessary columns surviving? □ Is important logic hidden from Catalyst? □ Where does WholeStageCodegen begin/end? □ Is generated code huge or falling back? □ Am I comparing warm runs? □ Did Catalyst optimize away my benchmark?

Notice what is missing: randomly raise a codegen threshold.

🧪 Lab 11 — Reconstruct the Crime

11__reconstruct_the_crime.ipynb collects the optimized/executed plans, codegen output, and stage metrics for one intentionally imperfect query before changing anything. It also fixes a benchmark whose projection can be pruned.

Finish with: the strongest suspect is ___; the evidence is ___; therefore the first thing worth testing is ___.

Everything before the evidence is fan fiction.

🍎 The Plan Was Never Just A Plan

We started with *(1), Exchange, and a few suspicious names. We ended with a compiler pipeline: Catalyst analyzes and optimizes the query, the physical planner chooses operators, WholeStageCodegen turns compatible row-based fragments into Java, Janino compiles that Java, and the JVM may optimize hot bytecode again.

The lesson is not “codegen is fast” or “UDFs are bad.”
It’s simpler: know which layer owns the cost before you try to fix it.

Read the plan. Check the metrics. Inspect generated code when the evidence points there.

Once you know the names, the physical plan stops looking supernatural. It becomes evidence.

📓 A Name Worth Writing In My Thankful Note

This article probably wouldn’t exist without a fascinating conversation I had with Sem Sinchenko about his work on GraphFrames. That discussion sent me digging into FiniteAXPlusB, CodegenFallback, and eventually the much larger WholeStageCodegen rabbit hole.

So, Sem: thank you for the spark. The rest of the investigation — and the suspicious amount of generated Java — followed naturally.

Read the plan, trust the metrics, and don’t let Ryuk benchmark the first run.