from pyspark.sql import SparkSession

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

# This DataFrame will be stored in the Storage Region when cached
customers = spark.read.parquet("/data/customers")
customers.cache()  # → data goes into Storage Region on first action

# Trigger the cache with a count
customers.count()  # after this, data is in Storage Region

# All subsequent uses read from memory, not disk
orders = spark.read.parquet("/data/orders")
result1 = orders.join(customers, "customer_id")  # reads from memory
result2 = orders.join(customers, "customer_id").filter("amount > 1000")

# Broadcast a Python dict — lives in Storage Region on each executor
country_map = {"IN": "India", "US": "United States", "GB": "UK"}
bc_map = spark.sparkContext.broadcast(country_map)

# Each executor accesses bc_map.value from its OWN Storage Region copy
# No network call needed per task
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

country_lookup = udf(lambda code: bc_map.value.get(code, "Unknown"), StringType())
df.withColumn("country_name", country_lookup(col("country_code")))

from pyspark import StorageLevel

# Default: MEMORY_AND_DISK (deserialized, prefers memory)
df.persist(StorageLevel.MEMORY_AND_DISK)

# Serialized: smaller memory footprint, slower reads due to deserialization
df.persist(StorageLevel.MEMORY_AND_DISK_SER)

# Storage Levels comparison:
# MEMORY_ONLY          → fastest, most memory
# MEMORY_ONLY_SER      → 2x smaller, slower deserialization
# MEMORY_AND_DISK      → spills to disk if not enough memory
# MEMORY_AND_DISK_SER  → serialized + disk fallback
# DISK_ONLY            → always disk (slowest)
# OFF_HEAP             → Tungsten off-heap (no GC)

# This join triggers a shuffle → execution region gets shuffle buffers
df1.join(df2, "key")  # wide transformation → both sides shuffled

# Config that controls shuffle buffer size (per reduce task)
spark.conf.set("spark.shuffle.file.buffer", "64k")  # default 32k

# Controls in-memory sort buffer for shuffle output
spark.conf.set("spark.shuffle.sort.bypassMergeThreshold", "200")

from pyspark.sql.functions import broadcast

# Broadcast Join: small table's hash table built IN EXECUTION MEMORY
# on each executor. Must fit: spark.sql.autoBroadcastJoinThreshold (10MB default)
result = large_df.join(broadcast(small_df), "id")

# Adjust threshold to allow larger broadcast
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50m")  # 50 MB

# If hash table doesn't fit → falls back to Sort-Merge Join
# (which uses execution memory for sorting)

from pyspark.sql.functions import sum, count, avg

# This aggregation uses an in-memory hash map in the Execution Region
# key = (department, region), value = running sum/count/avg
result = df.groupBy("department", "region").agg(
    sum("sales").alias("total_sales"),
    count("*").alias("row_count"),
    avg("amount").alias("avg_amount")
)
# If distinct (dept, region) combos are huge → execution memory fills → spill

spark = SparkSession.builder \
    .appName("MemoryTuning") \
    .config("spark.executor.memory", "4g") \
    # 60% of (executor.memory - 300MB reserved) = Spark pool
    .config("spark.memory.fraction", "0.6") \
    # 50% of Spark pool = initial storage region (can shrink)
    .config("spark.memory.storageFraction", "0.5") \
    .getOrCreate()

# Quick math for 4GB executor:
# Usable heap        = 4096 - 300 = 3796 MB
# Spark memory pool  = 3796 * 0.6 = ~2278 MB
# Initial storage    = 2278 * 0.5 = ~1139 MB
# Initial execution  = 2278 * 0.5 = ~1139 MB
# User memory        = 3796 * 0.4 = ~1518 MB (UDFs, Python etc.)

# Size of each page in the internal memory pool (default 64MB)
spark.conf.set("spark.buffer.pageSize", "64m")

# Tune local disk directories where spill files are written
spark.conf.set("spark.local.dir", "/fast-ssd/tmp,/fast-ssd2/tmp")
# Use fast SSDs for local dir to reduce spill cost!

# See if spills are happening — check Spark UI Stages tab
# Look for: "Spill (Memory)" and "Spill (Disk)" columns

# After a job runs, check Spark UI REST API for spill info
import requests

# Spark UI REST API (adjust host/port)
url = "http://localhost:4040/api/v1/applications"
apps = requests.get(url).json()
app_id = apps[0]["id"]

stages_url = f"http://localhost:4040/api/v1/applications/{app_id}/stages"
stages = requests.get(stages_url).json()

for stage in stages:
    mem_spill = stage.get("memoryBytesSpilled", 0)
    disk_spill = stage.get("diskBytesSpilled", 0)
    if mem_spill > 0:
        print(f"Stage {stage['stageId']}: Spill Memory={mem_spill/1e9:.1f}GB, Disk={disk_spill/1e9:.1f}GB")

# Number of reduce partitions during shuffle (default 200)
# Too few → each partition is huge → spills
# Too many → too many small tasks → scheduling overhead
spark.conf.set("spark.sql.shuffle.partitions", "400")

# Memory buffer for shuffle writes per task (default 32k)
spark.conf.set("spark.shuffle.file.buffer", "64k")

# Max memory used by in-memory shuffle aggregation
# before spilling — controlled by execution memory fraction
spark.conf.set("spark.memory.fraction", "0.7")  # give more to Spark pool

# spark-submit
spark-submit \
  --executor-memory 8g \          # was 4g, now doubled
  --executor-cores 4 \
  --num-executors 20 \
  my_job.py

# In SparkSession config
spark = SparkSession.builder \
    .config("spark.executor.memory", "8g") \
    .config("spark.executor.memoryOverhead", "2g") \  # for off-heap
    .getOrCreate()

# Rule of thumb: each shuffle partition should be ~128–200 MB of data
# If total shuffle data = 100 GB → need ~500–800 partitions

spark.conf.set("spark.sql.shuffle.partitions", "800")

# OR use AQE to set this dynamically (recommended!)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
# AQE will coalesce small partitions, keeping each at ideal size

# Before a heavy join, increase partitions to reduce per-task data
df = df.repartition(400)  # now 400 tasks each handle smaller slice
result = df.join(other_df, "key")

# Even better: repartition by the join key to co-locate data
df = df.repartition(400, "key")
other_df = other_df.repartition(400, "key")
result = df.join(other_df, "key")  # much less shuffle needed!

# This orderBy sorts using Tungsten's binary UnsafeRow sort
# No Java object creation during sort — pure binary comparison
df.orderBy("amount", ascending=False)

# Even window function ordering uses UnsafeRow sort
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

w = Window.partitionBy("dept").orderBy("salary".desc())
df.withColumn("rank", row_number().over(w))
# Internally sorts each partition using UnsafeRow binary comparators

spark = SparkSession.builder \
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "4g") \   # 4 GB off-heap per executor
    .config("spark.executor.memory", "4g") \         # 4 GB on-heap
    .config("spark.executor.memoryOverhead", "4g") \ # must match offHeap.size
    .getOrCreate()

# Total executor memory consumption = 4g (heap) + 4g (off-heap) + overhead
# IMPORTANT: set memoryOverhead >= offHeap.size so YARN/K8s allocates enough

# WSCG is enabled by default — it turns chains of operators into one function
spark.conf.set("spark.sql.codegen.wholeStage", "true")  # default true

# See WSCG in action via explain
df.filter(col("age") > 30).select("name", "age").explain(True)
# Physical plan shows: *(1) Filter, *(1) Project
# The *(1) means they're fused into WholeStageCodegen stage 1
# One JIT-compiled loop handles both — no per-row function call overhead

spark-submit \
  --conf "spark.executor.extraJavaOptions=\
    -XX:+UseG1GC \
    -XX:G1HeapRegionSize=16m \
    -XX:MaxGCPauseMillis=500 \
    -XX:InitiatingHeapOccupancyPercent=35 \
    -XX:+G1PrintRegionLivenessInfo \
    -XX:+PrintGCDetails \
    -XX:+PrintGCDateStamps" \
  --conf "spark.driver.extraJavaOptions=\
    -XX:+UseG1GC \
    -XX:G1HeapRegionSize=16m" \
  my_job.py

# Full recommended G1GC config for Spark executors
-XX:+UseG1GC
-XX:G1HeapRegionSize=16m          # reduces humongous allocations
-XX:MaxGCPauseMillis=500           # target max pause time
-XX:InitiatingHeapOccupancyPercent=35  # start GC at 35% heap use
-XX:ConcGCThreads=4                # threads for concurrent GC work
-XX:+G1PrintRegionLivenessInfo     # debug: shows region usage
-XX:+PrintGCDetails                # verbose GC logs
-XX:+PrintGCDateStamps             # timestamp each GC event
-Xloggc:/tmp/gc-executor.log       # write GC log to file

# Use ParallelGC for batch jobs where throughput > latency
-XX:+UseParallelGC
-XX:ParallelGCThreads=8    # number of GC threads

# BAD: Python UDF creates Python objects for every row → GC pressure
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf(StringType())
def slow_udf(name):
    return name.upper()  # creates Python str objects per row

# GOOD: Use built-in Spark function — runs in JVM/Tungsten, no Python overhead
from pyspark.sql.functions import upper
df.withColumn("name_upper", upper(col("name")))  # zero GC pressure

# BAD: Large collect() → pulls everything to driver heap
all_data = df.collect()  # massive GC on driver

# GOOD: Use take() for sampling, or write to storage
sample = df.take(100)
df.write.parquet("/output")  # no driver heap usage

spark = SparkSession.builder \
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
    .config("spark.kryo.registrationRequired", "false") \
    # Optionally register your custom classes for max efficiency:
    .config("spark.kryo.classesToRegister", "com.mycompany.MyClass") \
    .getOrCreate()

# Note: DataFrame API already uses Tungsten binary format by default
# Kryo matters most for RDD operations on complex Scala/Java objects

# Example GC log line (from -XX:+PrintGCDetails -XX:+PrintGCDateStamps)
# 2024-03-15T10:23:45.123+0530: [GC pause (G1 Evacuation Pause) (young) 512M->128M(4096M), 0.234 secs]
#     ↑ Timestamp                ↑ GC type + trigger    ↑ before→after (heap) ↑ pause

# WARNING signs:
# "Full GC" — stop-the-world, all threads frozen, very expensive
# Pause > 5 seconds — task will appear hung in Spark UI
# Heap after GC ≈ Heap before GC — memory leak or high retention

# Quick check: grep GC log for Full GC events
grep "Full GC" /tmp/gc-executor.log | wc -l  # how many full GCs?
grep "secs" /tmp/gc-executor.log | awk '{print $NF}' | sort -n | tail -10

# G1GC manages young/old regions dynamically — less manual tuning needed
# For legacy collectors, you might tune:
-XX:NewRatio=3          # old:new = 3:1 (young gen = 25% of heap)
-XX:SurvivorRatio=8     # Eden:Survivor = 8:1
-XX:MaxTenuringThreshold=15  # survive 15 minor GCs before promotion

# For G1GC (recommended): let G1 manage regions automatically
# Just set -XX:MaxGCPauseMillis=500 and let G1 tune itself

# BAD — pulls ALL data to driver heap
result = df.collect()
pandas_df = df.toPandas()

# GOOD — write to storage, driver stays lean
df.write.mode("overwrite").parquet("/output/result")

# GOOD — take small sample to driver
sample = df.sample(0.01).toPandas()  # only 1% of data

# GOOD — increase driver memory if you must collect
spark = SparkSession.builder \
    .config("spark.driver.memory", "8g") \
    .config("spark.driver.maxResultSize", "4g") \
    .getOrCreate()

# COMMON CAUSE 1: Data skew — one partition is 100x bigger than others
# Fix: salt the key or use AQE skew join
from pyspark.sql.functions import concat, lit, floor, rand

# Add salt to even out skewed join key
salt_factor = 10
df_salted = df.withColumn("salted_key",
    concat(col("key"), lit("_"), (rand() * salt_factor).cast("int").cast("string"))
)

# COMMON CAUSE 2: UDF leaks memory or creates large objects
# Fix: switch to built-in functions

# COMMON CAUSE 3: Executor memory too low for workload
# Fix: increase executor memory
spark = SparkSession.builder \
    .config("spark.executor.memory", "16g") \
    .config("spark.executor.memoryOverhead", "4g") \
    .getOrCreate()

# COMMON CAUSE 4: Too many cores per executor → too much data per JVM
# Fix: reduce cores per executor (more, smaller executors)
# e.g., 4 executors × 5 cores each = bad (5 tasks compete for same heap)
# e.g., 10 executors × 2 cores each = better (2 tasks share heap)

# LEAK: broadcast variable never unpersisted
bc_data = spark.sparkContext.broadcast(large_dict)  # stays in memory forever
# FIX: destroy when done
bc_data.unpersist()
bc_data.destroy()

# LEAK: cache never unpersisted
df.cache()
# ... use df ...
# FIX: unpersist when done
df.unpersist()

# LEAK: accumulating static Python dict in UDF closure
growing_dict = {}  # grows per call in driver context
def leaky_udf(val):
    growing_dict[val] = 1  # driver memory leak!
    return val

# FIX: use broadcast for lookup data, never mutate driver-side state

# Add to executor and/or driver Java options
-XX:+HeapDumpOnOutOfMemoryError
-XX:HeapDumpPath=/tmp/heap-dump-executor.hprof

# In SparkSession config:
spark = SparkSession.builder \
    .config("spark.executor.extraJavaOptions",
            "-XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp/exec-heap.hprof") \
    .getOrCreate()

# After OOM, copy .hprof file from executor node
# Open with Eclipse MAT (free): https://www.eclipse.org/mat/
# Look for: "Leak Suspects" report → shows biggest retained objects

# Step 1: Enable AQE (helps auto-fix many spill causes)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

# Step 2: Check partition size distribution before heavy op
from pyspark.sql.functions import spark_partition_id, count

df.groupBy(spark_partition_id()).agg(count("*").alias("rows")) \
  .orderBy("rows", ascending=False) \
  .show(10)
# If one partition has 100x more rows → SKEW → use salting
# If all partitions are large → MEMORY → increase executor memory or partitions

# Step 3: Use explain to understand what operation causes spill
df.groupBy("key").agg(count("*")).explain(True)
# Look for SortAggregate (spills) vs HashAggregate (faster, uses memory)
# SortAggregate is used when data doesn't fit in hash table