Pyspark Etl

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

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.

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.

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.

# 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.

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.

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.

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.

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:

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).