--- name: pyspark-etl description: "Best practices for building performant, testable PySpark ETL pipelines with Spark SQL and Apache Iceberg. Use when writing or reviewing PySpark jobs, designing joins and window functions, working with map/array higher-order functions, or building idempotent cumulative/snapshot table merges." --- # PySpark ETL This skill covers patterns for building production-grade, testable ETL pipelines with PySpark, Spark SQL, and Apache Iceberg, including project structure, join and window-function idioms, and safe cumulative-table merge patterns. ## Workflow for Building a PySpark ETL Job 1. **Scaffold the job class** — Create a class that manages the `SparkSession` lifecycle, accepts an injectable session for testing, and exposes an abstract `run_job` method. 2. **Define config via a factory function** — Keep config as a plain dataclass; parse CLI args in a separate factory function so tests can construct configs without touching `sys.argv`. 3. **Read source data with a shared, partition-aware reader** — Use a generic reader utility for date filters, hour ranges, and latest-partition lookups; keep business filters in the ETL class. 4. **Compose the pipeline with `.transform()`** — Chain named methods (`read_source().transform(self.enrich).transform(self.merge_with_existing)`) so `run_job` stays pure orchestration. 5. **Apply transformations idiomatically** — Use `select` over `withColumn` chains, explicit join types, explicit window frames, and native functions instead of UDFs. 6. **Write with schema-evolution safety** — Use `.byName()` when writing to Iceberg tables so column order doesn't matter. 7. **Validate output** — Check primary-key uniqueness and null counts on key columns after every write. 8. **Test locally** — Unit test transformation methods against a local `SparkSession` with small, hand-built DataFrames. ## Project Structure ### ETL class scaffold ```python from abc import ABC, abstractmethod import logging from pyspark.sql import SparkSession class BaseETL(ABC): def __init__(self, config, app_name="ETL Job", spark_session=None): self.spark = spark_session or SparkSession.builder.appName(app_name).getOrCreate() self.config = config self.logger = logging.getLogger(self.__class__.__name__) @abstractmethod def run_job(self): ... def stop(self): self.spark.stop() ``` ### Config as a factory function Keep the dataclass as pure data; put CLI parsing in a standalone factory so configs are easy to build in tests. ```python import argparse from dataclasses import dataclass @dataclass class MyConfig: read_date: int = 20260101 def create_config() -> MyConfig: parser = argparse.ArgumentParser() parser.add_argument("--read_date", type=int, default=20260101) args = parser.parse_args() return MyConfig(read_date=args.read_date) ``` ### Partition-aware shared reader Build one generic reader for partition mechanics; keep domain-specific filters visible in the ETL, not buried in a one-off reader class. ```python import pyspark.sql.functions as F class PartitionedReader: @staticmethod def read_latest(spark, table_name, partition_col): row = spark.read.table(table_name).agg(F.max(partition_col)).first() if row is None or row[0] is None: return spark.createDataFrame([], spark.read.table(table_name).schema) return spark.read.table(table_name).filter(F.col(partition_col) == row[0]) @staticmethod def read_by_date(spark, table_name, partition_col, date_value): return spark.read.table(table_name).filter(F.col(partition_col) == date_value) events = PartitionedReader.read_by_date(spark, "catalog.my_table", "event_date", 20260319) events = events.filter(F.col("event_type").isin("login", "purchase")) ``` ## Code Style - Import functions as `import pyspark.sql.functions as F` and always use `F.col()` instead of `df.colA` attribute access — attribute access binds a column to a specific DataFrame variable and breaks after joins or reassignment. - Extract complex boolean logic inside `.filter()` or `F.when()` into named variables once it exceeds 3 expressions. - Prefer `select` over chains of `withColumn` — `select` states the output schema in one pass, while each `withColumn` call adds a projection to the query plan. - Use `.alias()` instead of `withColumnRenamed`. - Limit chained method calls to 5 per statement; separate `select`/`filter` chains from `withColumn` chains from `join` chains by operation type. ```python # BAD — 3 intermediate DataFrames, one projection per call df = df.withColumn("a", F.col("a").cast("double")) df = df.withColumn("b", F.upper(F.col("b"))) df = df.withColumn("c", F.lit(1)) # GOOD — one DataFrame, explicit schema contract df = df.select( F.col("a").cast("double"), F.upper(F.col("b")).alias("b"), F.lit(1).alias("c"), ) ``` ## Joins - Always specify `how=` explicitly — never rely on the default. - Prefer left joins over right joins for readability; flip DataFrame order instead of using `how="right"`. - Alias whole DataFrames for disambiguation after joins rather than renaming every column with `withColumnRenamed`. - Wrap small dimension/lookup tables in `F.broadcast()` when joining against a large fact table, especially after filters or transformations that prevent Spark from inferring the size automatically (`spark.sql.autoBroadcastJoinThreshold` only auto-broadcasts tables Spark can size, typically under 10MB). Confirm with `df.explain()` — look for `BroadcastHashJoin` vs `SortMergeJoin`. - Never reach for `.dropDuplicates()` to mask unexpected duplicate rows — find the root cause; it also adds shuffle overhead. ```python flights = flights.alias("f") parking = parking.alias("p") result = flights.join(F.broadcast(parking), "code", how="left").select( F.col("f.start_time").alias("flight_start"), F.col("p.total_time").alias("parking_total"), ) ``` ## Window Functions Use `from pyspark.sql import Window as W`. - Always specify an explicit frame — without one, `F.sum().over(w)` behaves differently depending on whether `orderBy` is present (running sum vs. total). - Know the difference between `row_number()` + filter (drops rows, keeps the best one) and `first()` over a window (overwrites a column, keeps all rows). - Pass `ignorenulls=True` to `F.first()`/`F.last()` — otherwise a null in the first row of a partition nulls the entire partition's result. - Avoid empty `partitionBy()`; it forces all data into a single partition. Use `.agg()` for global aggregations instead. ```python w = W.partitionBy("key").orderBy("num").rowsBetween(W.unboundedPreceding, W.unboundedFollowing) df = df.withColumn("version", F.first("version", ignorenulls=True).over(w)) ``` ## Map & Array Higher-Order Functions - Use `map_zip_with` instead of `map_concat` when merging maps needs per-key conflict resolution (e.g., keep the entry with the later timestamp) rather than one side blindly winning. - Use `transform` + `array_max`/`array_min` to extract values out of nested structs without a UDF. - Avoid UDFs — check for a built-in Spark function or higher-order function first. UDFs break Catalyst optimization and add serialization overhead. ```python merged = F.map_zip_with( new_map, existing_map, lambda key, v1, v2: ( F.when(v1.isNull(), v2) .when(v2.isNull(), v1) .otherwise(F.when(v1.event_ts >= v2.event_ts, v1).otherwise(v2)) ), ) ``` ## Cumulative / Snapshot Table Patterns - Merges must be idempotent — re-running with the same input data must not create duplicates. - Merges must be order-independent — backfilling old data must never overwrite newer data. Resolve conflicts with an explicit criterion (event timestamp, version number, partition date), never positional precedence like `coalesce` argument order. - Validate primary-key uniqueness and null counts on key columns as an audit step after every write. ## Data Quality & Performance - Use `F.lit(None)` for empty columns — never empty strings or sentinel values like `"NA"`. - Avoid `.otherwise()` as a catch-all in `F.when()` chains for categorical mappings; an unmapped value should surface as null, not silently collapse into "Other". - Never leave `.show()`, `.collect()`, or `.printSchema()` in production code — they force full materialization or add driver overhead. `.count()` is fine when used intentionally for row-count logging or to force materialization before a DAG fork. - Use `.persist()` only when a DataFrame is referenced by multiple subsequent actions. Choose the storage level deliberately: `MEMORY_AND_DISK` (safe default), `MEMORY_ONLY` (fastest, risks recompute on eviction), `DISK_ONLY` (for DataFrames too large for memory). ## Iceberg Write Patterns - Use `.byName()` when writing so Spark matches columns by name, not position — this keeps writes safe across schema evolution. ```python df.write.byName().mode("overwrite").insertInto("catalog.my_table") ``` - Use the `__partitions` Iceberg metadata table to find the latest snapshot instead of scanning the full table: ```python partition_df = spark.read.table("catalog.my_table__partitions").select( "partition.partition_date", "partition.partition_hour" ) max_partition = partition_df.orderBy( F.col("partition_date").desc(), F.col("partition_hour").desc() ).first() if max_partition is None: raise ValueError("No partitions found in catalog.my_table") ``` - Choose `write.distribution-mode` deliberately: `"none"` (fastest, no re-shuffle, file sizes depend on upstream partitioning), `"hash"` (shuffles by partition key for evenly sized files), `"range"` (sorts before writing, best scan performance but most expensive).