13 Transformation, Loading, and ETL Pipelines
Note
Continuation of the previous chapter. Assumes you have the DataFrame df loaded from DBFS.
13.1 1. Transformation (Transform)
13.1.1 1.1 Column Functions (pyspark.sql.functions)
from pyspark.sql.functions import (
col, lit, when, coalesce,
to_timestamp, date_format, hour, dayofweek, month,
regexp_replace, trim, upper, lower,
round as spark_round, abs as spark_abs,
log1p, sqrt,
sum as spark_sum, avg, count, countDistinct,
lag, lead, rank, dense_rank,
window
)13.1.2 1.2 Data Cleaning at Scale
# Filter out invalid trip records (negative distances, negative fares, etc.)
df_clean = (df
.filter(col("trip_distance") > 0)
.filter(col("fare_amount") > 0)
.filter(col("total_amount") > 0)
.filter(col("passenger_count").between(1, 6))
# Remove extreme outliers (99.9th percentile ≈ $120 USD total fare)
.filter(col("total_amount") < 500)
)
# Impute missing values with default defaults
df_clean = df_clean.fillna({
"passenger_count": 1,
"trip_distance": 0.0,
})
# Feature engineering and derived columns
df_feat = (df_clean
.withColumn("pickup_hour", hour("tpep_pickup_datetime"))
.withColumn("day_of_week", dayofweek("tpep_pickup_datetime"))
.withColumn("month", month("tpep_pickup_datetime"))
.withColumn("duration_min",
(col("tpep_dropoff_datetime").cast("long") -
col("tpep_pickup_datetime").cast("long")) / 60)
.withColumn("fare_per_km",
when(col("trip_distance") > 0,
spark_round(col("fare_amount") / col("trip_distance"), 2))
.otherwise(None))
.withColumn("tip_percentage",
spark_round(col("tip_amount") / col("fare_amount") * 100, 1))
)13.1.3 1.3 Aggregations and Window Functions
from pyspark.sql import Window
# Simple aggregation: hourly revenue totals
hourly_revenue = (df_feat
.groupBy("pickup_hour")
.agg(
count("*").alias("trip_count"),
avg("total_amount").alias("avg_ticket"),
spark_sum("total_amount").alias("total_revenue")
)
.orderBy("pickup_hour")
)
hourly_revenue.show(24)
# Window function: rank pickup zones by daily revenue
wspec = Window.partitionBy("month","day_of_week").orderBy(col("total_revenue").desc())
ranking = (df_feat
.groupBy("month","day_of_week","PULocationID")
.agg(spark_sum("total_amount").alias("total_revenue"))
.withColumn("rank", dense_rank().over(wspec))
.filter(col("rank") <= 5)
.orderBy("month","day_of_week","rank")
)
ranking.show(20)13.1.4 1.4 Joins at Scale
# Read small lookup table (fits entirely in memory → broadcast join)
from pyspark.sql.functions import broadcast
zones = spark.read.csv("data/taxi_zone_lookup.csv",
header=True, inferSchema=True)
# Broadcast join: Spark distributes small table to all executor nodes
df_with_zone = df_feat.join(
broadcast(zones),
df_feat.PULocationID == zones.LocationID,
"left"
)
# For large tables: sort-merge join (Spark's default strategy)
df_large_join = df_feat.join(other_large_table, on="join_key", how="inner")
Warning
Avoid Cartesian Products. Unconditioned joins on large DataFrames create N×M rows and exhaust cluster memory. Spark will raise an error unless .crossJoin() is explicitly invoked.
13.1.5 1.5 Exercise E-2 — Complete Cleaning Pipeline
13.2 2. SQL Transformations with Spark
Spark SQL allows running standard ANSI SQL queries over DataFrames:
# Register DataFrame as a temporary view
df_feat.createOrReplaceTempView("trips")
# SQL query execution — identical to queries written in BigQuery or Snowflake
result = spark.sql("""
SELECT
pickup_hour,
day_of_week,
COUNT(*) AS trip_count,
ROUND(AVG(total_amount), 2) AS avg_ticket,
ROUND(AVG(tip_percentage), 1) AS avg_tip_pct,
ROUND(SUM(total_amount), 0) AS total_revenue
FROM trips
WHERE total_amount BETWEEN 2.5 AND 200
GROUP BY pickup_hour, day_of_week
ORDER BY pickup_hour, day_of_week
""")
result.show(48)
Tip
If you already know ANSI SQL, you can begin querying with Spark SQL right away. DataFrame methods and SQL queries compile down to identical Catalyst physical plans.
13.3 3. Loading (Load)
13.3.1 3.1 Writing to Parquet (Standard Lakehouse Format)
# Basic overwrite write
df_feat.write.mode("overwrite").parquet("output/clean_trips/")
# Partitioned write — essential for query partition pruning
df_feat.write\
.mode("overwrite")\
.partitionBy("month", "day_of_week")\
.parquet("output/trips_by_date/")
# Control number of output files on disk
df_feat.coalesce(8).write.mode("overwrite").parquet("output/trips_8files/")13.3.2 3.2 Delta Lake — The Lakehouse Format with ACID
from delta import configure_spark_with_delta_pip, DeltaTable
spark = configure_spark_with_delta_pip(
SparkSession.builder.appName("ETL-Delta")
.config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
.config("spark.sql.catalog.spark_catalog",
"org.apache.spark.sql.delta.catalog.DeltaCatalog")
).getOrCreate()
# Initial table write
df_feat.write.format("delta").mode("overwrite").save("output/delta/trips/")
# Upsert (MERGE) — core pattern for incremental data loading
delta_table = DeltaTable.forPath(spark, "output/delta/trips/")
new_data = spark.read.parquet("data/yellow_tripdata_2023-02.parquet")
(delta_table.alias("target")
.merge(new_data.alias("source"),
"target.VendorID = source.VendorID AND "
"target.tpep_pickup_datetime = source.tpep_pickup_datetime")
.whenMatchedUpdateAll()
.whenNotMatchedInsertAll()
.execute())
# Time Travel — Query past versions of the table
df_v0 = spark.read.format("delta").option("versionAsOf", 0)\
.load("output/delta/trips/")13.3.3 3.3 Writing to a Relational Database (JDBC)
(df_feat.write
.format("jdbc")
.option("url", "jdbc:postgresql://dw-host:5432/datawarehouse")
.option("dbtable", "facts.taxi_trips")
.option("user", "etl_user")
.option("password", "secret")
.option("batchsize", 10000) # rows per batch INSERT
.option("numPartitions", 8) # parallel connections
.mode("append")
.save())13.4 4. End-to-End ETL Pipeline
Connecting Extract, Transform, and Load into a reusable pipeline function:
from pyspark.sql import SparkSession, DataFrame
from pyspark.sql.functions import (
col, hour, dayofweek, month, when,
spark_round as _round
)
import logging
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("ETL-Taxi")
def extract(spark: SparkSession, path: str) -> DataFrame:
log.info(f"Extracting raw data from {path}")
return spark.read.parquet(path)
def transform(df: DataFrame) -> DataFrame:
log.info("Transforming datasets...")
return (df
.filter(col("trip_distance") > 0)
.filter(col("fare_amount").between(2.5, 500))
.filter(col("passenger_count").between(1, 6))
.fillna({"passenger_count": 1})
.withColumn("pickup_hour", hour("tpep_pickup_datetime"))
.withColumn("day_of_week", dayofweek("tpep_pickup_datetime"))
.withColumn("month", month("tpep_pickup_datetime"))
.withColumn("duration_min",
(col("tpep_dropoff_datetime").cast("long") -
col("tpep_pickup_datetime").cast("long")) / 60)
.filter(col("duration_min").between(1, 180))
.withColumn("tip_percentage",
when(col("fare_amount") > 0,
_round(col("tip_amount")/col("fare_amount")*100, 1))
.otherwise(0.0))
)
def load(df: DataFrame, output_path: str) -> None:
log.info(f"Loading cleaned output into {output_path}")
(df.write
.mode("overwrite")
.partitionBy("month", "day_of_week")
.parquet(output_path))
log.info("ETL pipeline execution complete!")
if __name__ == "__main__":
spark = (SparkSession.builder
.appName("ETL-NYC-Taxi")
.master("local[*]")
.config("spark.sql.shuffle.partitions", "16")
.getOrCreate())
raw = extract(spark, "data/yellow_tripdata_2023-01.parquet")
clean = transform(raw)
load(clean, "output/taxi_etl/")
spark.stop()