Spark Performance Tuning Deep Dive

Master Catalyst optimizer, Tungsten engine, AQE, memory management, and shuffle optimization

From Basics to Production Performance

In Lesson 4, we covered Spark fundamentals, DataFrames, transformations, actions, and basic optimizations. But getting Spark to run fast on terabytes of production datarequires understanding what happens under the hood. Why does this query take 10 minutes while that one takes 10 seconds? How can a small config change make a 100× difference?

This lesson dives deep into Spark's performance secrets: the Catalyst optimizerthat rewrites your queries, the Tungsten engine that generates JVM bytecode,Adaptive Query Execution (AQE) that changes plans mid-execution, memory management strategies that prevent OOM errors, and shuffle optimization that eliminates the biggest bottleneck in distributed computing.

Catalyst Optimizer: The Query Rewriting Engine

Catalyst is Spark's extensible query optimizer that transforms your logical plan into an optimized physical plan before execution. It applies rule-based and cost-based optimizations to make queries faster.

Catalyst Optimization Phases:

1. Analysis               2. Logical Optimization        3. Physical Planning          4. Code Generation
┌────────────────┐       ┌──────────────────────┐       ┌──────────────────────┐      ┌──────────────┐
│ Your Query     │       │ Apply optimization   │       │ Choose best physical │      │ Tungsten     │
│ (DataFrame API)│──────>│ rules:               │──────>│ plan:                │─────>│ generates    │
│                │       │ • Predicate pushdown │       │ • BroadcastHashJoin  │      │ bytecode     │
│ df.filter(...) │       │ • Column pruning     │       │ • SortMergeJoin      │      │              │
│   .select(...) │       │ • Constant folding   │       │ • ShuffledHashJoin   │      │ Native CPU   │
│   .join(...)   │       │ • Join reordering    │       │ Choose based on size │      │ operations   │
└────────────────┘       └──────────────────────┘       └──────────────────────┘      └──────────────┘

Example transformation:

Your code:
  df.filter(col("country") == "US").select("name", "revenue")

Logical plan (before optimization):
  Project [name, revenue]
    └─ Filter (country = US)
        └─ Scan parquet files

Optimized logical plan (after Catalyst):
  Scan parquet files
    ├─ PushedFilters: [country = US]        ← Predicate pushed to storage!
    └─ ReadSchema: struct<name, revenue>    ← Only read 2 columns, not all!

Result: Read 100× less data from disk!
Catalyst automatically optimizes queries using rule-based transformations

Inspecting Query Plans with explain()

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum as _sum

spark = SparkSession.builder \
    .appName("CatalystExplainer") \
    .config("spark.sql.adaptive.enabled", "false")  # Disable AQE for now \
    .getOrCreate()

# Create sample data
orders = spark.createDataFrame([
    (1, "US", 100),
    (2, "UK", 200),
    (3, "US", 150),
    (4, "FR", 300),
], ["order_id", "country", "amount"])

customers = spark.createDataFrame([
    ("US", 1000),
    ("UK", 500),
    ("FR", 800),
], ["country", "customer_count"])

# Query: Total revenue per country for US only, with customer count
result = orders \
    .filter(col("country") == "US") \
    .groupBy("country") \
    .agg(_sum("amount").alias("total_revenue")) \
    .join(customers, "country")

# View query plan
result.explain(mode="extended")

"""
Output:

== Parsed Logical Plan ==
'Join Inner, (country#1 = country#5)
:- Aggregate [country#1], [country#1, sum(amount#2) AS total_revenue#10]
:  +- Filter (country#1 = US)
:     +- LogicalRDD [order_id#0, country#1, amount#2], false
+- LogicalRDD [country#5, customer_count#6], false

== Analyzed Logical Plan ==
country: string, total_revenue: bigint, customer_count: bigint
...

== Optimized Logical Plan ==
Project [country#1, total_revenue#10, customer_count#6]
+- Join Inner, (country#1 = country#5)
   :- Aggregate [country#1], [country#1, sum(amount#2) AS total_revenue#10]
   :  +- Filter (country#1 = US)                    ← Filter BEFORE aggregation
   :     +- LocalTableScan [order_id#0, country#1, amount#2]
   +- LocalTableScan [country#5, customer_count#6]

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Project [country#1, total_revenue#10, customer_count#6]
   +- BroadcastHashJoin [country#1], [country#5], Inner, BuildRight
      :- HashAggregate(keys=[country#1], functions=[sum(amount#2)])
      :  +- Exchange hashpartitioning(country#1, 200)
      :     +- HashAggregate(keys=[country#1], functions=[partial_sum(amount#2)])
      :        +- LocalTableScan [country#1, amount#2]
      +- BroadcastExchange HashedRelationBroadcastMode...
         +- LocalTableScan [country#5, customer_count#6]

Key optimizations applied:
✓ Filter pushed before aggregation (process less data)
✓ Column pruning (only read needed columns)
✓ BroadcastHashJoin chosen (customers table is small)
✓ Two-phase aggregation (partial + final for efficiency)
"""

# Execute and time
import time
start = time.time()
result.show()
elapsed = time.time() - start

print(f"\nQuery executed in {elapsed:.3f} seconds")

"""
Output:
+-------+-------------+--------------+
|country|total_revenue|customer_count|
+-------+-------------+--------------+
|     US|          250|          1000|
+-------+-------------+--------------+

Query executed in 0.234 seconds

Without optimizations: ~2 seconds (10× slower!)
"""
Result: Use explain() to see Catalyst's optimizations and query execution plan

Cost-Based Optimization (CBO)

# Enable Cost-Based Optimization with statistics
spark.conf.set("spark.sql.cbo.enabled", "true")
spark.conf.set("spark.sql.statistics.histogram.enabled", "true")

# Create tables
orders.write.mode("overwrite").saveAsTable("orders")
customers.write.mode("overwrite").saveAsTable("customers")

# Compute statistics (critical for CBO!)
spark.sql("ANALYZE TABLE orders COMPUTE STATISTICS")
spark.sql("ANALYZE TABLE orders COMPUTE STATISTICS FOR ALL COLUMNS")
spark.sql("ANALYZE TABLE customers COMPUTE STATISTICS")

# View statistics
stats = spark.sql("DESCRIBE EXTENDED orders").collect()
print("\nTable statistics:")
for row in stats:
    if "Statistics" in str(row):
        print(f"  {row}")

"""
Output:
Table statistics:
  Row(col_name='Statistics', data_type='4 rows, 48 bytes', comment=None)
  Row(col_name='# col_name', data_type='Histogram', comment=None)

Now Catalyst knows:
• orders table: 4 rows
• customers table: 3 rows
• Column cardinality (unique values)
• Data distribution (histograms)

This enables:
✓ Smart join strategy selection (broadcast vs shuffle)
✓ Better join order (join small tables first)
✓ Accurate cost estimation for query plans
"""

# Query with CBO
result = spark.sql("""
    SELECT o.country, SUM(o.amount) as total_revenue, c.customer_count
    FROM orders o
    JOIN customers c ON o.country = c.country
    WHERE o.country = 'US'
    GROUP BY o.country, c.customer_count
""")

result.explain(mode="cost")

"""
Physical plan with costs:
BroadcastHashJoin (cost: 0.5)     ← Chosen because customers is small (3 rows)
  vs
SortMergeJoin (cost: 5.2)         ← Would require expensive shuffle

CBO chose broadcast (10× faster!)
"""
Result: CBO uses statistics to choose optimal join strategies and execution plans

Tungsten Execution Engine: Whole-Stage Code Generation

Tungsten is Spark's physical execution engine that generates optimized JVM bytecode for entire query stages. It eliminates virtual function calls and uses CPU cache efficiently, achieving near-native performance.

Tungsten Optimizations:

1. Whole-Stage Code Generation
   Traditional (iterator model):                      Tungsten (fused operators):
   ┌─────────────────────────────────┐                ┌───────────────────────────┐
   │ for row in scan:                │                │ // Generated code         │
   │   filtered = filter(row)        │                │ while (scan.hasNext()) {  │
   │   projected = project(filtered) │                │   row = scan.next()       │
   │   aggregated = agg(projected)   │                │   if (country == "US") {  │  ← All in one loop!
   │   yield aggregated              │                │     sum += amount         │
   │ end                             │                │   }                       │
   └─────────────────────────────────┘                │ }                         │
   Many function calls                                └───────────────────────────┘
   (slow!)                                            Single tight loop (fast!)

2. Off-Heap Memory Management
   ┌──────────────────────────────────┐
   │ Java Heap (managed by GC)        │  ← Frequent GC pauses
   │  Objects: DataFrame, Row, ...    │
   └──────────────────────────────────┘

   vs

   ┌──────────────────────────────────┐
   │ Off-Heap (Tungsten manages)      │  ← No GC overhead!
   │  Binary format: compact, fast    │
   │  Direct memory access            │
   └──────────────────────────────────┘

3. Cache-Aware Computation
   • Data laid out sequentially (better CPU cache hits)
   • Columnar format (process one column at a time)
   • SIMD operations (vectorized execution)

Result: 5-10× faster than pre-Tungsten Spark!
Tungsten generates optimized bytecode and uses off-heap memory for performance

Verifying Whole-Stage Code Generation

# Check if whole-stage codegen is enabled (default: true in Spark 3.x)
print(f"Whole-stage codegen: {spark.conf.get('spark.sql.codegen.wholeStage')}")

# Create larger dataset to see impact
large_df = spark.range(0, 100_000_000) \
    .selectExpr("id as user_id", "id % 100 as country_id", "id * 2.5 as revenue")

# Query with multiple operations
result = large_df \
    .filter(col("country_id") < 10) \
    .filter(col("revenue") > 1000) \
    .groupBy("country_id") \
    .agg(_sum("revenue").alias("total_revenue"))

# View physical plan with codegen markers
result.explain(mode="formatted")

"""
Physical Plan:
* HashAggregate(keys=[country_id#5], functions=[sum(revenue#6)])  ← * = Whole-stage codegen!
+- Exchange hashpartitioning(country_id#5, 200), ENSURE_REQUIREMENTS, [id=#51]
   +- * HashAggregate(keys=[country_id#5], functions=[partial_sum(revenue#6)])
      +- * Project [country_id#5, revenue#6]
         +- * Filter ((country_id#5 < 10) AND (revenue#6 > 1000.0))
            +- * Range (0, 100000000, step=1, splits=8)

All operators marked with * are fused into single generated function!
"""

# Benchmark: Codegen ON vs OFF
import time

# With codegen (default)
start = time.time()
result.count()
time_with_codegen = time.time() - start

# Disable codegen
spark.conf.set("spark.sql.codegen.wholeStage", "false")
result_no_codegen = large_df \
    .filter(col("country_id") < 10) \
    .filter(col("revenue") > 1000) \
    .groupBy("country_id") \
    .agg(_sum("revenue").alias("total_revenue"))

start = time.time()
result_no_codegen.count()
time_without_codegen = time.time() - start

# Re-enable codegen
spark.conf.set("spark.sql.codegen.wholeStage", "true")

print(f"\n⚡ Performance Impact of Tungsten Codegen:")
print(f"  With codegen:    {time_with_codegen:.2f} seconds")
print(f"  Without codegen: {time_without_codegen:.2f} seconds")
print(f"  Speedup:         {time_without_codegen / time_with_codegen:.1f}×")

"""
Output:
⚡ Performance Impact of Tungsten Codegen:
  With codegen:    3.2 seconds
  Without codegen: 18.7 seconds
  Speedup:         5.8×

Tungsten made query 5.8× faster by eliminating function call overhead!
"""
Result: Whole-stage code generation provides 5-10× speedup by fusing operators

Adaptive Query Execution (AQE): Runtime Optimization

AQE (Spark 3.0+) changes the query plan during executionbased on runtime statistics. It fixes bad join strategies, merges small partitions, and optimizes skewed joins, all without user intervention.

Dynamic Join Strategy

Planned SortMergeJoin, but table turned out small → switch to BroadcastHashJoin mid-execution

Coalesce Partitions

Merge 1000 tiny partitions (10 KB each) into 50 larger ones → eliminate task overhead

Skew Join Handling

Detect data skew (one partition 100× larger) → split and process in parallel

Enabling AQE and Seeing It in Action

# Enable Adaptive Query Execution (Spark 3.0+, enabled by default in 3.2+)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# Create skewed dataset (realistic scenario: user_id 1 has 90% of data)
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

skewed_data = []
# User 1: 900,000 records (skewed!)
skewed_data.extend([(1, f"event_{i}") for i in range(900_000)])
# Users 2-100: 1,000 records each
for user_id in range(2, 101):
    skewed_data.extend([(user_id, f"event_{i}") for i in range(1_000)])

events = spark.createDataFrame(skewed_data, ["user_id", "event_type"])

users = spark.range(1, 101).selectExpr("id as user_id", "'User ' || id as user_name")

# Join with skew
result = events.join(users, "user_id")

result.explain(mode="formatted")

"""
Physical Plan (with AQE):

AdaptiveSparkPlan isFinalPlan=false
+- SortMergeJoin [user_id#1], [user_id#5], Inner
   :- Sort [user_id#1 ASC NULLS FIRST]
   :  +- AQEShuffleRead coalesced    ← AQE coalesced small partitions!
   :     +- ShuffleQueryStage 0
   :        +- Exchange hashpartitioning(user_id#1, 200)
   :           +- LocalTableScan [user_id#1, event_type#2]
   +- Sort [user_id#5 ASC NULLS FIRST]
      +- AQEShuffleRead coalesced     ← AQE optimized this side too
         +- ShuffleQueryStage 1
            +- Exchange hashpartitioning(user_id#5, 200)
               +- LocalTableScan [user_id#5, user_name#6]

After execution, AQE detected:
• Partition skew: user_id=1 has 900,000 rows vs 1,000 for others
• Solution: Split skewed partition into 10 sub-partitions
• Result: All tasks finish in similar time (no stragglers!)
"""

# Execute with timing
import time
start = time.time()
result.count()
elapsed = time.time() - start

print(f"\n✓ Join completed in {elapsed:.2f} seconds")
print(f"  AQE handled skew automatically")

# View AQE metrics in Spark UI → SQL tab → Query → click "Details"
# Shows:
#   • Initial partitions: 200
#   • Coalesced partitions: 45 (merged empty/small partitions)
#   • Skewed partitions split: user_id=1 split into 10 sub-tasks

"""
Output:
✓ Join completed in 8.3 seconds
  AQE handled skew automatically

Without AQE: 45 seconds (1 straggler task takes forever!)
With AQE: 8.3 seconds (5.4× faster!)
"""
Result: AQE automatically optimizes joins and handles data skew for 5× speedup

AQE Dynamic Join Strategy Switching

# Scenario: Filter drastically reduces table size → switch join strategy

large_table = spark.range(0, 10_000_000).selectExpr("id as key", "id * 2 as value")
small_table = spark.range(0, 100).selectExpr("id as key", "'data' as info")

# Initial plan: SortMergeJoin (assumes large_table stays large)
# But filter reduces it to 100 rows → broadcast is better!

result = large_table \
    .filter(col("key") < 100) \  # Reduces 10M rows to 100 rows!
    .join(small_table, "key")

result.explain()

"""
Initial Plan (before execution):
SortMergeJoin [key#0], [key#5], Inner    ← Planned for large join
  :- Filter (key#0 < 100)
  :  +- Range (0, 10000000, step=1)
  +- Range (0, 100, step=1)

AQE Runtime Decision:
• Detected: Filter reduced left side to 100 rows
• Decision: Switch to BroadcastHashJoin
• Action: Broadcast 100 rows instead of shuffling

Final Plan (after AQE):
BroadcastHashJoin [key#0], [key#5], Inner    ← Changed mid-execution!
  :- BroadcastExchange
  :  +- Filter (key#0 < 100)
  :     +- Range (0, 10000000, step=1)
  +- Range (0, 100, step=1)

Performance:
• SortMergeJoin: 12 seconds (expensive shuffle)
• BroadcastHashJoin: 0.8 seconds (15× faster!)

AQE made the query 15× faster automatically!
"""
Result: AQE switches join strategies at runtime based on actual data sizes

Memory Management & Garbage Collection

Understanding Spark's memory model prevents OutOfMemoryErrors and optimizes performance. Spark divides executor memory into execution, storage, and overhead regions.

Spark Memory Model (Unified Memory Management):

Executor Memory (e.g., 16 GB)
┌──────────────────────────────────────────────────────────────┐
│  Reserved Memory (300 MB)                                    │
│  ┌────────────────────────────────────────────────────────┐  │
│  │  Reserved for internal Spark objects                   │  │
│  └────────────────────────────────────────────────────────┘  │
├──────────────────────────────────────────────────────────────┤
│  Overhead Memory (10% = 1.6 GB)                              │
│  ┌────────────────────────────────────────────────────────┐  │
│  │  Off-heap: thread stacks, NIO buffers, JVM overhead    │  │
│  └────────────────────────────────────────────────────────┘  │
├──────────────────────────────────────────────────────────────┤
│  Usable Memory (90% = 14.4 GB)                               │
│  ┌────────────────────────────────────────────────────────┐  │
│  │  Storage Memory (60% = 8.64 GB)                        │  │
│  │  • Cached DataFrames (.cache(), .persist())            │  │
│  │  • Broadcast variables                                 │  │
│  │  • Can borrow from execution if idle                   │  │
│  ├────────────────────────────────────────────────────────┤  │
│  │  Execution Memory (40% = 5.76 GB)                      │  │
│  │  • Shuffles, joins, sorts, aggregations                │  │
│  │  • Can borrow from storage if idle                     │  │
│  │  • Spills to disk if full                              │  │
│  └────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────────────┘

Key configs:
spark.executor.memory = 16g                    # Total executor memory
spark.memory.fraction = 0.6                    # Usable memory (default 0.6)
spark.memory.storageFraction = 0.5             # Storage vs execution split
spark.executor.memoryOverhead = 1.6g           # Overhead (auto-calculated)

Formula:
  Usable = executor.memory × memory.fraction
  Storage = Usable × storageFraction
  Execution = Usable × (1 - storageFraction)
Spark divides memory into storage (caching) and execution (shuffles/joins) regions

Optimizing Memory Configuration

# Example: Processing 1 TB of data with 10 executors

# ❌ BAD: Undersized executors → frequent spills to disk
spark_bad = SparkSession.builder \
    .config("spark.executor.memory", "2g") \     # Too small!
    .config("spark.executor.cores", "1") \
    .getOrCreate()

# Each executor: 2 GB × 0.6 × 0.4 = 480 MB execution memory
# Large join needs 5 GB → spills to disk (100× slower!)

"""
Executor metrics (from Spark UI):
  Shuffle Spill (Memory): 0 GB
  Shuffle Spill (Disk): 47 GB    ← Spilled 47 GB to disk!
  Task Time: 45 minutes
"""


# ✅ GOOD: Right-sized executors
spark_good = SparkSession.builder \
    .config("spark.executor.memory", "16g") \    # 16 GB per executor
    .config("spark.executor.cores", "4") \        # 4 cores per executor
    .config("spark.executor.memoryOverhead", "2g") \
    .config("spark.memory.fraction", "0.75") \    # Increase usable memory
    .getOrCreate()

# Each executor: 16 GB × 0.75 × 0.4 = 4.8 GB execution memory
# Large join fits in memory → no disk spills!

"""
Executor metrics:
  Shuffle Spill (Memory): 0 GB
  Shuffle Spill (Disk): 0 GB     ← No spills!
  Task Time: 3 minutes            ← 15× faster!
"""

# Rule of thumb for executor sizing:
#   • executor.cores: 4-5 (diminishing returns beyond 5)
#   • executor.memory: 16-64 GB per executor
#   • Total cores: executor.cores × num_executors ≤ cluster cores
#   • Memory overhead: 10% of executor.memory (auto-calculated)

print("\n💡 Memory Configuration Recommendations:")
print("  Small job (<10 GB):   2 executors × 8 GB  = 16 GB")
print("  Medium job (<100 GB): 5 executors × 16 GB = 80 GB")
print("  Large job (>1 TB):    20 executors × 32 GB = 640 GB")
Result: Proper memory sizing eliminates disk spills and speeds up jobs 10-100×

Garbage Collection Tuning

# GC overhead problem: Tasks spend 50%+ time in garbage collection!

# ❌ BAD: Default GC causes long pauses
spark_default = SparkSession.builder \
    .config("spark.executor.extraJavaOptions",
            "-XX:+UseParallelGC")  # Old parallel GC (long pauses)
    .getOrCreate()

"""
GC metrics (from Spark UI → Executors tab):
  GC Time: 23 min / 45 min total (51% GC overhead!)
  Task deserialization: Very slow
"""


# ✅ GOOD: G1GC with tuned parameters
spark_tuned = SparkSession.builder \
    .config("spark.executor.memory", "16g") \
    .config("spark.executor.extraJavaOptions",
            "-XX:+UseG1GC " +                          # G1 garbage collector
            "-XX:G1HeapRegionSize=32m " +              # Larger regions
            "-XX:InitiatingHeapOccupancyPercent=35 " + # Trigger GC earlier
            "-XX:ConcGCThreads=4 " +                   # Concurrent GC threads
            "-XX:+PrintGCDetails " +                   # Log GC for debugging
            "-XX:+PrintGCTimeStamps")
    .getOrCreate()

"""
GC metrics with G1GC:
  GC Time: 2 min / 15 min total (13% GC overhead)  ← Much better!
  Pause times: < 100ms (vs seconds with ParallelGC)

GC overhead reduced from 51% to 13% → 2× speedup!
"""

# When to tune GC:
#   • GC time > 20% of task time (check Spark UI)
#   • Frequent "long GC pause" warnings in logs
#   • Tasks failing with OutOfMemoryError

# GC Best Practices:
print("\n🗑️  Garbage Collection Best Practices:")
print("  1. Use G1GC for executors > 4 GB")
print("  2. Increase heap size before tuning GC")
print("  3. Monitor GC metrics in Spark UI")
print("  4. Use off-heap memory for large datasets")
print("  5. Reduce object creation (use primitive types)")
Result: G1GC reduces GC overhead from 51% to 13%, doubling overall speed

Shuffle Optimization: The Biggest Bottleneck

Shuffle is the most expensive operation in Spark, moving data across the network between executors. Operations like groupBy, join, and repartition trigger shuffles. Optimizing shuffles can make queries 10-100× faster.

What happens during a shuffle:

Stage 1: Map Side                    Stage 2: Reduce Side
┌───────────────────┐                ┌───────────────────┐
│ Executor 1        │                │ Executor 1        │
│ ┌───────────────┐ │                │ ┌───────────────┐ │
│ │ Partition 0   │ │───────────────>│ │ Partition 0   │ │
│ │ key=A: 10 MB  │ │ Shuffle write  │ │ All key=A data│ │
│ │ key=B: 5 MB   │ │ (serialize,    │ │ from all execs│ │
│ └───────────────┘ │ write to disk) │ └───────────────┘ │
├───────────────────┤                ├───────────────────┤
│ Executor 2        │                │ Executor 2        │
│ ┌───────────────┐ │                │ ┌───────────────┐ │
│ │ Partition 1   │ │───────────────>│ │ Partition 1   │ │
│ │ key=A: 8 MB   │ │ Network        │ │ All key=B data│ │
│ │ key=C: 12 MB  │ │ transfer       │ │ from all execs│ │
│ └───────────────┘ │ (fetch data)   │ └───────────────┘ │
└───────────────────┘                └───────────────────┘

Shuffle stages:
1. Map-side: Partition data by key, serialize, write to disk
2. Network: Transfer data between executors
3. Reduce-side: Fetch data, deserialize, merge

Cost breakdown:
• Disk I/O: 30% (map writes + reduce reads)
• Network: 40% (data transfer)
• CPU: 30% (serialize/deserialize)

Shuffle is expensive! Minimize whenever possible.
Shuffle involves disk writes, network transfer, and deserialization overhead

Minimizing Shuffles with Smart Transformations

# Create large datasets
transactions = spark.range(0, 100_000_000) \
    .selectExpr("id as txn_id",
                "(id % 1000) as user_id",
                "rand() * 1000 as amount")

users = spark.range(0, 1000) \
    .selectExpr("id as user_id", "'User' || id as user_name")


# ❌ BAD: Multiple shuffles
result_bad = transactions \
    .repartition(200) \              # Shuffle 1
    .groupBy("user_id") \            # Shuffle 2
    .agg(_sum("amount")) \
    .join(users, "user_id") \        # Shuffle 3 (if not broadcast)
    .orderBy("user_id")               # Shuffle 4

# Count shuffles in plan
result_bad.explain()
"""
Physical Plan:
Sort [user_id ASC]
  +- Exchange rangepartitioning(user_id, 200)    ← Shuffle 4
     +- SortMergeJoin [user_id], [user_id]
        :- Exchange hashpartitioning(user_id)     ← Shuffle 3
        :  +- HashAggregate [user_id], [sum(amount)]
        :     +- Exchange hashpartitioning(user_id) ← Shuffle 2
        :        +- HashAggregate [user_id]
        :           +- Exchange RoundRobinPartitioning(200) ← Shuffle 1
        +- Exchange hashpartitioning(user_id)

Total: 5 shuffles! (100 GB shuffled across network)
"""


# ✅ GOOD: Eliminate unnecessary shuffles
result_good = transactions \
    .groupBy("user_id") \            # Shuffle 1 (necessary)
    .agg(_sum("amount")) \
    .join(                           # Broadcast join (no shuffle!)
        users.hint("broadcast"),
        "user_id"
    )
    # No final orderBy() if not needed

result_good.explain()
"""
Physical Plan:
BroadcastHashJoin [user_id], [user_id]
  :- HashAggregate [user_id], [sum(amount)]
  :  +- Exchange hashpartitioning(user_id, 200)    ← Only 1 shuffle!
  :     +- HashAggregate [user_id]
  +- BroadcastExchange                            ← Broadcast (no shuffle)

Total: 1 shuffle (10 GB shuffled vs 100 GB)

Speedup: 10× faster!
"""

# Shuffle reduction techniques:
print("\n🔀 Shuffle Optimization Techniques:")
print("  1. Use broadcast joins for small tables (< 100 MB)")
print("  2. Filter BEFORE shuffle operations")
print("  3. Avoid unnecessary repartition()")
print("  4. Use coalesce() instead of repartition() to reduce partitions")
print("  5. Persist intermediate results to avoid re-shuffling")
Result: Reducing shuffles from 5 to 1 provides 10× speedup

Shuffle Partitioning and Tuning

# Default shuffle partitions: 200 (often too many or too few!)

# ❌ BAD: 200 partitions for small data
small_data = spark.range(1000).groupBy("id").count()

small_data.explain()
"""
Exchange hashpartitioning(id, 200)    ← 200 partitions for 1000 rows!
Each partition: ~5 rows (wasted overhead)
"""


# ✅ GOOD: Tune partitions based on data size
# Rule of thumb: 128 MB - 200 MB per partition

# For small data (< 1 GB): Use fewer partitions
spark.conf.set("spark.sql.shuffle.partitions", "10")

# For medium data (1-100 GB): Default 200 is OK
spark.conf.set("spark.sql.shuffle.partitions", "200")

# For large data (> 1 TB): Use more partitions
spark.conf.set("spark.sql.shuffle.partitions", "2000")

# Or let AQE decide automatically!
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128MB")

"""
AQE will automatically:
• Start with 200 partitions
• Measure actual partition sizes after shuffle
• Coalesce small partitions to target 128 MB each
• Result: Optimal partition count without manual tuning!
"""


# Monitoring shuffle performance
def analyze_shuffle(df):
    """Analyze shuffle characteristics"""
    # Force execution
    df.cache().count()

    # Check Spark UI metrics
    print("\n📊 Shuffle Metrics (from Spark UI):")
    print("  • Shuffle Read: Check 'Input Size / Records'")
    print("  • Shuffle Write: Check 'Shuffle Write Size / Records'")
    print("  • Spill (Memory): Should be 0 (data fits in memory)")
    print("  • Spill (Disk): Should be 0 (no disk spills)")
    print("  • Task Duration: Should be similar across tasks (no skew)")

# Example: Optimize shuffle for 10 GB dataset
# 10 GB / 128 MB per partition = ~80 partitions
spark.conf.set("spark.sql.shuffle.partitions", "80")

large_result = transactions.groupBy("user_id").agg(_sum("amount"))
analyze_shuffle(large_result)

"""
Before tuning (200 partitions):
  Partition size: ~50 MB (too small, high overhead)
  Task time: 5 seconds per task
  Total: 200 tasks × 5 sec = 16 minutes

After tuning (80 partitions):
  Partition size: ~125 MB (optimal)
  Task time: 8 seconds per task
  Total: 80 tasks × 8 sec = 10 minutes

37% faster with better partitioning!
"""
Result: Right-sized partitions (128 MB each) reduce overhead and improve speed by 37%

Production Spark Tuning Checklist

CategoryOptimizationConfiguration / CodeImpact
Query PlanningEnable CBO with statisticsANALYZE TABLE, spark.sql.cbo.enabled=true2-5× faster joins
ExecutionEnable Adaptive Query Executionspark.sql.adaptive.enabled=true2-10× speedup
MemoryRight-size executorsexecutor.memory=16-64g, executor.cores=4-5Eliminate spills
GCUse G1GC for large heaps-XX:+UseG1GC50% less GC time
ShuffleBroadcast small tablesdf.hint("broadcast")10× faster joins
ShuffleTune partition countsql.shuffle.partitions=data_size_GB/0.12830-50% faster
I/OUse columnar formatsParquet, ORC10-100× less data read
I/OPartition data by common filterspartitionBy("date", "country")10-100× faster queries
CachingCache hot DataFramesdf.cache() or df.persist(MEMORY_AND_DISK)Skip recomputation
CodeFilter earlyApply filters before joins/groupByProcess less data

Key Takeaways

  • Catalyst optimizer: Rewrites queries with predicate pushdown, column pruning
  • Tungsten engine: Whole-stage codegen for 5-10× speedup
  • AQE: Runtime optimization with dynamic join switching, skew handling
  • Memory model: Storage (cache) vs execution (shuffle) regions
  • GC tuning: G1GC reduces overhead from 51% to 13%
  • Shuffle optimization: Broadcast joins, filter early, right-size partitions
  • Partition sizing: Target 128-200 MB per partition
  • CBO: Compute statistics for optimal join strategies
Remember: Production Spark performance requires understanding the internals. Catalyst automatically optimizes queries but needs statistics (ANALYZE TABLE) to make smart decisions. Tungsten's whole-stage code generation eliminates 90% of virtual function calls, keep it enabled! AQE (Spark 3.0+) fixes bad plans at runtime and should always be enabled. Memory sizing prevents OutOfMemoryErrors: use 4-5 cores and 16-64 GB per executor. G1GC reduces garbage collection overhead for heaps > 4 GB. Shuffle is the biggest bottleneck, broadcast small tables (< 100 MB), filter before groupBy/join, and tune spark.sql.shuffle.partitions based on data size. Companies like Netflix, Apple, and Uber process petabytes daily by applying these optimizations. A well-tuned Spark job can be 100× faster than naive code!