# SparkContext (Spark 1.x) — entry point for RDD API
from pyspark import SparkContext
sc = SparkContext("local[*]", "MyApp")

# SparkSession (Spark 2.0+) — entry point for EVERYTHING
# Wraps SparkContext + SQLContext + HiveContext in one object
from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("MyApp") \
    .master("yarn") \
    .config("spark.executor.memory", "4g") \
    .config("spark.executor.cores", "2") \
    .config("spark.sql.shuffle.partitions", "200") \
    .enableHiveSupport() \
    .getOrCreate()               # creates new or returns existing session

sc = spark.sparkContext          # access SparkContext from SparkSession

# EXAMPLE: Lazy evaluation in action
df = spark.read.csv("/huge/file.csv")       # ← NO execution (just reads metadata)
df2 = df.filter("age > 30")                 # ← NO execution (builds DAG node)
df3 = df2.select("name", "age")             # ← NO execution (builds DAG node)
df4 = df3.withColumn("age_group",           # ← NO execution (builds DAG node)
        when(col("age") > 60, "senior").otherwise("adult"))

# ONLY NOW does Spark execute — when you call an ACTION:
result = df4.count()    # ← TRIGGERS execution of the ENTIRE plan
df4.show()              # ← TRIGGERS execution again (recomputes from scratch unless cached!)

sc = spark.sparkContext

# From a Python collection
rdd = sc.parallelize([1, 2, 3, 4, 5], numSlices=4)  # 4 partitions
print(rdd.getNumPartitions())   # 4

# From a file
rdd = sc.textFile("hdfs:///path/to/file.txt", minPartitions=10)

# From another RDD (transformation)
filtered = rdd.filter(lambda x: x > 2)  # new RDD, not executed yet

# ─── map() — apply function to EACH element, one-to-one ───
rdd = sc.parallelize([1, 2, 3, 4])
doubled = rdd.map(lambda x: x * 2)
doubled.collect()   # [2, 4, 6, 8]

# ─── filter() — keep elements matching condition ───
evens = rdd.filter(lambda x: x % 2 == 0)
evens.collect()     # [2, 4]

# ─── flatMap() — one input → ZERO OR MORE outputs (flattened) ───
sentences = sc.parallelize(["hello world", "foo bar"])
words = sentences.flatMap(lambda s: s.split(" "))
words.collect()     # ['hello', 'world', 'foo', 'bar']
# map() would give: [['hello', 'world'], ['foo', 'bar']]  (nested lists!)
# flatMap() flattens: ['hello', 'world', 'foo', 'bar']     (flat list!)

# ─── mapPartitions() — function runs ON ENTIRE PARTITION (not row-by-row) ───
# MORE EFFICIENT for expensive initialization (DB connections, ML models)
def process_partition(iterator):
    # Setup once per partition (not once per row!)
    connection = create_db_connection()
    for record in iterator:
        yield connection.lookup(record)
    connection.close()

result = rdd.mapPartitions(process_partition)

# ─── reduceByKey() vs groupByKey() — CRITICAL interview question ───

pairs = sc.parallelize([("a", 1), ("b", 2), ("a", 3), ("b", 4)])

# reduceByKey — aggregates LOCALLY first, then shuffles partial results
total = pairs.reduceByKey(lambda a, b: a + b)
total.collect()     # [('a', 4), ('b', 6)]
# What happens: Each partition reduces locally, only small partial sums are shuffled

# groupByKey — shuffles ALL raw values, THEN groups them
grouped = pairs.groupByKey().mapValues(list)
grouped.collect()   # [('a', [1, 3]), ('b', [2, 4])]
# What happens: ALL individual values sent across network → much more data shuffled!

# ─── aggregateByKey — when output type differs from input ───
# Use case: computing average (need sum AND count, not just sum)
pairs = sc.parallelize([("a", 1), ("b", 2), ("a", 3), ("b", 4)])

zero_value = (0, 0)  # (sum, count)
seq_func = lambda acc, val: (acc[0] + val, acc[1] + 1)     # within partition
comb_func = lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])  # across partitions

avg = pairs.aggregateByKey(zero_value, seq_func, comb_func) \
           .mapValues(lambda x: x[0] / x[1])
avg.collect()   # [('a', 2.0), ('b', 3.0)]

# ─── join() — join two key-value RDDs by key ───
rdd1 = sc.parallelize([("a", 1), ("b", 2)])
rdd2 = sc.parallelize([("a", "x"), ("b", "y")])
joined = rdd1.join(rdd2)
joined.collect()    # [('a', (1, 'x')), ('b', (2, 'y'))]

# ─── coalesce() vs repartition() ───
rdd.repartition(10)    # full shuffle, increases OR decreases to 10 equal partitions
rdd.coalesce(4)         # no full shuffle, reduces (merges adjacent partitions)
# RULE: Use coalesce() to DECREASE partitions (no shuffle)
#       Use repartition() to INCREASE partitions (requires shuffle)

# Actions TRIGGER computation and return results to Driver
rdd = sc.parallelize([1, 2, 3, 4, 5])

rdd.collect()               # returns ALL elements as Python list ⚠️ OOM risk on large RDDs!
rdd.count()                 # count of elements → 5
rdd.first()                 # first element → 1
rdd.take(3)                 # first N elements → [1, 2, 3]
rdd.top(3)                  # top N (sorted descending) → [5, 4, 3]
rdd.takeSample(False, 3)    # 3 random samples (False = without replacement)
rdd.reduce(lambda a, b: a + b)    # aggregate all elements → 15
rdd.fold(0, lambda a, b: a + b)   # like reduce but with zero value → 15
rdd.foreach(print)          # apply function to each element (no return value)
rdd.saveAsTextFile("/path/")    # write each partition as a text file
rdd.countByValue()          # count occurrences of each unique value
rdd.countByKey()            # for key-value RDDs

# On a 100 GB dataset:
result = df.collect()  # ← tries to send ALL 100 GB to Driver memory → OutOfMemoryError!

# Safe alternatives:
df.show(20)              # show first 20 rows (small data to Driver)
df.take(100)             # take 100 rows only
df.write.parquet("/out") # write to storage (no data to Driver)
df.count()               # just returns a number

from pyspark import StorageLevel

# cache() = persist(MEMORY_AND_DISK) for DataFrames
# cache() = persist(MEMORY_ONLY) for RDDs
df.cache()

# persist() lets you choose storage level explicitly
rdd.persist(StorageLevel.MEMORY_ONLY)             # JVM heap, deserialized (fastest access)
rdd.persist(StorageLevel.MEMORY_AND_DISK)         # Memory first, spill to disk if full
rdd.persist(StorageLevel.MEMORY_ONLY_SER)         # Serialized in memory (more compact)
rdd.persist(StorageLevel.MEMORY_AND_DISK_SER)     # Serialized + disk spill
rdd.persist(StorageLevel.DISK_ONLY)               # Always on disk (slowest)
rdd.persist(StorageLevel.OFF_HEAP)                # Tungsten off-heap memory

# Unpersist when done (free up memory for other operations)
rdd.unpersist()

rdd.cache()              # ← nothing happens yet! Just marks for caching
rdd.count()              # ← NOW the RDD is computed AND cached
rdd.collect()            # ← reads from cache (fast!)

# View lineage of an RDD
rdd1 = sc.textFile("/data/logs.txt")
rdd2 = rdd1.filter(lambda line: "ERROR" in line)
rdd3 = rdd2.map(lambda line: line.split(","))
rdd4 = rdd3.map(lambda parts: (parts[0], parts[1]))
rdd5 = rdd4.reduceByKey(lambda a, b: a + b)

print(rdd5.toDebugString())
# Shows: rdd5 ← rdd4 ← rdd3 ← rdd2 ← rdd1 ← textFile(logs.txt)

# PROBLEM: After 100 ML iterations, lineage is 100 steps deep
# If one partition fails → recompute ALL 100 steps from scratch!

# SOLUTION: Checkpoint! Saves to HDFS, truncates lineage
sc.setCheckpointDir("hdfs:///checkpoints/")
rdd5.checkpoint()   # mark for checkpoint
rdd5.count()        # triggers: compute, save to HDFS, new lineage starts from checkpoint

# Best practice: cache THEN checkpoint (avoids computing twice)
rdd5.cache()
rdd5.checkpoint()
rdd5.count()   # computes once → stores in cache AND saves to HDFS

# ─── BROADCAST VARIABLES ───
# USE CASE: lookup tables, configuration data, ML model parameters

# WITHOUT broadcast (BAD for large lookups)
country_codes = {"US": "United States", "IN": "India", "GB": "United Kingdom"}
rdd.map(lambda code: country_codes.get(code, "Unknown"))
# ^ Python closure: serialized and sent with EVERY task!
# 1000 tasks × 100 MB table = 100 GB network traffic!

# WITH broadcast (GOOD)
bc_codes = sc.broadcast(country_codes)
rdd.map(lambda code: bc_codes.value.get(code, "Unknown"))
# ^ Sent ONCE to each executor, all tasks on that executor share it

bc_codes.unpersist()    # release from executor memory
bc_codes.destroy()      # remove from ALL executors immediately

# ─── ACCUMULATORS ───
# USE CASE: counters (bad records, events, debug metrics)

error_count = sc.accumulator(0)
malformed_count = sc.accumulator(0)

def parse_line(line):
    try:
        parts = line.split(",")
        if len(parts) != 5:
            malformed_count.add(1)
            return None
        return parts
    except Exception:
        error_count.add(1)
        return None

results = rdd.map(parse_line).filter(lambda x: x is not None)
results.count()    # trigger action

print(f"Errors: {error_count.value}")
print(f"Malformed: {malformed_count.value}")

# Accumulators in TRANSFORMATIONS can be counted MULTIPLE TIMES
# because Spark may retry failed tasks!
# If a map() task fails and retries → accumulator incremented AGAIN

# SAFE: Use accumulators inside foreach() (an action — runs exactly once per record)
rdd.foreach(lambda x: error_count.add(1) if x is None else None)

# Push an eligible filter to a columnar source and read only needed columns.
result = (
    spark.read.parquet("data/")
         .filter(col("year") == 2024)
         .select("name", "amount")
)

# A Python UDF is opaque to Catalyst and can block source pushdown.
blocked = df.filter(my_udf(col("year")) == 2024)

df.explain()             # physical plan
df.explain(True)         # parsed, analyzed, optimized, and physical plans
df.explain("formatted")  # structured physical plan
df.explain("cost")       # plan plus available statistics
df.explain("codegen")    # generated code where supported

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("DataPipeline") \
    .master("yarn") \
    .config("spark.executor.memory", "8g") \
    .config("spark.executor.cores", "4") \
    .config("spark.executor.instances", "10") \
    .config("spark.sql.shuffle.partitions", "400") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.autoBroadcastJoinThreshold", str(50 * 1024 * 1024)) \
    .enableHiveSupport() \
    .getOrCreate()

# getOrCreate() -> creates new session OR returns existing one
# Prevents creating multiple SparkSessions in the same application

# --- FROM A PYTHON LIST ---
data = [("Alice", 30, "Engineering"), ("Bob", 25, "Marketing")]
columns = ["name", "age", "department"]
df = spark.createDataFrame(data, columns)

# --- FROM AN RDD ---
rdd = spark.sparkContext.parallelize([("Alice", 30), ("Bob", 25)])
df = rdd.toDF(["name", "age"])

# --- FROM A PANDAS DATAFRAME ---
import pandas as pd
pdf = pd.DataFrame({"name": ["Alice", "Bob"], "age": [30, 25]})
df = spark.createDataFrame(pdf)

# --- FROM AN EXISTING TABLE (Hive/Delta metastore) ---
df = spark.table("database_name.table_name")

# --- FROM A RANGE ---
df = spark.range(0, 100, 1)  # single column 'id' with values 0-99

from pyspark.sql.types import (
    StructType, StructField,
    StringType, IntegerType, LongType, DoubleType, FloatType,
    BooleanType, DateType, TimestampType, ArrayType, MapType
)

# Define schema explicitly (no inferSchema scan needed)
booking_schema = StructType([
    StructField("booking_id",   StringType(),    nullable=False),
    StructField("customer_id",  LongType(),      nullable=True),
    StructField("flight_code",  StringType(),    nullable=True),
    StructField("amount",       DoubleType(),    nullable=True),
    StructField("booking_date", DateType(),      nullable=True),
    StructField("is_cancelled", BooleanType(),   nullable=True),
    # Nested struct:
    StructField("address", StructType([
        StructField("city",    StringType(), True),
        StructField("country", StringType(), True)
    ]), True),
    # Array column:
    StructField("tags", ArrayType(StringType()), True),
    # Map column:
    StructField("metadata", MapType(StringType(), StringType()), True)
])

df = spark.read.schema(booking_schema).csv("/data/bookings.csv", header=True)

# --- CSV ---
df_csv = spark.read \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .option("delimiter", ",") \
    .option("nullValue", "N/A") \
    .option("dateFormat", "yyyy-MM-dd") \
    .csv("/data/bookings.csv")
# WARNING: inferSchema triggers extra read pass — use explicit schema in prod

# --- JSON ---
df_json = spark.read \
    .option("multiLine", "true") \
    .option("mode", "PERMISSIVE") \
    .json("/data/events/*.json")
# mode options: PERMISSIVE (default, nulls bad fields) | DROPMALFORMED | FAILFAST

# --- PARQUET (PREFERRED FORMAT) ---
df_parquet = spark.read.parquet("/data/warehouse/bookings/")
# Schema embedded in file (no inferSchema needed)
# Columnar format with predicate pushdown
# Efficient compression (Snappy by default)

# --- ORC ---
df_orc = spark.read.orc("/data/hive_table/")

# --- JDBC (relational database) ---
df_jdbc = spark.read \
    .format("jdbc") \
    .option("url", "jdbc:postgresql://db-host:5432/mydb") \
    .option("dbtable", "bookings") \
    .option("user", "etl_user") \
    .option("password", "secret") \
    .option("driver", "org.postgresql.Driver") \
    .option("numPartitions", "10") \
    .option("partitionColumn", "booking_id") \
    .option("lowerBound", "1") \
    .option("upperBound", "1000000") \
    .load()

# --- AVRO ---
df_avro = spark.read.format("avro").load("/data/kafka_output/")

# --- DELTA LAKE ---
df_delta = spark.read.format("delta").load("/data/delta/bookings/")
# Or using table name if registered in metastore:
df_delta = spark.table("bookings_db.bookings")

# --- GLOB PATTERN (most common) ---
df = spark.read.csv("/data/2024/*/bookings_*.csv", header=True)
# Reads: /data/2024/jan/bookings_1.csv, /data/2024/feb/bookings_2.csv etc.

# --- LIST OF SPECIFIC FILES ---
df = spark.read.parquet(
    "/data/2024/jan/bookings.parquet",
    "/data/2024/feb/bookings.parquet",
    "/data/2024/mar/bookings.parquet"
)
# Or pass as Python list:
files = ["/data/2024/jan/bookings.parquet", "/data/2024/feb/bookings.parquet"]
df = spark.read.parquet(*files)   # unpack list

# --- DIRECTORY (all files in dir with same format) ---
df = spark.read.parquet("/data/2024/")   # reads ALL parquet files in dir

# --- ADD SOURCE FILE NAME COLUMN ---
from pyspark.sql.functions import input_file_name
df = spark.read.option("header", "true") \
    .csv("/data/2024/*/bookings_*.csv") \
    .withColumn("source_file", input_file_name())

# --- READ MULTIPLE FORMATS AND UNION ---
df_csv = spark.read.option("header", True).csv("/landing/batch/")
df_json = spark.read.json("/landing/api/events/")
df_parquet = spark.read.parquet("/landing/warehouse/")

# Align columns then union
cols = ["booking_id", "customer_id", "amount", "booking_date"]
df_all = df_csv.select(cols) \
    .union(df_json.select(cols)) \
    .union(df_parquet.select(cols))

# Better: unionByName (aligns by column name, not position)
df_all = df_csv.unionByName(df_json, allowMissingColumns=True) \
               .unionByName(df_parquet, allowMissingColumns=True)

# --- READ FROM DIFFERENT STORAGE SYSTEMS ---
df_s3 = spark.read.parquet("s3a://my-bucket/data/")
df_adls = spark.read.parquet("abfss://container@storage.dfs.core.windows.net/data/")
df_hdfs = spark.read.parquet("hdfs://namenode:8020/data/")
df_local = spark.read.parquet("file:///local/path/data/")

# Combine multiple storage locations:
df_combined = spark.read.parquet(
    "s3a://bucket/data/2023/",
    "s3a://bucket/data/2024/",
    "hdfs://namenode/archive/2022/"
)

df.select("id", "name", "amount")                    # by name
df.select(col("id"), col("amount") * 1.1)            # with expression
df.select("*", (col("amount") * 1.1).alias("new_amount"))  # all + new column

df.filter(col("age") > 30)
df.filter("age > 30")                                 # SQL string form
df.where((col("country") == "IN") & (col("amount") > 100))
df.filter(col("status").isin("CONFIRMED", "PENDING"))
df.filter(col("name").startswith("A"))
df.filter(col("name").like("A%"))
df.filter(~col("is_cancelled"))                       # NOT condition

df.withColumn("tax", col("amount") * 0.18)
df.withColumn("amount", col("amount").cast("double"))  # modify existing column
df.withColumn("status", when(col("amount") > 1000, "high")
                        .when(col("amount") > 100, "medium")
                        .otherwise("low"))

df.drop("col1", "col2")
df.drop(col("temporary_column"))

df.distinct()                                     # all columns must match
df.dropDuplicates(["customer_id", "booking_date"])  # subset columns

# --- RENAME ---
df.withColumnRenamed("old_name", "new_name")

# --- SORT ---
df.orderBy("amount")                              # ascending
df.orderBy(col("amount").desc())                  # descending
df.orderBy(col("date").asc(), col("amount").desc())  # multi-column

from pyspark.sql.functions import col, lit

# col() — reference an existing column
df.select(col("name"), col("amount") * 2)

# lit() — create a constant value column
df.withColumn("country", lit("India"))
df.withColumn("multiplier", lit(1.18))

# Common mistake: trying to use a Python variable directly
tax_rate = 0.18
# WRONG: df.withColumn("tax", col("amount") * tax_rate)  # works but unclear
# RIGHT: df.withColumn("tax", col("amount") * lit(tax_rate))  # explicit

from pyspark.sql.functions import when

# Simple condition
df.withColumn("category",
    when(col("amount") > 1000, "high")
    .when(col("amount") > 100, "medium")
    .otherwise("low")
)

# Multiple conditions combined
df.withColumn("flag",
    when((col("status") == "CANCELLED") & (col("amount") > 500), "refund_priority")
    .when(col("status") == "CANCELLED", "standard_refund")
    .otherwise("no_action")
)

# Nested when for complex logic
df.withColumn("tier",
    when(col("total_spend") > 10000, "platinum")
    .when((col("total_spend") > 5000) & (col("years") > 3), "gold")
    .when(col("total_spend") > 1000, "silver")
    .otherwise("bronze")
)

df.withColumn("amount", col("amount").cast("double"))
df.withColumn("booking_date", col("booking_date").cast("date"))
df.withColumn("amount_int", col("amount").cast(IntegerType()))

# Common cast operations:
# "string" -> "integer", "long", "double", "float", "date", "timestamp", "boolean"

from pyspark.sql.functions import count, sum, avg, max, min, countDistinct, collect_list, collect_set

# --- BASIC GROUPBY + AGG ---
df.groupBy("country") \
  .agg(
      count("*").alias("booking_count"),
      sum("amount").alias("total_amount"),
      avg("amount").alias("avg_amount"),
      max("amount").alias("max_amount"),
      min("amount").alias("min_amount"),
      countDistinct("customer_id").alias("unique_customers")
  )

# --- MULTIPLE GROUPBY COLUMNS ---
df.groupBy("country", "booking_year") \
  .agg(
      count("*").alias("bookings"),
      sum("amount").alias("revenue")
  )

# --- COLLECT VALUES INTO A LIST/SET ---
df.groupBy("customer_id") \
  .agg(
      collect_list("product").alias("all_products"),     # allows duplicates
      collect_set("product").alias("unique_products")    # no duplicates
  )

# --- SHORTCUT METHODS (less flexible but quick) ---
df.groupBy("country").count()
df.groupBy("country").sum("amount")
df.groupBy("country").avg("amount")
df.groupBy("country").max("amount")

from pyspark.sql.functions import broadcast, col

# --- INNER JOIN (only matching rows from both sides) ---
result = df1.join(df2, on="customer_id", how="inner")

# --- LEFT JOIN (all rows from left, matching from right, null if no match) ---
result = df1.join(df2, on=["customer_id", "date"], how="left")

# --- RIGHT JOIN (all rows from right, matching from left) ---
result = df1.join(df2, on="customer_id", how="right")

# --- FULL OUTER JOIN (all rows from both, nulls where no match) ---
result = df1.join(df2, on="customer_id", how="full")

# --- CROSS JOIN (cartesian product — every row with every row) ---
result = df1.crossJoin(df2)
# WARNING: N x M rows — use only when intentional (e.g., generating combinations)

# --- LEFT SEMI JOIN (rows in df1 that HAVE a match in df2, NO df2 columns) ---
result = df1.join(df2, on="customer_id", how="left_semi")
# Equivalent to: WHERE customer_id IN (SELECT customer_id FROM df2)

# --- LEFT ANTI JOIN (rows in df1 that do NOT have a match in df2) ---
result = df1.join(df2, on="customer_id", how="left_anti")
# Equivalent to: WHERE customer_id NOT IN (SELECT customer_id FROM df2)

# --- JOIN WITH DIFFERENT COLUMN NAMES ---
result = df1.join(df2, df1["cust_id"] == df2["customer_id"], how="inner")
# After join, drop the duplicate column:
result = result.drop(df2["customer_id"])

# --- BROADCAST JOIN (force small table broadcast — no shuffle) ---
result = large_df.join(broadcast(small_df), on="airport_code")
# Spark sends the small table to every executor (avoids shuffle of large table)
# Default broadcast threshold: 10 MB (spark.sql.autoBroadcastJoinThreshold)

from pyspark.sql.window import Window
from pyspark.sql.functions import (
    row_number, rank, dense_rank, lag, lead,
    sum, avg, max, min, count,
    ntile, percent_rank, cume_dist,
    first, last
)

# --- DEFINE A WINDOW ---
windowSpec = Window \
    .partitionBy("department") \
    .orderBy(col("salary").desc())

df = df.withColumn("row_num", row_number().over(windowSpec))  # 1,2,3,4,5 (no ties)
df = df.withColumn("rank",    rank().over(windowSpec))         # 1,2,2,4,5 (ties skip)
df = df.withColumn("drank",   dense_rank().over(windowSpec))   # 1,2,2,3,4 (ties no skip)

# Top 2 highest-paid employees per department
window = Window.partitionBy("department").orderBy(col("salary").desc())

top2 = df.withColumn("rank", dense_rank().over(window)) \
         .filter(col("rank") <= 2) \
         .drop("rank")

lag_spec = Window.partitionBy("customer_id").orderBy("purchase_date")

df = df.withColumn("prev_purchase", lag("amount", 1, 0).over(lag_spec))
#                                   column, offset, default_if_null

df = df.withColumn("next_purchase", lead("amount", 1).over(lag_spec))

# Month-over-month change
df = df.withColumn("mom_change",
    round((col("revenue") - lag("revenue", 1).over(lag_spec))
          / lag("revenue", 1).over(lag_spec) * 100, 2))

# --- RUNNING TOTAL ---
running_spec = Window.partitionBy("region") \
    .orderBy("sale_date") \
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)

df = df.withColumn("running_total", sum("revenue").over(running_spec))

# --- ROLLING 7-DAY AVERAGE ---
rolling_spec = Window.partitionBy("user_id") \
    .orderBy("event_date") \
    .rowsBetween(-6, 0)  # current + 6 before = 7 rows

df = df.withColumn("rolling_7day_avg", avg("daily_count").over(rolling_spec))

# --- NTILE (divide into N equal buckets) ---
quartile_spec = Window.orderBy("spend")
df = df.withColumn("quartile", ntile(4).over(quartile_spec))
df = df.withColumn("percentile", percent_rank().over(quartile_spec))

# --- FIRST / LAST ---
first_spec = Window.partitionBy("customer_id").orderBy("purchase_date")
df = df.withColumn("first_purchase_amount",
    first("amount", ignorenulls=True).over(first_spec))

# --- SESSION-SCOPED VIEW (most common) ---
df.createOrReplaceTempView("bookings")

# --- GLOBAL VIEW (visible across sessions) ---
df.createOrReplaceGlobalTempView("bookings")
# Access global views with: global_temp.bookings

# --- RUNNING SQL ---
result = spark.sql("""
    SELECT
        country,
        COUNT(*) as booking_count,
        SUM(amount) as total_revenue,
        ROUND(AVG(amount), 2) as avg_booking
    FROM bookings
    WHERE booking_date >= '2024-01-01'
      AND is_cancelled = false
    GROUP BY country
    HAVING COUNT(*) > 100
    ORDER BY total_revenue DESC
""")

# --- SQL WITH WINDOW FUNCTIONS ---
spark.sql("""
    SELECT *,
        ROW_NUMBER() OVER (PARTITION BY department ORDER BY salary DESC) as rank
    FROM employees
""")

# --- MIXING SQL AND DATAFRAME API ---
spark.sql("SELECT * FROM bookings WHERE amount > 1000").groupBy("country").count()

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, DoubleType, ArrayType

# --- REGULAR PYTHON UDF ---
def clean_phone(phone):
    """Remove non-digit chars from phone number"""
    import re
    return re.sub(r'\D', '', phone) if phone else None

# Register as UDF
clean_phone_udf = udf(clean_phone, StringType())

# Use with DataFrame API:
df = df.withColumn("clean_phone", clean_phone_udf(col("phone")))

# Register for SQL use:
spark.udf.register("clean_phone", clean_phone, StringType())
spark.sql("SELECT clean_phone(phone) FROM customers")

@udf(returnType=DoubleType())
def calculate_tax(amount, rate):
    return amount * rate if amount and rate else 0.0

df = df.withColumn("tax", calculate_tax(col("amount"), col("tax_rate")))

# Uses Apache Arrow for zero-copy data transfer
# Processes data in COLUMNAR BATCHES (not row by row)
# 3x to 100x faster than regular Python UDF!

from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf(StringType())
def normalize_name(names: pd.Series) -> pd.Series:
    """Vectorized: processes entire column as pandas Series"""
    return names.str.strip().str.title().fillna("Unknown")

df = df.withColumn("clean_name", normalize_name(col("name")))

# Pandas UDF with multiple columns:
@pandas_udf(DoubleType())
def calculate_discount(amount: pd.Series, tier: pd.Series) -> pd.Series:
    discount_map = {"gold": 0.15, "silver": 0.10, "bronze": 0.05}
    discount = tier.map(discount_map).fillna(0)
    return amount * discount

df = df.withColumn("discount", calculate_discount(col("amount"), col("tier")))

# --- WRITE MODES ---
# overwrite: replace existing data entirely
# append: add to existing data
# ignore: don't write if destination already exists
# error / errorIfExists: fail if destination exists (default)

df.write.mode("overwrite").parquet("/output/bookings/")

# --- WRITE PARTITIONED (Hive-style partitioning) ---
df.write \
  .mode("overwrite") \
  .partitionBy("booking_year", "booking_month") \
  .parquet("/output/bookings_partitioned/")
# Creates: /output/bookings_partitioned/booking_year=2024/booking_month=01/part-00000.parquet

# --- WRITE BUCKETED (for joins -- avoids shuffle!) ---
df.write \
  .mode("overwrite") \
  .bucketBy(32, "customer_id") \
  .sortBy("customer_id") \
  .saveAsTable("default.bookings_bucketed")
# When two tables bucketed by same column + same N buckets: JOIN has NO SHUFFLE!

# --- WRITE AS DELTA ---
df.write \
  .format("delta") \
  .mode("overwrite") \
  .save("/delta/bookings/")

# --- WRITE CSV WITH OPTIONS ---
df.write \
  .option("header", "true") \
  .option("delimiter", "|") \
  .mode("overwrite") \
  .csv("/output/bookings.csv")

# --- WRITE TO JDBC ---
df.write \
  .format("jdbc") \
  .option("url", "jdbc:postgresql://db:5432/mydb") \
  .option("dbtable", "bookings") \
  .option("user", "user") \
  .option("password", "pass") \
  .mode("append") \
  .save()

# --- CONTROL NUMBER OF OUTPUT FILES ---
df.repartition(10).write.parquet("/output/")     # exactly 10 files
df.coalesce(1).write.csv("/output/single_file/") # single file (avoid for large data)

from pyspark.sql.functions import col, coalesce, lit, isnull

# --- FINDING NULLS ---
df.filter(col("amount").isNull())                  # rows where amount is null
df.filter(col("amount").isNotNull())               # rows where amount is not null

# --- DROPPING ROWS WITH NULLS ---
df.na.drop()                                       # drop rows with ANY null
df.na.drop(how="all")                              # drop only if ALL columns null
df.na.drop(subset=["customer_id", "amount"])       # only check these columns
df.na.drop(how="any", thresh=3)                    # keep rows with at least 3 non-null values

# --- FILLING NULLS ---
df.na.fill(0)                                      # fill all numeric nulls with 0
df.na.fill("")                                     # fill all string nulls with empty string
df.na.fill({"amount": 0, "country": "UNKNOWN"})    # per-column fill

# --- COALESCE (first non-null value from multiple columns) ---
df.withColumn("phone",
    coalesce(col("mobile_phone"), col("home_phone"), col("work_phone"), lit("N/A"))
)

# --- SAFE DIVISION (avoid null/zero errors) ---
from pyspark.sql.functions import when
df.withColumn("rate",
    when(col("impressions") != 0,
         col("clicks") / col("impressions"))
    .otherwise(None))

# --- NULL-SAFE EQUALITY (<=>) ---
# Regular ==: NULL == NULL returns NULL (not True!)
# Null-safe <=>: NULL <=> NULL returns True
df.filter(col("a").eqNullSafe(col("b")))

from pyspark.sql.functions import (
    concat, concat_ws, substring, trim, ltrim, rtrim,
    lower, upper, initcap, length, lpad, rpad,
    regexp_replace, regexp_extract, split,
    instr, locate, translate, reverse,
    col, lit
)

# --- CONCATENATION ---
df.withColumn("full_name", concat(col("first_name"), lit(" "), col("last_name")))
df.withColumn("full_name", concat_ws(" ", col("first_name"), col("middle"), col("last_name")))
# concat_ws skips NULLs; concat returns NULL if any input is NULL

# --- CASE CONVERSION ---
df.withColumn("upper_name", upper(col("name")))
df.withColumn("lower_name", lower(col("name")))
df.withColumn("title_name", initcap(col("name")))  # "alice smith" -> "Alice Smith"

# --- TRIMMING ---
df.withColumn("clean", trim(col("name")))           # both sides
df.withColumn("clean", ltrim(col("name")))          # left only
df.withColumn("clean", rtrim(col("name")))          # right only

# --- SUBSTRING ---
df.withColumn("area_code", substring(col("phone"), 1, 3))  # first 3 chars (1-indexed)

# --- REGEX REPLACE ---
df.withColumn("clean_phone", regexp_replace(col("phone"), r"[^0-9]", ""))
df.withColumn("no_special", regexp_replace(col("text"), r"[^a-zA-Z0-9 ]", ""))

# --- REGEX EXTRACT ---
df.withColumn("domain", regexp_extract(col("email"), r"@(.+)", 1))
# Group 1 extracts the domain from "user@example.com" -> "example.com"

# --- SPLIT ---
df.withColumn("parts", split(col("full_name"), " "))       # returns array
df.withColumn("first", split(col("full_name"), " ")[0])    # first element
df.withColumn("last", split(col("full_name"), " ")[1])     # second element

# --- PADDING ---
df.withColumn("padded_id", lpad(col("id"), 10, "0"))   # "42" -> "0000000042"

# --- LENGTH ---
df.withColumn("name_len", length(col("name")))

from pyspark.sql.functions import (
    current_date, current_timestamp,
    datediff, months_between, date_add, date_sub,
    date_format, to_date, to_timestamp,
    year, month, dayofmonth, dayofweek, dayofyear,
    hour, minute, second, weekofyear, quarter,
    last_day, next_day, trunc, date_trunc,
    col, lit
)

# --- CURRENT DATE/TIMESTAMP ---
df.withColumn("today", current_date())
df.withColumn("now", current_timestamp())

# --- STRING TO DATE/TIMESTAMP ---
df.withColumn("parsed_date", to_date(col("date_str"), "yyyy-MM-dd"))
df.withColumn("parsed_ts", to_timestamp(col("ts_str"), "yyyy-MM-dd HH:mm:ss"))
# Common formats: "yyyy-MM-dd", "MM/dd/yyyy", "dd-MMM-yyyy", "yyyyMMdd"

# --- DATE FORMATTING ---
df.withColumn("formatted", date_format(col("booking_date"), "MMM dd, yyyy"))
# "2024-03-15" -> "Mar 15, 2024"
df.withColumn("year_month", date_format(col("booking_date"), "yyyy-MM"))

# --- DATE ARITHMETIC ---
df.withColumn("next_week", date_add(col("booking_date"), 7))
df.withColumn("last_week", date_sub(col("booking_date"), 7))
df.withColumn("days_diff", datediff(col("end_date"), col("start_date")))
df.withColumn("months_diff", months_between(col("end_date"), col("start_date")))

# --- EXTRACTING PARTS ---
df.withColumn("year", year(col("booking_date")))
df.withColumn("month", month(col("booking_date")))
df.withColumn("day", dayofmonth(col("booking_date")))
df.withColumn("dow", dayofweek(col("booking_date")))     # 1=Sunday, 7=Saturday
df.withColumn("quarter", quarter(col("booking_date")))
df.withColumn("week", weekofyear(col("booking_date")))

# --- TRUNCATION ---
df.withColumn("month_start", trunc(col("booking_date"), "month"))   # first day of month
df.withColumn("year_start", trunc(col("booking_date"), "year"))     # first day of year
df.withColumn("hour_start", date_trunc("hour", col("timestamp_col")))

# --- LAST DAY OF MONTH ---
df.withColumn("month_end", last_day(col("booking_date")))

# --- NEXT SPECIFIC DAY ---
df.withColumn("next_monday", next_day(col("booking_date"), "Monday"))

from pyspark.sql.functions import explode, explode_outer, col, flatten, map_keys, map_values

# --- NESTED STRUCT ---
# Schema: address STRUCT<city STRING, country STRING, zip STRING>
df.select(
    col("id"),
    col("address.city").alias("city"),
    col("address.country").alias("country"),
    col("address.zip").alias("zip")
)

# --- ARRAY COLUMN ---
# Schema: tags ARRAY<STRING>
# explode: creates one row per element (NULL arrays -> zero rows)
df.withColumn("tag", explode(col("tags"))) \
  .select("id", "tag")

# explode_outer: creates one row per element (NULL arrays -> one row with NULL)
df.withColumn("tag", explode_outer(col("tags"))) \
  .select("id", "tag")

# size: number of elements in array
from pyspark.sql.functions import size
df.withColumn("tag_count", size(col("tags")))

# contains: check if array contains value
from pyspark.sql.functions import array_contains
df.filter(array_contains(col("tags"), "business"))

# flatten: flatten nested arrays
from pyspark.sql.functions import flatten
df.withColumn("flat_tags", flatten(col("nested_tags")))

# --- MAP COLUMN ---
# Schema: metadata MAP<STRING, STRING>
# Get value by key:
df.withColumn("status", col("metadata")["status"])
df.withColumn("status", col("metadata").getItem("status"))

# Explode map into key-value rows:
df.select("id", explode("metadata").alias("key", "value"))

# --- READING NESTED JSON ---
json_data = """
{"booking_id": "B001", "customer": {"name": "Alice", "email": "a@b.com"}, "tags": ["biz", "premium"]}
"""
df = spark.read.json(sc.parallelize([json_data]))

# Flatten it:
df_flat = df.select(
    col("booking_id"),
    col("customer.name").alias("customer_name"),
    col("customer.email").alias("customer_email"),
    explode(col("tags")).alias("tag")
)

# repartition(N) — WIDE transformation (full shuffle)
df.repartition(200)           # increase OR decrease partitions
df.repartition(200, "col")    # hash partition by column (for joins!)
df.repartition(200, "col1", "col2")  # hash by multiple columns

# coalesce(N) — NARROW transformation (no shuffle)
df.coalesce(10)               # only DECREASE partitions (merge adjacent)
                              # cannot increase partition count

# The static default is commonly 200, but size it from observed shuffle bytes.
# At 100-200 MB per partition, 200 partitions covers roughly 20-40 GB.

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

# TOO FEW partitions (e.g., 10 for 1 TB data):
#   → Each partition too large → OOM, slow, no parallelism

# TOO MANY partitions (e.g., 200 for 1 MB data):
#   → Up to 200 tiny tasks/files, plus scheduler and file-system overhead

# AQE automatically adjusts this! (see Part 5)

# Check current setting:
spark.conf.get("spark.sql.shuffle.partitions")

from pyspark.sql.functions import broadcast

# Method 1: Hint in code (override Catalyst's decision)
result = large_df.join(broadcast(small_df), "customer_id")

# Method 2: SQL hint
spark.sql("""
    SELECT /*+ BROADCAST(countries) */ *
    FROM orders o
    JOIN countries c ON o.country_code = c.code
""")

# Threshold: auto-broadcast tables smaller than this
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50m")  # 50 MB
# Set to -1 to DISABLE auto-broadcast (for testing SMJ)

# When to manually broadcast:
# 1. Table is >10 MB but <driver_memory (auto-threshold is too conservative)
# 2. AQE hasn't kicked in yet (early in job)
# 3. You know table is small but stats aren't computed (no ANALYZE TABLE)

# ⚠️ DANGER: Broadcasting a 10 GB table → Driver/Executor OOM

# Enable AQE (default ON in Spark 3.2+)
spark.conf.set("spark.sql.adaptive.enabled", "true")

spark = SparkSession.builder \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64m") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5") \
    .config("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256m") \
    .getOrCreate()

# Method 1: Spark UI → Stages → look for one task MUCH slower than median
# "Task Duration: median=5s, max=500s" → SKEW

# Method 2: Check key distribution
df.groupBy("country_code").count().orderBy(col("count").desc()).show(20)
# "US: 50M rows, UK: 2M, FR: 1M" → US is skewed key

# Method 3: Check partition sizes
from pyspark.sql.functions import spark_partition_id, count
df.withColumn("pid", spark_partition_id()) \
  .groupBy("pid").count() \
  .orderBy(col("count").desc()) \
  .show(20)
# Large variance in count → skew

from pyspark.sql.functions import col, concat, lit, monotonically_increasing_id, rand
import math

# Step 1: Add salt to the large table's skewed key
SALT_FACTOR = 10
large_df_salted = large_df.withColumn(
    "salted_key",
    concat(col("country_code"), lit("_"), (rand() * SALT_FACTOR).cast("int").cast("string"))
)

# Step 2: Explode the small table to match all salt values
from pyspark.sql.functions import array, explode

small_df_exploded = small_df.withColumn(
    "salt", array([lit(i) for i in range(SALT_FACTOR)])
).withColumn("salt", explode("salt")) \
 .withColumn("salted_key",
     concat(col("country_code"), lit("_"), col("salt").cast("string")))

# Step 3: Join on salted key
result = large_df_salted.join(small_df_exploded, "salted_key", "inner")

# Step 4: Drop salt columns
result = result.drop("salted_key", "salt")

# If the skewed join is small-large, just broadcast the small side
result = large_df.join(broadcast(small_df), "country_code")

# Enable AQE skew join handling (splits hot partitions automatically)
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
# No code change! AQE detects and handles at runtime.

# Phase 1: Group by (key, salt) to partially aggregate within partitions
import math
from pyspark.sql.functions import rand, col, sum

SALT = 10
partial = df.withColumn("salt", (rand() * SALT).cast("int")) \
            .groupBy("country_code", "salt") \
            .agg(sum("amount").alias("partial_sum"))

# Phase 2: Final aggregation drops salt
final = partial.groupBy("country_code") \
               .agg(sum("partial_sum").alias("total_amount"))

from pyspark.storagelevel import StorageLevel

df.cache()                                              # MEMORY_AND_DISK
df.persist(StorageLevel.MEMORY_ONLY)                    # fastest, may recompute
df.persist(StorageLevel.MEMORY_AND_DISK)                # safest
df.persist(StorageLevel.DISK_ONLY)                      # slowest, reliable
df.persist(StorageLevel.MEMORY_ONLY_SER)                # smaller memory, slower
df.persist(StorageLevel.OFF_HEAP)                       # Tungsten managed
df.unpersist()                                          # release cache

# For RDD:
rdd.cache()
rdd.persist(StorageLevel.MEMORY_AND_DISK)
rdd.unpersist()

# Average per key using aggregateByKey (output type ≠ input type)
rdd = sc.parallelize([("a", 10), ("b", 20), ("a", 30), ("b", 40)])

seqOp  = lambda acc, val: (acc[0] + val, acc[1] + 1)   # (sum, count)
combOp = lambda a, b: (a[0] + b[0], a[1] + b[1])
result = rdd.aggregateByKey((0, 0), seqOp, combOp) \
            .mapValues(lambda x: x[0] / x[1])
# → [("a", 20.0), ("b", 30.0)]

# vs DataFrame way (preferred):
df.groupBy("key").agg(avg("value"))  # Let Catalyst handle it

# Solution 1: Coalesce before writing
df.coalesce(10).write.parquet("output/")  # reduce to 10 output files

# Solution 2: Repartition by actual data volume
target_size_mb = 128
data_size_mb = df.count() * schema.avg_row_size_mb  # estimate
n_partitions = max(1, int(data_size_mb / target_size_mb))
df.repartition(n_partitions).write.parquet("output/")

# Solution 3: Delta OPTIMIZE (compact small files into larger ones)
# OPTIMIZE delta.`/path/to/delta` ZORDER BY (col)
spark.sql("OPTIMIZE delta.`/path/` ZORDER BY (customer_id)")

# Solution 4: Hive/Spark Config (auto merge small files)
spark.conf.set("spark.sql.files.maxPartitionBytes", "134217728")  # 128 MB per partition
spark.conf.set("spark.sql.files.openCostInBytes", "4194304")      # 4 MB open cost estimate

# --- EXECUTOR MEMORY ---
spark.executor.memory = "4g"             # JVM heap for executors
spark.executor.memoryOverhead = "512m"   # Off-heap overhead (Python, native)
                                         # Default: max(executor.memory * 0.1, 384 MB)
spark.memory.fraction = 0.6              # fraction of heap for Spark (default: 0.6)
spark.memory.storageFraction = 0.5       # within Spark memory: storage vs execution

# Memory layout (4g executor):
# Total heap = 4g
# User memory = 4g × (1 - 0.6) = 1.6g  (UDFs, user data structures)
# Spark memory = 4g × 0.6 = 2.4g
#   Storage (cache) = 2.4g × 0.5 = 1.2g
#   Execution (joins, shuffles) = 2.4g × 0.5 = 1.2g (both can borrow from each other)

# --- DRIVER MEMORY ---
spark.driver.memory = "4g"               # Driver JVM heap
spark.driver.maxResultSize = "2g"        # max result from collect() etc.

# --- OFF-HEAP (Tungsten) ---
spark.memory.offHeap.enabled = "true"
spark.memory.offHeap.size = "4g"

spark.executor.cores = "4"               # cores per executor (2-5 recommended)
spark.executor.instances = "10"          # number of executors (static allocation)
# Dynamic allocation config below

# Rule of thumb:
# cores: 4-5 per executor (too many → memory contention, too few → underutilized)
# memory: 4g per core (so 4 cores → 16g executor)

spark.sql.shuffle.partitions = "200"     # default (tune: 1 partition per 100-200 MB)
spark.shuffle.compress = "true"          # compress shuffle data (default: true)
spark.shuffle.spill.compress = "true"    # compress spilled shuffle data
spark.io.compression.codec = "snappy"   # snappy (fast), lz4, zstd (best ratio)

spark.default.parallelism = "200"        # default for RDD operations (= shuffle partitions)
spark.sql.shuffle.partitions = "200"     # for DataFrame/SQL operations
spark.sql.files.maxPartitionBytes = "134217728"  # 128 MB max per input partition
spark.sql.files.openCostInBytes = "4194304"      # 4 MB estimated open cost

spark.dynamicAllocation.enabled = "true"          # auto scale executors
spark.dynamicAllocation.minExecutors = "1"
spark.dynamicAllocation.maxExecutors = "50"
spark.dynamicAllocation.initialExecutors = "5"
spark.dynamicAllocation.executorIdleTimeout = "60s"   # remove idle after 60s
spark.dynamicAllocation.schedulerBacklogTimeout = "1s" # add executor after 1s backlog
spark.shuffle.service.enabled = "true"                # required for dynamic allocation

spark.serializer = "org.apache.spark.serializer.KryoSerializer"  # faster than Java
spark.kryo.registrationRequired = "false"
spark.kryo.registrator = "com.myapp.MyKryoRegistrator"  # custom for custom classes

# Kryo vs Java serialization:
# Java: safe, works with any class, 3-4x larger output
# Kryo: faster, ~3-10x smaller, needs class registration for best performance

# G1GC is recommended (better for large heaps)
spark.executor.extraJavaOptions = "-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35 -XX:ConcGCThreads=4"
spark.driver.extraJavaOptions = "-XX:+UseG1GC"

spark = SparkSession.builder \
    .appName("ProductionJob") \
    .config("spark.executor.memory", "8g") \
    .config("spark.executor.cores", "4") \
    .config("spark.executor.memoryOverhead", "1g") \
    .config("spark.driver.memory", "4g") \
    .config("spark.driver.maxResultSize", "2g") \
    .config("spark.sql.shuffle.partitions", "400") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
    .config("spark.sql.autoBroadcastJoinThreshold", "50m") \
    .config("spark.dynamicAllocation.enabled", "true") \
    .config("spark.dynamicAllocation.minExecutors", "2") \
    .config("spark.dynamicAllocation.maxExecutors", "50") \
    .config("spark.shuffle.service.enabled", "true") \
    .config("spark.executor.extraJavaOptions", "-XX:+UseG1GC") \
    .getOrCreate()

# Dynamic Allocation = Spark requests/releases Executors based on workload

# ENABLE:
spark.dynamicAllocation.enabled = true
spark.shuffle.service.enabled = true          # REQUIRED: keeps shuffle files after executor released

# SCALING UP:
spark.dynamicAllocation.schedulerBacklogTimeout = "1s"   # add exec if backlog > 1s
spark.dynamicAllocation.sustainedSchedulerBacklogTimeout = "5s"  # keep adding every 5s

# SCALING DOWN:
spark.dynamicAllocation.executorIdleTimeout = "60s"   # release executor idle >60s

# LIMITS:
spark.dynamicAllocation.minExecutors = "2"   # always keep at least 2
spark.dynamicAllocation.maxExecutors = "100" # never exceed 100

# WHY shuffle service is required:
# When executor is released, shuffle files it wrote are LOST
# External shuffle service (running on YARN NM) keeps those files available
# Without it: releasing executors causes downstream stages to fail

# WHEN TO USE dynamic allocation:
# ✓ Shared clusters (don't hog resources when idle)
# ✓ Jobs with varying workload stages (small start, large shuffle middle)
# ✓ Interactive workloads (notebooks)

# WHEN TO PREFER STATIC:
# ✓ Consistent workload (same data size every run)
# ✓ SLA-critical jobs (dynamic scaling adds latency)

from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("LazyLab").getOrCreate()

import time

# STEP 1: create a DataFrame
print("Creating DataFrame...")
t0 = time.time()
df = spark.range(1, 10_000_000)           # ← transformation: LAZY
print(f"Create time: {time.time() - t0:.3f}s")
# EXPECTED: Create time: 0.050s   ← ALMOST ZERO! Nothing ran.

# STEP 2: add transformations
print("Adding transformations...")
t0 = time.time()
df2 = (
    df.filter("id % 2 = 0")  # ← LAZY
      .withColumn("doubled", df.id * 2)  # ← LAZY
      .filter("doubled > 100")  # ← LAZY
)
print(f"Transform time: {time.time() - t0:.3f}s")
# EXPECTED: Transform time: 0.020s   ← STILL nothing executed!

# STEP 3: trigger an action
print("Calling action (count)...")
t0 = time.time()
result = df2.count()                     # ← ACTION — NOW everything runs
print(f"Action time: {time.time() - t0:.3f}s")
print(f"Result: {result}")
# EXPECTED:
# Action time: 2.500s    ← THIS is where the work happened
# Result: 4999950

# STEP 4: see the plan that Spark built
df2.explain(True)
# You'll see:
# == Parsed Logical Plan ==
# == Analyzed Logical Plan ==
# == Optimized Logical Plan ==       ← Catalyst OPTIMIZED here
# == Physical Plan ==                ← This is what actually runs
#
# 💡 Catalyst COMBINES both filters into one and pushes it down!

# SETUP
df = spark.range(1, 1_000_000) \
          .selectExpr("id", "id % 10 AS bucket")

# STEP 1: narrow transformation — check partition count
print("Original partitions:", df.rdd.getNumPartitions())
# EXPECTED: 8  (depends on cluster config)

df_filtered = df.filter("bucket > 5")
print("After filter:", df_filtered.rdd.getNumPartitions())
# EXPECTED: 8  ← same partitions (narrow, no shuffle)

# STEP 2: wide transformation — shuffle happens
df_grouped = df.groupBy("bucket").count()
print("After groupBy:", df_grouped.rdd.getNumPartitions())
# EXPECTED: 200  ← jumped to 200! (default shuffle partitions)
#                  This number comes from spark.sql.shuffle.partitions

df_grouped.show()
# EXPECTED:
# +------+------+
# |bucket| count|
# +------+------+
# |     1|100000|
# |     2|100000|
# ...

# STEP 3: see the stage boundary in the plan
df_grouped.explain()
# EXPECTED: Look for "Exchange hashpartitioning(bucket, 200)"
# THAT is the shuffle.

# Visit Spark UI (http://localhost:4040) and look at DAG:
# Stage 0: read + filter (NARROW)
#   │
#   ▼ Exchange (SHUFFLE — disk+network)
#   │
# Stage 1: groupBy + count (NARROW again, on shuffled data)

# SETUP: one large table, one small table
big   = spark.range(1, 10_000_000).selectExpr("id AS big_id", "id % 1000 AS key")
small = spark.range(1, 100).selectExpr("id AS key", "id * 10 AS multiplier")

# STEP 1: regular join (may auto-broadcast since small is tiny)
joined = big.join(small, "key")
joined.explain()
# Look for "BroadcastHashJoin" in the plan
# → Spark auto-broadcast the small side (under 10 MB threshold)

# STEP 2: force sort-merge join (disable auto-broadcast)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)

joined = big.join(small, "key")
joined.explain()
# Look for "SortMergeJoin" in the plan
# → Both sides get SHUFFLED and SORTED — expensive even though "small" is tiny!

# STEP 3: force broadcast with hint
from pyspark.sql.functions import broadcast

joined = big.join(broadcast(small), "key")
joined.explain()
# Look for "BroadcastHashJoin" again — hint overrides threshold
# → Small table sent to every executor. NO shuffle of big table.

# STEP 4: time the difference
import time

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)  # force sort-merge
t0 = time.time()
big.join(small, "key").count()
print(f"Sort-merge: {time.time() - t0:.2f}s")

t0 = time.time()
big.join(broadcast(small), "key").count()
print(f"Broadcast:  {time.time() - t0:.2f}s")

# EXPECTED: Broadcast is 2-5x faster when one side is small.

# ❌ DANGER: on a 1B row DataFrame, this will OOM the driver
results = df.collect()              # brings ENTIRE df to driver JVM
for row in results:
    print(row)

# ✅ FIX: use show() or write or toLocalIterator
df.show(10)                         # only 10 rows to driver
df.take(100)                        # first 100 rows to driver
for row in df.toLocalIterator():    # streams rows one-by-one
    print(row)

# ❌ BAD: 50 withColumn calls → huge logical plan, slow analyzer
for col_name in columns_to_process:
    df = df.withColumn(col_name, f.upper(f.col(col_name)))

# ✅ GOOD: single select with a list of expressions
df = df.select(*[f.upper(f.col(c)).alias(c) if c in columns_to_process
                 else f.col(c) for c in df.columns])

# ❌ SLOW: row-at-a-time Python UDF, serialization per row
from pyspark.sql.functions import udf
upper_udf = udf(lambda s: s.upper())
df.withColumn("name_upper", upper_udf("name"))

# ✅ FAST: use built-in SQL functions (run in JVM, optimized)
from pyspark.sql.functions import upper
df.withColumn("name_upper", upper("name"))

# ✅ ALSO FAST: Pandas UDF (vectorized, batch-based)
import pandas as pd
from pyspark.sql.functions import pandas_udf
@pandas_udf("string")
def upper_pandas(s: pd.Series) -> pd.Series:
    return s.str.upper()

# After streaming writes, compact hourly:
df.repartition(10).write.mode("overwrite").parquet(path)

# Or use Delta Lake OPTIMIZE for Delta tables
spark.sql(f"OPTIMIZE delta.`{path}`")

from pyspark.sql.functions import rand, concat, lit
from pyspark.sql.types import IntegerType

# Add random salt to skewed side
big = big.withColumn("salt", (rand() * 10).cast(IntegerType()))
big = big.withColumn("key_salted", concat("key", lit("_"), "salt"))

# Explode small side with all possible salts
small_exploded = small.crossJoin(
    spark.range(10).toDF("salt")
).withColumn("key_salted", concat("key", lit("_"), "salt"))

# Now join on key_salted (distributes work evenly)
big.join(small_exploded, "key_salted")

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

df_expensive = df.filter(...).join(...).withColumn(...)
df_expensive.cache()                     # mark for caching (still lazy!)

# ❌ If you immediately do:
df_expensive.groupBy("a").count().show()  # triggers cache + computes
df_expensive.groupBy("b").count().show()  # reuses cache ✅
# But the first .show() had to BOTH compute AND materialize cache — slow

# ✅ GOOD: force cache warmup first
df_expensive.count()                     # materializes cache quickly
df_expensive.groupBy("a").count().show() # fast (from cache)
df_expensive.groupBy("b").count().show() # fast (from cache)

from pyspark.sql.functions import col

employees = spark.read.parquet("employees/")

result = employees \
    .filter(col("salary") > 50000) \
    .select("emp_id", "name", "dept", "salary") \
    .orderBy(col("salary").desc())

result.show()

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

result = orders.groupBy("customer_id") \
    .agg(
        count("order_id").alias("total_orders"),
        sum("amount").alias("total_spent")
    ) \
    .orderBy(col("total_spent").desc())

result.show()

from pyspark.sql.functions import count, col

# Method 1: groupBy + filter
duplicates = users.groupBy("email") \
    .agg(count("*").alias("cnt")) \
    .filter(col("cnt") > 1)

# Method 2: join back to get full rows
result = users.join(duplicates, "email", "inner") \
    .select("user_id", "email", "created_at", "cnt") \
    .orderBy("email")

result.show()

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

result = transactions \
    .withColumn("date", to_date(col("date"))) \
    .groupBy("date") \
    .agg(sum("amount").alias("daily_sales")) \
    .orderBy("date")

result.show()

from pyspark.sql.functions import col, when

result = employees.withColumn(
    "salary_band",
    when(col("salary") > 100000, "High")
    .when(col("salary") >= 50000, "Mid")
    .otherwise("Low")
)

result.show()

from pyspark.sql.types import StructType, StructField, IntegerType, StringType, DoubleType

schema = StructType([
    StructField("emp_id",  IntegerType(), nullable=True),
    StructField("name",    StringType(),  nullable=True),
    StructField("dept",    StringType(),  nullable=True),
    StructField("salary",  DoubleType(),  nullable=True)
])

df = spark.read \
    .schema(schema) \
    .option("header", True) \
    .csv("employees.csv")

result = (
    df.na.drop(subset=["emp_id"])  # drop rows where emp_id is null
      .na.fill({"salary": 0.0, "dept": "Unknown"})  # fill remaining nulls
)

result.show()

# RDD approach (classic interview question)
rdd = sc.textFile("data/text_file.txt")

word_counts = rdd \
    .flatMap(lambda line: line.lower().split()) \
    .map(lambda word: (word, 1)) \
    .reduceByKey(lambda a, b: a + b) \
    .sortBy(lambda x: x[1], ascending=False)

word_counts.take(10)

# DataFrame approach (prefer in practice)
from pyspark.sql.functions import explode, split, lower, col

df = spark.read.text("data/text_file.txt")
result = df \
    .withColumn("word", explode(split(lower(col("value")), "\\s+"))) \
    .groupBy("word") \
    .count() \
    .orderBy(col("count").desc()) \
    .limit(10)

result.show()

from pyspark.sql.functions import max, col

result = employees.groupBy("dept") \
    .agg(max("salary").alias("max_salary")) \
    .orderBy(col("max_salary").desc())

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import dense_rank, col

w = Window.partitionBy("dept").orderBy(col("salary").desc())

result = employees.withColumn("rank", dense_rank().over(w)) \
    .select("emp_id", "name", "dept", "salary", "rank") \
    .orderBy("dept", "rank")

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import dense_rank, col

w = Window.partitionBy("dept").orderBy(col("salary").desc())

result = employees \
    .withColumn("rnk", dense_rank().over(w)) \
    .filter(col("rnk") <= 3) \
    .select("dept", "name", "salary", "rnk") \
    .orderBy("dept", "rnk")

result.show()

# KEY INSIGHT: dense_rank → tied salaries get same rank, all appear
# If you used rank() → gaps: ranks 1,1,3 (skips 2)
# If you used row_number() → only 3 rows max, misses ties

from pyspark.sql.window import Window
from pyspark.sql.functions import dense_rank, col

w = Window.orderBy(col("salary").desc())

result = employees \
    .withColumn("rnk", dense_rank().over(w)) \
    .filter(col("rnk") == 2) \
    .select("salary") \
    .distinct()

# Handle case where no second salary exists
if result.count() == 0:
    result = spark.createDataFrame([(None,)], ["salary"])

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import sum, col

w = Window.orderBy("date").rowsBetween(Window.unboundedPreceding, Window.currentRow)

result = revenue \
    .withColumn("running_total", sum("amount").over(w)) \
    .select("date", "amount", "running_total") \
    .orderBy("date")

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import avg, col, round

w = Window.orderBy("date").rowsBetween(-6, Window.currentRow)  # 6 prior + current

result = daily_sales \
    .withColumn("rolling_7d_avg", round(avg("sales").over(w), 2)) \
    .select("date", "sales", "rolling_7d_avg") \
    .orderBy("date")

result.show()

from pyspark.sql.functions import col

emp = employees.alias("e")
mgr = employees.alias("m")

result = emp.join(mgr, col("e.manager_id") == col("m.emp_id"), "inner") \
    .filter(col("e.salary") > col("m.salary")) \
    .select(
        col("e.emp_id"),
        col("e.name").alias("employee"),
        col("e.salary").alias("emp_salary"),
        col("m.name").alias("manager"),
        col("m.salary").alias("mgr_salary")
    )

result.show()

# Method 1: left_anti join (cleanest, most efficient)
result = customers.join(orders, "customer_id", "left_anti") \
    .select("customer_id", "name")

# Method 2: left join + filter null
result = customers.join(orders, "customer_id", "left") \
    .filter(col("order_id").isNull()) \
    .select(customers.customer_id, customers.name)

result.show()

# ⚡ Prefer left_anti — cleaner, Catalyst handles it well

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col

w = Window.partitionBy("user_id", "event_type").orderBy(col("created_at").desc())

result = events \
    .withColumn("rn", row_number().over(w)) \
    .filter(col("rn") == 1) \
    .drop("rn")

result.show()

# ⚡ Use row_number (not rank/dense_rank) — guarantees exactly 1 row per group

from pyspark.sql.window import Window
from pyspark.sql.functions import lag, col, round

w = Window.orderBy("year_month")

result = monthly_revenue \
    .withColumn("prev_revenue", lag("revenue", 1).over(w)) \
    .withColumn(
        "mom_pct_change",
        round(
            (col("revenue") - col("prev_revenue")) / col("prev_revenue") * 100,
            2
        )
    ) \
    .select("year_month", "revenue", "prev_revenue", "mom_pct_change") \
    .orderBy("year_month")

result.show()

# ⚠️ First row: prev_revenue = null → mom_pct_change = null (expected)
# ⚠️ If prev_revenue can be 0: wrap denominator in nullif logic manually

from pyspark.sql.functions import sum

result = sales.groupBy("product").pivot("month").agg(sum("revenue"))
result.show()

# Output:
# product | Jan  | Feb  | Mar
# --------|------|------|-----
# iPhone  | 1000 | 1200 | 900

# Optimization: provide list of pivot values to avoid extra scan
months = ["Jan", "Feb", "Mar", "Apr", "May"]
result = sales.groupBy("product").pivot("month", months).agg(sum("revenue"))

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

result = (
    posts.withColumn("tag", explode(col("tags")))  # one row per tag
         .groupBy("tag")
         .agg(count("*").alias("usage_count"))
         .orderBy(col("usage_count").desc())
         .limit(5)
)

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import lag, col, datediff, count
from pyspark.sql.functions import row_number

# Step 1: Deduplicate (1 row per customer per date)
deduped = orders.select("customer_id", "order_date").distinct()

# Step 2: Assign row number per customer ordered by date
w = Window.partitionBy("customer_id").orderBy("order_date")
numbered = deduped.withColumn("rn", row_number().over(w))

# Step 3: Island key = date - rn (constant for consecutive dates)
from pyspark.sql.functions import date_sub, expr
islands = numbered.withColumn(
    "island_key",
    expr("date_sub(order_date, rn)")   # Spark SQL: date - rn = island constant
)

# Step 4: Group by island, find streaks >= 3
result = islands.groupBy("customer_id", "island_key") \
    .agg(count("*").alias("streak_len")) \
    .filter(col("streak_len") >= 3) \
    .select("customer_id") \
    .distinct()

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import lag, col, sum as _sum, unix_timestamp, when

w_order = Window.partitionBy("user_id").orderBy("event_time")

# Step 1: Flag new session start (gap > 30 min or first event)
flagged = events.withColumn(
    "prev_time", lag("event_time", 1).over(w_order)
).withColumn(
    "new_session",
    when(
        col("prev_time").isNull() |
        ((unix_timestamp("event_time") - unix_timestamp("prev_time")) > 1800),
        1
    ).otherwise(0)
)

# Step 2: Cumulative sum of new_session flags = session_id
w_cum = Window.partitionBy("user_id").orderBy("event_time") \
              .rowsBetween(Window.unboundedPreceding, Window.currentRow)

result = flagged.withColumn(
    "session_id",
    _sum("new_session").over(w_cum)
).select("user_id", "event_time", "session_id")

result.show()

from pyspark.sql.window import Window
from pyspark.sql.functions import lag, col, datediff

w = Window.orderBy("record_date")

result = weather \
    .withColumn("prev_temp", lag("temperature", 1).over(w)) \
    .withColumn("prev_date", lag("record_date", 1).over(w)) \
    .filter(
        (col("temperature") > col("prev_temp")) &
        (datediff(col("record_date"), col("prev_date")) == 1)  # must be consecutive day
    ) \
    .select("id", "record_date", "temperature")

result.show()

# ⚠️ KEY: must check datediff == 1 — data may have gaps (missing dates)
#         Without this check, you'd compare non-adjacent days

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col, count, max as _max
from pyspark.sql.functions import expr

# Step 1: Deduplicate (1 row per user per date)
deduped = logins.select("user_id", "login_date").distinct()

# Step 2: Row number per user ordered by date
w = Window.partitionBy("user_id").orderBy("login_date")
numbered = deduped.withColumn("rn", row_number().over(w))

# Step 3: island_key = date - rn (constant for consecutive sequences)
islands = numbered.withColumn(
    "island_key",
    expr("date_sub(login_date, rn)")
)

# Step 4: Count streak length per island
streak_lengths = islands.groupBy("user_id", "island_key") \
    .agg(count("*").alias("streak_len"))

# Step 5: Max streak per user
result = streak_lengths.groupBy("user_id") \
    .agg(_max("streak_len").alias("longest_streak")) \
    .orderBy(col("longest_streak").desc())

result.show()

from pyspark.sql.functions import col, count

o1 = order_items.alias("o1")
o2 = order_items.alias("o2")

result = o1.join(
    o2,
    (col("o1.order_id") == col("o2.order_id")) &
    (col("o1.product_id") < col("o2.product_id")),   # avoid duplicates + self-pairs
    "inner"
) \
.groupBy(col("o1.product_id").alias("product_a"),
         col("o2.product_id").alias("product_b")) \
.agg(count("*").alias("co_purchase_count")) \
.orderBy(col("co_purchase_count").desc()) \
.limit(5)

result.show()

# ⚡ KEY: o1.product_id < o2.product_id ensures each pair appears once
#         Without this: (A,B) and (B,A) both appear = duplicates

from pyspark.sql.functions import countDistinct, col, round, when, max as _max

# Method: MAX(CASE WHEN) per user — count users, not events
funnel = events.groupBy("user_id").agg(
    _max(when(col("event_type") == "view",         1).otherwise(0)).alias("viewed"),
    _max(when(col("event_type") == "add_to_cart",  1).otherwise(0)).alias("carted"),
    _max(when(col("event_type") == "purchase",     1).otherwise(0)).alias("purchased")
)

from pyspark.sql.functions import sum as _sum, lit

summary = funnel.agg(
    _sum("viewed").alias("total_views"),
    _sum("carted").alias("total_carted"),
    _sum("purchased").alias("total_purchased")
)

# Compute conversion rates
total_views  = summary.collect()[0]["total_views"]
total_carted = summary.collect()[0]["total_carted"]
total_purch  = summary.collect()[0]["total_purchased"]

print(f"View → Cart:     {total_carted/total_views*100:.1f}%")
print(f"Cart → Purchase: {total_purch/total_carted*100:.1f}%")
print(f"Overall:         {total_purch/total_views*100:.1f}%")

# ⚡ KEY: MAX(CASE WHEN) per user ensures each user counted once per stage
#         Even if user viewed 10 times, they count as 1 viewer

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col

w = Window.partitionBy("id", "dept").orderBy(col("updated_at").desc())

result = cdc_events \
    .withColumn("rn", row_number().over(w)) \
    .filter(col("rn") == 1) \
    .drop("rn")

result.show()

# ⚡ row_number guarantees exactly 1 row per partition (unlike rank/dense_rank)
# Always use row_number for deduplication, not rank

from pyspark.sql.functions import rand, concat, lit, col, explode, array

SALT_FACTOR = 10

# APPROACH 1: Broadcast (if countries table is small — BEST approach)
from pyspark.sql.functions import broadcast
result = orders.join(broadcast(countries), "country_code")

# APPROACH 2: Salting (when both tables are large and skewed)

# Step 1: Salt the large orders table
orders_salted = orders.withColumn(
    "salted_key",
    concat(col("country_code"), lit("_"),
           (rand() * SALT_FACTOR).cast("int").cast("string"))
)

# Step 2: Explode countries to match all salt values (0 to SALT_FACTOR-1)
countries_exploded = countries.withColumn(
    "salt", explode(array([lit(i) for i in range(SALT_FACTOR)]))
).withColumn(
    "salted_key",
    concat(col("country_code"), lit("_"), col("salt").cast("string"))
)

# Step 3: Join on salted key
result = orders_salted.join(countries_exploded, "salted_key", "inner") \
    .drop("salted_key", "salt")

result.show()

# APPROACH 3: Enable AQE (automatic, no code change)
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
# AQE detects hot partitions at runtime and splits them

from pyspark.sql.functions import broadcast

# Problem: Catalyst doesn't auto-broadcast (store_metadata stats not computed)
# Default autoBroadcastJoinThreshold = 10 MB

# SOLUTION 1: Manual broadcast hint (works immediately)
result = transactions.join(
    broadcast(store_metadata),  # forces BroadcastHashJoin, no shuffle on store_metadata
    "store_id"
)

# SOLUTION 2: Increase threshold (if you want auto-broadcast for similar tables)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "52428800")  # 50 MB

# SOLUTION 3: Cache small table + broadcast
store_metadata.cache()
store_metadata.count()  # trigger cache
result = transactions.join(broadcast(store_metadata), "store_id")

# VERIFICATION: Check join type in plan
result.explain("formatted")
# Look for: BroadcastHashJoin (not SortMergeJoin)

# IMPACT: Sort-Merge Join = 2 shuffles + 2 sorts + merge per 100M rows
#         Broadcast Hash Join = 0 shuffles + hash lookup per 100M rows
#         Speedup: often 10-50x for large fact + small dim joins

from pyspark.sql.functions import input_file_name, col, lit

# APPROACH 1: Glob pattern (all files match same schema)
df = spark.read.parquet("s3://bucket/data/year=2024/month=*/") \
    .withColumn("source_file", input_file_name())

# APPROACH 2: Read separately + unionByName (different schemas)
jan = spark.read.parquet("s3://bucket/data/year=2024/month=01/") \
    .withColumn("source_month", lit("Jan"))
feb = spark.read.parquet("s3://bucket/data/year=2024/month=02/") \
    .withColumn("source_month", lit("Feb"))
mar = spark.read.parquet("s3://bucket/data/year=2024/month=03/") \
    .withColumn("source_month", lit("Mar"))

# allowMissingColumns=True fills missing cols with null instead of failing
combined = jan.unionByName(feb, allowMissingColumns=True) \
              .unionByName(mar, allowMissingColumns=True)

combined.show()

# APPROACH 3: Python list of paths
paths = [
    "s3://bucket/2024-01/",
    "s3://bucket/2024-02/",
    "s3://bucket/2024-03/"
]
df = spark.read.parquet(*paths).withColumn("source_file", input_file_name())

from pyspark.sql.functions import col, explode_outer

# Read raw JSON (schema inferred or explicit)
raw = spark.read.option("multiLine", True).json("orders/*.json")

# raw schema:
# order_id: long
# customer: struct<id: long, name: string, city: string>
# items: array<struct<sku: string, qty: int, price: double>>

# Step 1: Explode items array (one row per item)
exploded = raw.withColumn("item", explode_outer("items"))  # outer = keep null items

# Step 2: Select and flatten nested fields with dot notation
result = exploded.select(
    col("order_id"),
    col("customer.id").alias("customer_id"),
    col("customer.name").alias("customer_name"),
    col("customer.city").alias("city"),
    col("item.sku").alias("sku"),
    col("item.qty").alias("qty"),
    col("item.price").alias("price")
)

result.show()

# ⚡ explode_outer vs explode:
#   explode → drops orders with null/empty items array
#   explode_outer → keeps the order row with null item fields (safer)

# ── STEP 1: ESTIMATE RIGHT PARTITION COUNT ──────────────────────────────────
# Rule: each partition should be ~100-200 MB in memory
# 200 GB data → 200,000 MB / 128 MB per partition = ~1,600 partitions

data_size_gb = 200
target_partition_mb = 128
n_partitions = int((data_size_gb * 1024) / target_partition_mb)  # = 1600

spark.conf.set("spark.sql.shuffle.partitions", str(n_partitions))

# ── STEP 2: TUNE EXECUTOR MEMORY ─────────────────────────────────────────────
# With 16 GB executor, 4 cores → each core gets ~4 GB working memory
# Memory layout:
#   spark.executor.memory = 14g  (leave 2g for OS + overhead)
#   spark.executor.memoryOverhead = 2g  (native/Python overhead)
#   spark.memory.fraction = 0.6  → 8.4g for Spark execution + storage
#   Each core → ~2g execution memory for shuffle/joins

# ── STEP 3: AVOID CACHING 200 GB ──────────────────────────────────────────────
# DO NOT: df.cache()  ← 200 GB won't fit, causes eviction thrashing
# DO: Process in a streaming fashion (pipeline partitions through)

# ── STEP 4: FILTER EARLY — reduce data before joins ──────────────────────────
df = (
    spark.read.parquet("s3://bucket/data/")
         .filter(col("date") >= "2024-01-01")  # push filter to Parquet reader
         .filter(col("region") == "US")  # reduce data ASAP
         .select("id", "amount", "date")  # column pruning
)

# ── STEP 5: AVOID WIDE OPERATIONS ON FULL 200 GB ──────────────────────────────
# Bad: join 200 GB table with another 200 GB table (400 GB shuffle)
# Good: pre-aggregate one side first, then join reduced result
pre_agg = large_df.groupBy("customer_id") \
    .agg(sum("amount").alias("total"))   # 200 GB → maybe 10 GB after aggregation
result = pre_agg.join(customers, "customer_id")  # join 10 GB, not 200 GB

# ── STEP 6: ENABLE AQE ───────────────────────────────────────────────────────
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")
# AQE will auto-split/merge partitions based on actual sizes

# ── STEP 7: WRITE IN PARTITIONS, NOT COLLECT ─────────────────────────────────
# NEVER: result.collect()  ← pulls 200 GB to Driver → Driver OOM
# DO: result.write.parquet("s3://output/")  ← each executor writes its partitions

# ── STEP 8: USE DISK SPILL (not OOM) ─────────────────────────────────────────
# If a shuffle really must exceed memory, configure spill to disk:
spark.conf.set("spark.executor.memory", "14g")
spark.conf.set("spark.memory.fraction", "0.6")
# Execution pool will spill to disk automatically (slowdown vs OOM)

# ── FULL CONFIG FOR THIS SCENARIO ────────────────────────────────────────────
spark = (
    SparkSession.builder
        .config("spark.executor.memory", "14g")
        .config("spark.executor.cores", "4")
        .config("spark.executor.memoryOverhead", "2g")
        .config("spark.sql.shuffle.partitions", "1600")
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        # 128 MB advisory shuffle and input partition sizes
        .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "134217728")
        .config("spark.sql.files.maxPartitionBytes", "134217728")
        .getOrCreate()
)

# ── THE 3 READ MODES ─────────────────────────────────────────────────────────
#
# PERMISSIVE  (default): Parse what you can, set corrupt record to null
#                        Adds _corrupt_record column with the bad raw line
# DROPMALFORMED:         Silently drop bad rows, keep good ones
# FAILFAST:             Throw exception on first bad record (fail the job)

# ── MODE 1: PERMISSIVE — capture bad rows for inspection ─────────────────────
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, DoubleType

schema = StructType([
    StructField("id",     IntegerType(),  nullable=True),
    StructField("name",   StringType(),   nullable=True),
    StructField("amount", DoubleType(),   nullable=True),
    StructField("_corrupt_record", StringType(), nullable=True)  # captures bad rows
])

df = spark.read \
    .schema(schema) \
    .option("mode", "PERMISSIVE") \
    .option("columnNameOfCorruptRecord", "_corrupt_record") \
    .option("header", True) \
    .csv("data/vendor_feed.csv")

# Separate good and bad records
good_records = df.filter(col("_corrupt_record").isNull()).drop("_corrupt_record")
bad_records  = df.filter(col("_corrupt_record").isNotNull())

print(f"Good rows: {good_records.count()}")
print(f"Bad rows:  {bad_records.count()}")

# Save bad records for investigation / reprocessing
bad_records.write.mode("overwrite").json("s3://bucket/bad-records/")

# ── MODE 2: DROPMALFORMED — silently drop bad rows ────────────────────────────
df_clean = spark.read \
    .schema(schema) \
    .option("mode", "DROPMALFORMED") \
    .option("header", True) \
    .csv("data/vendor_feed.csv")
# Good for: non-critical data where a few bad rows are acceptable

# ── MODE 3: FAILFAST — crash on any bad row ───────────────────────────────────
df_strict = spark.read \
    .schema(schema) \
    .option("mode", "FAILFAST") \
    .option("header", True) \
    .csv("data/vendor_feed.csv")
# Good for: financial/compliance data where ANY bad row = stop everything

# ── BADRECORDSPATH — Spark 2.2+ auto-save bad records ────────────────────────
df = spark.read \
    .schema(schema) \
    .option("badRecordsPath", "s3://bucket/bad-records/") \
    .option("header", True) \
    .csv("data/vendor_feed.csv")
# Spark automatically writes all bad records + reasons to badRecordsPath
# Each bad record file: { "path": "...", "reason": "...", "record": "..." }

# ── JSON-SPECIFIC: multiline + corrupt handling ───────────────────────────────
df_json = (
    spark.read
         .option("mode", "PERMISSIVE")
         .option("columnNameOfCorruptRecord", "_corrupt_record")
         .option("multiLine", True)  # JSON arrays/objects spanning multiple lines
         .json("data/*.json")
)

# ── PARQUET corrupt file handling ─────────────────────────────────────────────
spark.conf.set("spark.sql.files.ignoreCorruptFiles", "true")   # skip corrupt files
spark.conf.set("spark.sql.files.ignoreMissingFiles", "true")   # skip missing files
df_parquet = spark.read.parquet("s3://bucket/data/")
# Great for: reading historical S3 data where some files may be deleted/corrupt

# ── NULL HANDLING AFTER READ ──────────────────────────────────────────────────
# After PERMISSIVE read: types that failed cast → null
df_clean = (
    df.filter(col("id").isNotNull())  # drop rows where id failed to parse
      .na.fill({"amount": 0.0, "name": "UNKNOWN"})  # fill other nulls
)

# ── METHOD 1: recursiveFileLookup (Spark 3.0+) ────────────────────────────────
# Reads ALL files recursively under the root path — ignores partition structure
df = spark.read \
    .option("recursiveFileLookup", "true") \
    .parquet("s3://bucket/data/")
# Reads every .parquet file in the entire /data/ tree recursively

# ── METHOD 2: Glob patterns (most flexible) ───────────────────────────────────
# All files 2 levels deep:
df = spark.read.parquet("s3://bucket/data/*/*/")

# All files any depth (Hadoop glob):
df = spark.read.parquet("s3://bucket/data/")  # Spark auto-recurses for Parquet dirs

# Specific year range:
df = spark.read.parquet("s3://bucket/data/2024/*/")

# Multiple specific months:
df = spark.read.parquet(
    "s3://bucket/data/2024/01/",
    "s3://bucket/data/2024/02/",
    "s3://bucket/data/2024/03/"
)

# ── METHOD 3: pathGlobFilter — read only specific file types ──────────────────
# In a mixed directory (CSVs + Parquets + JSONs), read only CSVs:
df = spark.read \
    .option("pathGlobFilter", "*.csv") \
    .option("recursiveFileLookup", "true") \
    .csv("s3://bucket/mixed-data/")

# Read only files matching a date pattern:
df = spark.read \
    .option("pathGlobFilter", "2024-01-*.parquet") \
    .parquet("s3://bucket/data/")

# ── METHOD 4: modifiedAfter / modifiedBefore — time-based filtering ───────────
# Read only files modified in the last 24 hours (Spark 3.4+):
df = spark.read \
    .option("modifiedAfter", "2024-01-01T00:00:00") \
    .option("modifiedBefore", "2024-01-02T00:00:00") \
    .option("recursiveFileLookup", "true") \
    .parquet("s3://bucket/data/")

# ── METHOD 5: Python glob → collect paths → pass to Spark ────────────────────
import boto3, os

# For S3: list all parquet files recursively
s3 = boto3.client("s3")
bucket, prefix = "my-bucket", "data/2024/"
all_paths = [
    f"s3://{bucket}/{obj['Key']}"
    for obj in s3.list_objects_v2(Bucket=bucket, Prefix=prefix)["Contents"]
    if obj["Key"].endswith(".parquet")
]

df = spark.read.parquet(*all_paths)

# For local filesystem:
import glob
all_paths = glob.glob("/data/**/*.parquet", recursive=True)  # recursive=True is key!
df = spark.read.parquet(*all_paths)

# ── ADD SOURCE FILE TRACKING ──────────────────────────────────────────────────
from pyspark.sql.functions import input_file_name

df = spark.read \
    .option("recursiveFileLookup", "true") \
    .parquet("s3://bucket/data/") \
    .withColumn("source_file", input_file_name())

df.show(truncate=False)
# Each row knows exactly which file it came from

# ── VALIDATE WHAT WAS READ ────────────────────────────────────────────────────
df.select("source_file").distinct().show(100, truncate=False)
# Shows all unique file paths that contributed to the DataFrame

# ── DIAGNOSE THE PROBLEM ──────────────────────────────────────────────────────
# collect_list() = collects ALL values per key into memory on ONE executor
# If "US" has 400M events → collect_list creates 400M element list in 1 executor
# That's the OOM culprit, not groupBy itself

# ── SOLUTION 1: Don't use collect_list — aggregate instead ────────────────────
# Bad (OOM):
df.groupBy("country").agg(collect_list("event"))

# Good (if you need count/sum/stats, not the full list):
from pyspark.sql.functions import count, countDistinct, sum
df.groupBy("country").agg(
    count("event").alias("event_count"),
    countDistinct("event").alias("unique_events")
)

# ── SOLUTION 2: Use collect_set with limit ────────────────────────────────────
from pyspark.sql.functions import slice, collect_list, col

# Sample first N values per key (avoid unbounded collection)
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

w = Window.partitionBy("country").orderBy("event_time")
df_limited = df.withColumn("rn", row_number().over(w)) \
               .filter(col("rn") <= 1000)   # max 1000 events per country

result = df_limited.groupBy("country").agg(collect_list("event").alias("top_events"))

# ── SOLUTION 3: Increase shuffle partitions to reduce per-partition size ───────
spark.conf.set("spark.sql.shuffle.partitions", "2000")  # more, smaller partitions
spark.conf.set("spark.executor.memory", "16g")
# More partitions = less data per task = less memory per task

# ── SOLUTION 4: Pre-filter to reduce data volume ─────────────────────────────
# Before the groupBy, eliminate data you don't need
df_filtered = df.filter(col("date") >= "2024-01-01") \
                .filter(col("is_valid") == True) \
                .select("country", "event", "event_time")   # column pruning
result = df_filtered.groupBy("country").agg(count("event"))

# ── SOLUTION 5: Use aggregateByKey (RDD) for full control ─────────────────────
# When DataFrame API can't express what you need
rdd = df.select("country", "event").rdd.map(lambda r: (r.country, r.event))

# Pre-aggregate locally: build set of unique events per partition
from pyspark.sql.functions import col

result_rdd = rdd.aggregateByKey(
    set(),                                           # initial accumulator = empty set
    lambda acc, val: acc | {val},                    # within partition: union set
    lambda acc1, acc2: acc1 | acc2                   # cross partition: union sets
)
result = spark.createDataFrame(
    result_rdd.map(lambda x: (x[0], list(x[1]))),
    ["country", "unique_events"]
)

# ── SOLUTION 6: Enable AQE skew handling ──────────────────────────────────────
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "3")
# AQE splits the "US" partition into smaller sub-partitions automatically