Whole Stage Code Generation

Overview

Learn how Whole-Stage Code Generation collapses multiple physical execution steps into a single Java function to eliminate CPU virtual call overhead.

What You Will Learn

Detailed Concept Explanation

In older Spark versions, processing a row meant passing it through a chain of objects (e.g. Scan -> Filter -> Project), where each step called a virtual method (next()) to retrieve the next row. This created massive CPU virtual call overhead.

Whole-Stage Code Generation solves this by collapsing the physical steps of a query stage into a single Java function at runtime.

Key Benefits

• CPU register locality: Intermediate data is kept in CPU registers instead of being written to JVM memory arrays.
• Janino Compiler: Spark compiles this collapsed Java function directly into bytecode at runtime, running close to hand-written code speeds.

Code Examples

Python (PySpark) Implementation

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("Codegen").getOrCreate()

# Simple filter and projection query
df = spark.range(1, 1000)
result = df.filter(df.id > 10).selectExpr("id + 5")

# Codegen stages are marked by star symbols (*) in explain logs
result.explain()

Expected Output

== Physical Plan ==
*(1) Project [(id#0L + 5) AS (id + 5)#2L]
+- *(1) Filter (id#0L > 10)
   +- *(1) Range (1, 1000, step=1, splits=2)

Execution Plan Diagram (Python & Scala)

Scala Implementation

import org.apache.spark.sql.SparkSession

val spark = SparkSession.builder().appName("CodegenScala").getOrCreate()

val df = spark.range(1, 1000)
val result = df.filter($"id" > 10).selectExpr("id + 5")

result.explain()

Related Topics

• Tungsten Engine
• Logical Plan
• Physical Plan