Spark Engineer

from pyspark.sql import SparkSession, Window
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, TimestampType

spark = SparkSession.builder \
    .appName("etl_pipeline") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.shuffle.partitions", "auto") \
    .getOrCreate()

# Read with schema enforcement (avoid inferSchema in production)
schema = StructType([
    StructField("user_id", IntegerType(), False),
    # ... (condensed) ...
        F.sum("amount").alias("total_amount"),
        F.percentile_approx("amount", 0.5).alias("median_amount"),
    )
    .orderBy("date", "hour")
)

# Write partitioned output (Hive-style partitioning)
result.write \
    .partitionBy("year", "month", "day") \
    .mode("overwrite") \
    .parquet("s3://bucket/output/events/")

# Partition pruning: only reads relevant partitions
filtered = spark.read.parquet("s3://bucket/output/events/") \
    .filter(F.col("year") == 2024) \
    .filter(F.col("month") == 6)
# Spark reads only /year=2024/month=6/ directories

# Bucketing: pre-sort data for join optimization
events.write \
    .bucketBy(256, "user_id") \
    .sortBy("user_id", "timestamp") \
    .saveAsTable("events_bucketed")
# Joins on user_id between bucketed tables avoid shuffle

# Repartition: full shuffle, use when you need specific partitioning
df_repartitioned = df.repartition(200, "customer_id")

# Coalesce: reduce partitions without full shuffle (narrow dependency)
df_coalesced = df.coalesce(10)  # Only for reducing partition count

# Check current partitioning
print(f"Partitions: {df.rdd.getNumPartitions()}")

# Custom partitioning (RDD level)
rdd = df.rdd.partitionBy(100, lambda key: hash(key) % 100)

# Anti-pattern: multiple shuffles
result = (
    df.groupBy("user_id").agg(F.count("*").alias("cnt"))
      .filter(F.col("cnt") > 10)
      .join(user_details, "user_id")
)
# This does TWO shuffles: one for groupBy, one for join

# Optimized: pre-partition to align both operations
df_partitioned = df.repartition(200, "user_id")
result = (
    df_partitioned
    .groupBy("user_id").agg(F.count("*").alias("cnt"))
    .filter(F.col("cnt") > 10)
    .join(user_details.repartition(200, "user_id"), "user_id")
)
# Still two shuffles, but the repartition is shared

# Best: use Adaptive Query Execution (AQE) in Spark 3.x
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

from pyspark.sql.functions import broadcast

# Explicit broadcast hint
result = large_df.join(broadcast(small_df), "join_key")

# Auto-broadcast threshold (default 10MB)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50m")  # Increase to 50MB

# Check if broadcast was used
result.explain(True)
# Look for "BroadcastHashJoin" in the physical plan

# Broadcast variable for lookups (RDD-level)
lookup_dict = {"US": "United States", "UK": "United Kingdom"}
bc_lookup = spark.sparkContext.broadcast(lookup_dict)

@F.udf(StringType())
def resolve_country(code):
    return bc_lookup.value.get(code, "Unknown")

from pyspark import StorageLevel

# Cache levels (from fastest to most durable)
df.cache()                                              # MEMORY_AND_DISK (default)
df.persist(StorageLevel.MEMORY_ONLY)                    # Fastest, recompute on eviction
df.persist(StorageLevel.MEMORY_AND_DISK)                # Spill to disk
df.persist(StorageLevel.MEMORY_AND_DISK_SER)            # Serialized, less memory
df.persist(StorageLevel.DISK_ONLY)                      # No memory usage
df.persist(StorageLevel.OFF_HEAP)                       # Tungsten off-heap

# IMPORTANT: cache is lazy - trigger materialization
df.cache()
df.count()  # Forces caching

# ... (condensed) ...

# When NOT to cache:
# 1. DataFrame used only once
# 2. DataFrame is very large (will cause memory pressure)
# 3. Storage is the bottleneck (will slow down other tasks)

# Register DataFrame as temporary view
events.createOrReplaceTempView("events")
user_details.createOrReplaceTempView("users")

# Complex SQL with window functions
result = spark.sql("""
    WITH user_sessions AS (
        SELECT
            user_id,
            timestamp,
            event_type,
            LAG(timestamp) OVER (PARTITION BY user_id ORDER BY timestamp) AS prev_ts,
            CASE
                WHEN UNIX_TIMESTAMP(timestamp) -
                     # ... (condensed) ...
        COLLECT_SET(event_type) AS event_types
    FROM sessions
    GROUP BY user_id, session_id
    HAVING COUNT(*) > 1
""")

# Read from Kafka
stream_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092") \
    .option("subscribe", "events") \
    .option("startingOffsets", "latest") \
    .option("maxOffsetsPerTrigger", 100000) \
    .load()

# Parse and transform
parsed = (
    stream_df
    .selectExpr("CAST(value AS STRING) as json_str", "timestamp as kafka_ts")
    .select(
        # ... (condensed) ...
    .trigger(processingTime="30 seconds")
    .start("s3://bucket/output/windowed_events/")
)

query.awaitTermination()

# AVOID UDFs when possible - they disable Catalyst optimization

# Anti-pattern: UDF for simple logic
@F.udf(StringType())
def categorize_udf(amount):
    if amount > 1000: return "high"
    elif amount > 100: return "medium"
    return "low"

# Better: use built-in functions
df.withColumn("category",
    F.when(F.col("amount") > 1000, "high")
     .when(F.col("amount") > 100, "medium")
     .otherwise("low")
# ... (condensed) ...
    model.fit(pdf[['x1', 'x2']], pdf['y'])
    pdf['prediction'] = model.predict(pdf[['x1', 'x2']])
    return pdf

result = df.groupBy("segment").apply(train_model_per_group)

Per Executor:
  Total Memory = spark.executor.memory + spark.executor.memoryOverhead

  spark.executor.memory:
    - 300MB reserved for Spark internals
    - Remaining split: 60% execution (shuffles, joins, sorts, aggregations)
                       40% storage (cache, broadcast variables)
    - Controlled by spark.memory.fraction (default 0.6)
    - Controlled by spark.memory.storageFraction (default 0.5 of fraction)

  spark.executor.memoryOverhead:
    - Default: max(384MB, 0.10 * spark.executor.memory)
    - Increase for PySpark (Python processes), large broadcasts, or off-heap

Sizing formula:
  Data size (compressed on disk) * decompression ratio (~3-5x) * number of passes
  / target partition size (128-256MB)
  = minimum total executor memory needed

# Small job: 10-100 GB data
spark.conf.set("spark.executor.memory", "4g")
spark.conf.set("spark.executor.cores", "4")
spark.conf.set("spark.executor.instances", "10")
spark.conf.set("spark.sql.shuffle.partitions", "100")

# Medium job: 100 GB - 1 TB data
spark.conf.set("spark.executor.memory", "8g")
spark.conf.set("spark.executor.cores", "4")
spark.conf.set("spark.executor.instances", "50")
spark.conf.set("spark.sql.shuffle.partitions", "500")

# Large job: 1-10 TB data
spark.conf.set("spark.executor.memory", "16g")
# ... (condensed) ...
# Dynamic allocation (recommended for shared clusters)
spark.conf.set("spark.dynamicAllocation.enabled", "true")
spark.conf.set("spark.dynamicAllocation.minExecutors", "5")
spark.conf.set("spark.dynamicAllocation.maxExecutors", "200")
spark.conf.set("spark.dynamicAllocation.executorIdleTimeout", "120s")

# Detect skew: check partition sizes
df.groupBy(F.spark_partition_id().alias("partition")) \
  .count() \
  .describe("count") \
  .show()
# If max >> mean, you have skew

# Fix 1: Salting (add random prefix to skewed key)
salt_range = 10
df_salted = df.withColumn("salt", (F.rand() * salt_range).cast("int"))
df_salted = df_salted.withColumn("salted_key",
    F.concat(F.col("join_key"), F.lit("_"), F.col("salt"))
)

# ... (condensed) ...

# Fix 2: AQE skew join (Spark 3.0+) - automatic
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256m")

# Compaction: merge small files into optimal-sized files
df = spark.read.parquet("s3://bucket/many_small_files/")
df.coalesce(target_file_count).write \
    .mode("overwrite") \
    .parquet("s3://bucket/compacted/")

# Target file count calculation
total_size_bytes = sum(f.size for f in dbutils.fs.ls("s3://bucket/many_small_files/"))
target_file_size = 256 * 1024 * 1024  # 256 MB
target_file_count = max(1, total_size_bytes // target_file_size)

# Prevent small files on write
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "64m")

# Must-have configurations for production
configs = {
    # Adaptive Query Execution (Spark 3.0+)
    "spark.sql.adaptive.enabled": "true",
    "spark.sql.adaptive.coalescePartitions.enabled": "true",
    "spark.sql.adaptive.skewJoin.enabled": "true",

    # Serialization
    "spark.serializer": "org.apache.spark.serializer.KryoSerializer",

    # Parquet optimization
    "spark.sql.parquet.filterPushdown": "true",
    "spark.sql.parquet.mergeSchema": "false",
    "spark.hadoop.parquet.enable.summary-metadata": "false",
# ... (condensed) ...
    "spark.memory.storageFraction": "0.5",
}

for k, v in configs.items():
    spark.conf.set(k, v)

# Keep latest record per key
from pyspark.sql import Window

w = Window.partitionBy("user_id").orderBy(F.col("updated_at").desc())
deduped = (
    df.withColumn("rn", F.row_number().over(w))
      .filter(F.col("rn") == 1)
      .drop("rn")
)

# Explode: one row per array element
df.select("user_id", F.explode("tags").alias("tag"))

# Collect: aggregate back to arrays
df.groupBy("user_id").agg(
    F.collect_list("tag").alias("all_tags"),
    F.collect_set("tag").alias("unique_tags")
)

# Delta Lake: ACID transactions on data lakes
from delta.tables import DeltaTable

# Upsert (merge)
delta_table = DeltaTable.forPath(spark, "s3://bucket/delta/customers")
delta_table.alias("target").merge(
    updates_df.alias("source"),
    "target.customer_id = source.customer_id"
).whenMatchedUpdateAll() \
 .whenNotMatchedInsertAll() \
 .execute()

# Time travel
df_yesterday = spark.read.format("delta") \
    .option("timestampAsOf", "2024-06-14") \
    .load("s3://bucket/delta/customers")

# Optimize and Z-Order
spark.sql("OPTIMIZE delta.`s3://bucket/delta/customers` ZORDER BY (region, customer_id)")