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