Performance Tuning & Optimization

from pyspark.sql import functions as F
from pyspark.sql.functions import broadcast

# BAD: Shuffle join (both tables are large or Spark doesn't auto-broadcast)
result_slow = sales_df.join(products_df, "product_id")

# GOOD: Force broadcast join when you know the table is small
result_fast = sales_df.join(broadcast(products_df), "product_id")

# Configure broadcast threshold (default 10MB)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "100MB")

# Check the query plan to verify broadcast join is used
result_fast.explain(True)
# Look for "BroadcastHashJoin" instead of "SortMergeJoin"

# Identify skewed keys
skew_analysis = sales_df.groupBy("customer_id") \
    .count() \
    .orderBy(F.desc("count")) \
    .limit(20)
skew_analysis.show()

# Technique 1: Salt the skewed key
salt_range = 10
salted_sales = sales_df.withColumn(
    "salt", (F.rand() * salt_range).cast("int")
).withColumn(
    "salted_key",
    F.concat(F.col("customer_id"), F.lit("_"), F.col("salt"))
)

# Technique 2: Enable AQE skew join optimization (Databricks)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

# Check current partition count
print(f"Current partitions: {df.rdd.getNumPartitions()}")

# REPARTITION: Full shuffle, creates equal-sized partitions
# Use when you need to INCREASE partitions or partition by a column
df_repartitioned = df.repartition(200)
df_by_date = df.repartition("date")
df_by_date_200 = df.repartition(200, "date")

# COALESCE: No shuffle, merges existing partitions
# Use when you need to DECREASE partitions (e.g., before writing)
df_coalesced = df.coalesce(10)

# Common pattern: process with many partitions, coalesce before write
df.repartition(500) \
    .filter(F.col("status") == "active") \
    .groupBy("region").agg(F.sum("revenue")) \
    .coalesce(1) \
    .write.mode("overwrite").parquet("/output/regional_revenue")

# Delta Lake auto-optimize handles file sizing automatically
spark.conf.set("spark.databricks.delta.optimizeWrite.enabled", "true")
spark.conf.set("spark.databricks.delta.autoCompact.enabled", "true")

-- Run OPTIMIZE with Z-ORDER on frequently filtered columns
OPTIMIZE catalog.schema.sales
ZORDER BY (region, product_category);

-- Z-ORDER is most effective with:
-- 1. High-cardinality columns used in WHERE clauses
-- 2. Columns frequently used in JOIN conditions
-- 3. Up to 3-4 columns (diminishing returns after)

-- Check file-level stats to verify data skipping
DESCRIBE DETAIL catalog.schema.sales;

-- Schedule regular OPTIMIZE jobs
OPTIMIZE catalog.schema.sales
WHERE date >= current_date() - INTERVAL 7 DAYS
ZORDER BY (customer_id, product_id);

from pyspark import StorageLevel

# Cache in memory (default)
df.cache()  # equivalent to persist(StorageLevel.MEMORY_AND_DISK)

# Persist with specific storage level
df.persist(StorageLevel.MEMORY_ONLY)         # fastest, may lose partitions
df.persist(StorageLevel.MEMORY_AND_DISK)     # spills to disk if needed
df.persist(StorageLevel.DISK_ONLY)           # large datasets, memory pressure

# IMPORTANT: cache is lazy -- trigger with an action
df.cache()
df.count()  # materializes the cache

# Unpersist when done to free memory
df.unpersist()

# Delta Cache (Databricks-specific, automatic)
# Automatically caches frequently-read Parquet files on local SSD
spark.conf.set("spark.databricks.io.cache.enabled", "true")
spark.conf.set("spark.databricks.io.cache.maxDiskUsage", "50g")
spark.conf.set("spark.databricks.io.cache.maxMetaDataCache", "1g")

-- Check if Photon is enabled on the current cluster
SET spark.databricks.photon.enabled;

-- Enable Photon (if supported by cluster type)
SET spark.databricks.photon.enabled = true;

-- Run a query and check the physical plan for Photon operators
EXPLAIN SELECT region, SUM(revenue)
FROM catalog.schema.sales
WHERE date >= '2025-01-01'
GROUP BY region;
-- Look for "PhotonGroupingAgg" and "PhotonScan" in the plan

-- Photon works best with Delta Lake tables
CREATE TABLE catalog.schema.sales_delta
USING DELTA
AS SELECT * FROM catalog.schema.sales_parquet;

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

# Coalesce partitions: merge small partitions after shuffle
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "64MB")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128MB")

# Skew join: automatically handle data skew
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

# With AQE, you can set shuffle.partitions high and let AQE coalesce
spark.conf.set("spark.sql.shuffle.partitions", "auto")

# Verify AQE is active in the query plan
df = spark.sql("""
    SELECT region, SUM(revenue) as total
    FROM catalog.schema.sales
    GROUP BY region
""")
df.explain("formatted")
# Look for "AdaptiveSparkPlan isFinalPlan=true"

# Shuffle and Partitions
spark.sql.shuffle.partitions = auto
spark.sql.adaptive.enabled = true
spark.sql.autoBroadcastJoinThreshold = 100MB

# Delta Lake Optimization
spark.databricks.delta.optimizeWrite.enabled = true
spark.databricks.delta.autoCompact.enabled = true

# Caching
spark.databricks.io.cache.enabled = true

# Photon
spark.databricks.photon.enabled = true