# This entire Python file IS the driver process
# When you run: spark-submit my_job.py
# Python spawns a process that becomes the driver

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("MyApp") \
    .getOrCreate()

# sc.defaultParallelism — how many cores across all executors
print("Driver host:", spark.sparkContext.uiWebUrl)
print("Total cores:", spark.sparkContext.defaultParallelism)

# Driver collects this small result back to itself
df = spark.range(1000)
result = df.count()  # executors compute, driver receives the answer
print(f"Count: {result}")

spark = SparkSession.builder \
    .appName("InternalsDemo") \
    .master("local[4]") \          # 4 threads on local machine
    .config("spark.executor.memory", "4g") \
    .config("spark.executor.cores", "2") \
    .config("spark.sql.shuffle.partitions", "200") \
    .getOrCreate()   # reuses existing session if one exists

# SparkSession wraps SparkContext (lower-level API)
sc = spark.sparkContext

# Important SparkContext attributes
print(sc.appName)               # "InternalsDemo"
print(sc.applicationId)        # "app-20240101..."
print(sc.defaultParallelism)   # total cores across all executors
print(sc.master)               # "local[4]" or "yarn" etc.

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count

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

df = spark.range(1000000) \
    .filter(col("id") % 2 == 0) \          # narrow — same stage
    .withColumn("doubled", col("id") * 2) \  # narrow — same stage
    .groupBy("doubled") \                   # WIDE — causes shuffle → new stage
    .agg(count("*").alias("cnt"))

# Simple explain — shows physical plan
df.explain()

# Extended explain — shows all 4 plans
df.explain(extended=True)

# Formatted explain — best for reading (Spark 3.0+)
df.explain(mode="formatted")

# Cost-based explain (shows row estimates)
df.explain(mode="cost")

# Spark waits for better locality before falling back to ANY
# Default: spark will wait 3 seconds at each locality level
spark = SparkSession.builder \
    .config("spark.locality.wait", "3s") \      # total wait before ANY
    .config("spark.locality.wait.process", "0s") \ # no wait for process_local
    .config("spark.locality.wait.node", "2s") \   # wait 2s for node_local
    .getOrCreate()

# If using S3/cloud: disable locality wait — no point waiting
# .config("spark.locality.wait", "0s")  # skip to ANY immediately

from pyspark import StorageLevel

# MEMORY_ONLY: deserialized Java objects in heap. Fastest read.
# If partition doesn't fit, it's NOT stored (recomputed each time).
df.persist(StorageLevel.MEMORY_ONLY)

# MEMORY_AND_DISK: tries memory first, spills to disk if needed
# Most common choice in production
df.persist(StorageLevel.MEMORY_AND_DISK)

# MEMORY_ONLY_SER: serialized (byte arrays) — less memory, slower reads
df.persist(StorageLevel.MEMORY_ONLY_SER)

# DISK_ONLY: always on disk. Slowest. Use for very large rarely-used data.
df.persist(StorageLevel.DISK_ONLY)

# MEMORY_AND_DISK_2: replication factor of 2 (stores on 2 executors)
# Used when executor failure recovery is critical
df.persist(StorageLevel.MEMORY_AND_DISK_2)

# Release from Block Manager when done
df.unpersist()

# cache() is a shortcut for MEMORY_AND_DISK on DataFrames
df.cache()  # equivalent to df.persist(StorageLevel.MEMORY_AND_DISK)

# Most impactful shuffle config: number of shuffle partitions
# Default is 200 — often too high for small data, too low for large
spark.conf.set("spark.sql.shuffle.partitions", "400")

# Rule of thumb: aim for 128MB-256MB per partition after shuffle
# If total shuffle data = 80 GB → 80000 / 200 = ~400 MB per partition
# → set partitions to 400-640 for 128-200 MB each

# Compress shuffle data (saves network/disk, costs CPU)
spark.conf.set("spark.shuffle.compress", "true")       # default true
spark.conf.set("spark.shuffle.spill.compress", "true")  # default true

# Push-based shuffle (Spark 3.2+) — better for large shuffles
# Reduces small random reads on reduce side
spark.conf.set("spark.shuffle.push.enabled", "true")

# External shuffle service (important for YARN dynamic allocation)
# Without it, executor decommissioning loses shuffle data
spark.conf.set("spark.shuffle.service.enabled", "true")

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

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

df = spark.range(10_000_000) \
    .filter(col("id") > 100) \
    .groupBy((col("id") % 5).alias("bucket")) \
    .agg(sum("id").alias("total"))

# See the 4 plan stages in extended explain:
df.explain(extended=True)
# Output shows:
# == Parsed Logical Plan ==       (Step 2)
# == Analyzed Logical Plan ==     (Step 3)
# == Optimized Logical Plan ==    (Step 4)
# == Physical Plan ==             (Step 5)

# Look for these in physical plan:
# *(1) HashAggregate — partial agg (pre-shuffle)
# Exchange hashpartitioning — THE SHUFFLE
# *(2) HashAggregate — final agg (post-shuffle)
# The * means whole-stage code generation is active

spark = SparkSession.builder \
    .config("spark.executor.memory", "8g") \
    .config("spark.executor.memoryOverhead", "2g") \  # off-heap
    .config("spark.memory.fraction", "0.6") \          # default
    .config("spark.memory.storageFraction", "0.5") \   # default
    .getOrCreate()

# spark.executor.memoryOverhead is OFF-HEAP memory used for:
# - JVM overhead (thread stacks, code cache)
# - Python process memory (PySpark!)
# - Native memory (Arrow, Netty buffers)
# Rule: set to 10-20% of executor.memory, minimum 384MB

# For PySpark (Python UDFs): increase overhead significantly
# because Python process runs OUTSIDE the JVM heap
spark.conf.set("spark.executor.memoryOverhead", "3g")  # for heavy UDFs

spark = SparkSession.builder \
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
    .config("spark.kryo.registrationRequired", "false") \  # fallback allowed
    .getOrCreate()

# For RDD operations with custom classes:
# Register your classes with Kryo for best performance
# spark.kryo.classesToRegister = "com.mycompany.MyClass,..."

# Example RDD with custom objects
class SalesRecord:
    def __init__(self, id, amount):
        self.id = id
        self.amount = amount

# With Kryo: serialized as compact binary → faster shuffle
rdd = spark.sparkContext.parallelize([SalesRecord(1, 100.0)])
result = rdd.map(lambda r: r.amount * 1.1).collect()

# DataFrame UnsafeRow — automatic, no config needed:
df = spark.createDataFrame([(1, 100.0)], ["id", "amount"])
# ↑ Internally stored as UnsafeRow binary — already optimal

# DataFrames use schema-aware encoders automatically
# Each column is encoded by its type:
#   IntegerType  → 4 bytes inline in UnsafeRow
#   LongType     → 8 bytes inline
#   StringType   → offset+length pointer + UTF8 bytes
#   DoubleType   → 8 bytes inline
#   ArrayType    → offset + nested UnsafeArrayData

# WHY DataFrames >> RDDs with Python objects:
# RDD of Python dict: {"id": 1, "name": "Alice", "age": 30}
#   → Python object → pickle → bytes → JVM object → bytes → JVM
#   → 5 serialization steps per row per operation!

# DataFrame Row("Alice", 30):
#   → UnsafeRow [0xFF, 0x00, 41, 6C, ...] (binary once, stays binary)
#   → Zero deserialize steps between operators (Tungsten fuses them)

# Practical consequence:
data = [(i, f"user_{i}", i * 1.5) for i in range(1_000_000)]

# SLOW: Python objects through RDD
rdd = spark.sparkContext.parallelize(data)
result_rdd = rdd.map(lambda x: (x[0], x[2] * 2))

# FAST: UnsafeRow through DataFrame
df = spark.createDataFrame(data, ["id", "name", "val"])
result_df = df.select(col("id"), (col("val") * 2).alias("val2"))
# Typically 5-10x faster than RDD on same data

# ─── DRIVER OOM PREVENTION ───────────────────────────────────

# BAD: collect() pulls ALL data to driver — OOM if data is large
all_data = df.collect()  # ← DANGER: 10M rows → driver crash

# GOOD: Write to storage, don't bring to driver
df.write.parquet("s3://bucket/output")  # ← executors write directly

# GOOD: Take only what you need to driver
sample = df.limit(100).collect()  # ← only 100 rows, safe

# GOOD: Aggregate to small result before collecting
summary = df.groupBy("region").count().collect()  # ← few rows

# ─── EXECUTOR OOM PREVENTION ─────────────────────────────────

# BAD: One huge partition due to skew
df.groupBy("country").agg(F.collect_list("data"))  
# ↑ If 50% of data is "US", one executor gets 50% of all data → OOM

# GOOD: Use approx instead of exact collect
df.groupBy("country").agg(F.approx_count_distinct("data"))

# GOOD: Salt keys to distribute skewed data
df_fix = df.withColumn("key", 
    concat(col("country"), lit("_"), (rand() * 10).cast("int"))
)

from pyspark.sql.functions import broadcast, col

# ── Step 1: Check current plan ──────────────────────────────
sales = spark.read.parquet("s3://data/sales")      # 1 TB
products = spark.read.parquet("s3://data/products") # 500 MB

sales.join(products, "product_id").explain()
# If you see: SortMergeJoin → shuffle + sort of 1TB → very slow!

# ── Step 2: Force broadcast join for dimension table ────────
# Raise the auto-broadcast threshold to 1 GB
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "1073741824")  # 1GB

# OR use explicit broadcast hint (more reliable)
result = sales.join(broadcast(products), "product_id")
result.explain()
# Now shows: BroadcastHashJoin → products sent to all executors
# → zero shuffle of the 1TB table!

# ── Step 3: Partition pruning on date ───────────────────────
# Only read relevant partitions of the 1TB table
sales_q4 = spark.read.parquet("s3://data/sales") \
    .filter(col("sale_date") >= "2024-10-01")
# Spark only reads 3 month partitions instead of 12 → 4x less I/O

# ── Step 4: Enable AQE (handles runtime surprises) ──────────
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

# ── Step 5: Configure for large join ─────────────────────────
spark.conf.set("spark.sql.shuffle.partitions", "2000")
# 1TB / 2000 partitions ≈ 500MB per partition — manageable

# ── Final optimized join ─────────────────────────────────────
result = sales_q4.join(broadcast(products), "product_id") \
    .filter(col("amount") > 0) \
    .groupBy("product_name", "region") \
    .agg(F.sum("amount").alias("total"))

result.write.parquet("s3://output/q4_sales")