media_ml/etl_spark.md

Building an ETL Pipeline with PySpark and Spark SQL

Setup: pip install pyspark. Java 8, 11, or 17 must be installed.

---

Extract Phase

The extract phase is the first step in an ETL pipeline, where raw data is collected from files, databases, and APIs. In PySpark, every extract starts with a SparkSession, the entry point for DataFrames and Spark SQL.

In this task you will learn how to:

• Create a SparkSession
• Read flat files into Spark DataFrames
• Read data from a SQL database over JDBC
• Read data from REST APIs into Spark DataFrames

To get started open the extract.py file.

Task 1.0.1: Create a SparkSession

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder
    .appName("etl_pipeline")
    .config("spark.jars.packages", "com.mysql:mysql-connector-j:8.3.0")  # MySQL JDBC driver, used in Task 1.0.3
    .getOrCreate()
)

.master("local[*]") runs Spark on your machine. On a real cluster you would remove .master() and let the cluster manager supply it.

0 / 1 checks: Create a SparkSession named spark.

Task 1.0.2: Extracting Data from Flat Files

Spark reads these formats natively:

• .csv
• .json
• .parquet

Excel is not supported natively, so you would read it with Pandas and convert it (see Task 1.0.4).

For this task:

1. Read data/example.csv into df1 with spark.read.csv. Set header=True and inferSchema=True, otherwise every column is read as a string.
2. Read data/example.json into df2 with spark.read.json.
3. Read data/example.parquet into df3 with spark.read.parquet.

df1 = spark.read.csv("data/example.csv", header=True, inferSchema=True)
df2 = spark.read.json("data/example.json")
df3 = spark.read.parquet("data/example.parquet")

df1.printSchema()
df1.show(5)

0 / 1 checks: Read different file types into Spark DataFrames using spark.read.

Task 1.0.3: Loading Multiple Files with Wildcards

Spark is built for parallel processing, so reading many files at once is a core use case. Spark will split the files into partitions and process them across your cores or cluster.

1. Observe the dask_data_*.csv files in the data folder.
2. Assign df4 the result of:

df4 = spark.read.csv("data/dask_data_*.csv", header=True, inferSchema=True)

3. Optionally inspect metadata about the DataFrame:

print(df4.count())                  # row count (triggers a job)
print(len(df4.columns))             # column count
print(df4.rdd.getNumPartitions())   # how many partitions Spark created

0 / 1 checks: Read the dask_data files in data into the df4 DataFrame.

Task 1.0.4: Extract Data from Database Tables (JDBC)

Spark reads databases through JDBC. Instead of a SQLAlchemy engine, you pass connection options to spark.read.format("jdbc").

1. Assign the query to the query variable:

query = "SELECT TABLE_NAME FROM information_schema.TABLES"

2. Read the results into df5:

df5 = (
    spark.read.format("jdbc")
    .option("url", "jdbc:mysql://localhost:3306/mydb")
    .option("query", query)
    .option("user", "root")
    .option("password", "your_password")
    .option("driver", "com.mysql.cj.jdbc.Driver")
    .load()
)
df5.show()

For large tables, add partitionColumn, lowerBound, upperBound, and numPartitions so Spark reads in parallel rather than through a single connection.

0 / 1 checks: Read the provided query results into the df5 Spark DataFrame.

Task 1.0.5: Extracting REST API Data with Requests

Spark has no built-in REST reader. The common pattern is to fetch the JSON with requests on the driver, then convert it into a DataFrame.

1. Open a new terminal tab and run python3 scripts/app.py to start the Flask app.
2. Fetch the data and create a DataFrame:

import requests

url = "http://localhost:5000/random-user"
r = requests.get(url)
print(r.json())

df6 = spark.createDataFrame([r.json()])   # a dict becomes a one-row DataFrame
df6.show()

For a list of records, pass the list directly: spark.createDataFrame(r.json()).

Task 1.0.6: Authenticating with Requests

Add a Bearer token to the request headers to authenticate to the /protected endpoint.

1. In the headers dictionary, add the key Authorization.
2. Set the value to Bearer mysecrettoken123.
3. Assign the result to response:

headers = {"Authorization": "Bearer mysecrettoken123"}
response = requests.get("http://localhost:5000/protected", headers=headers)
df7 = spark.createDataFrame([response.json()])

0 / 1 checks: Assign the header values and authenticate to the /protected endpoint, assigning the results to response.

(Optional) Reading JSON strings with a schema

createDataFrame infers types, which can be fragile with nested or inconsistent API data. Define a schema explicitly for predictable results:

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

schema = StructType([
    StructField("name", StringType(), True),
    StructField("age", IntegerType(), True),
])
df8 = spark.createDataFrame(response.json(), schema=schema)

---

Transform Phase

The transform phase is where raw data is cleansed, standardized, aggregated, and joined. PySpark gives you two interchangeable APIs for this: the DataFrame API (pyspark.sql.functions) and Spark SQL. Both compile to the same optimized execution plan, so use whichever reads better.

In this task you will learn how to:

• Convert data types and date formats
• Replace values and clean dirty data with regex
• Filter, aggregate, and create calculated columns
• Join DataFrames
• Run the same logic with Spark SQL using temp views

*Example Data

data = {
    "Name": ["Adam", "Bob", "Charlie"],
    "Children": [1, 2, 0],
    "ZipCode":[10001,90001,60601],
    "City": ["New York", "Los Angeles", "Chicago"],
    "Birthdate": [pd.Timestamp('20000101'),pd.Timestamp('19900418'),pd.Timestamp('20180310')],
    "CharacterSeq": ['****','____','****'],
    "UserId":['A111','B222','C333'],
    "TShirtSize":['M','XL','XL'],
    "Salary":[50000,60000,70000]
}

from pyspark.sql import functions as F
from pyspark.sql.types import StringType, DoubleType

Task 2.0.1: Convert Column Data Types

Observe the data dictionary. Notice that ZipCode is an integer. Convert it to a string with cast:

df = df.withColumn("ZipCode", F.col("ZipCode").cast(StringType()))
df.printSchema()

DataFrames are immutable in Spark. withColumn returns a new DataFrame, so you must assign the result back.

0 / 1 checks: Convert the ZipCode column to a string data type.

Task 2.0.2: Convert Date Formats

Birthdate is in yyyy-MM-dd format. We want MM-dd-yyyy. Parse the string into a date, then format it back out:

df = df.withColumn(
    "Birthdate",
    F.date_format(F.to_date("Birthdate", "yyyy-MM-dd"), "MM-dd-yyyy")
)
df.select("Birthdate").show(5)

Spark date patterns are case-sensitive. MM is month and mm is minutes.

0 / 1 checks: Convert the Birthdate column to the MM-dd-yyyy format.

Task 2.0.3: Replace Values with a Mapping

1. Observe the replace_map dictionary. The value being replaced comes first, followed by the replacement.
2. Apply it to the City column:

replace_map = {"NYC": "New York", "LA": "Los Angeles"}
df = df.replace(replace_map, subset=["City"])

0 / 1 checks: Convert the existing City values using df.replace().

Task 2.0.4: Using Regex to Replace Dirty Data

Use regexp_replace to find and replace patterns. * is a special regex character, so it must be escaped.

1. Observe the CharacterSeq column.
2. Write a regex pattern that matches four asterisks: \*{4} (or \*\*\*\*).
3. Replace them with underscores:

df = df.withColumn("CharacterSeq", F.regexp_replace("CharacterSeq", r"\*{4}", "____"))

0 / 1 checks: Use a regex on CharacterSeq to replace **** with ____.

(Optional) DataFrame Filtering

# Column expression syntax
df.filter(F.col("Name") == "Adam")

# SQL string syntax
df.filter("Name = 'Adam'")

# Assign to a new DataFrame for a sub-dataset
adams = df.filter(F.col("Name") == "Adam")
adams.show()

where() is an alias for filter().

Task 2.0.5: Creating Custom Calculated Columns

1. Observe the cost_of_living_adjust function, which takes an original value and an inc_percent.
2. Wrap it in a UDF (user-defined function) and create Salary_COL:

def cost_of_living_adjust(value, inc_percent):
    return value * (1 + inc_percent / 100) if value is not None else None

col_udf = F.udf(lambda s: cost_of_living_adjust(s, 5), DoubleType())
df = df.withColumn("Salary_COL", col_udf(F.col("Salary")))

Python UDFs are slow because every row is serialized between the JVM and Python. When you can, use native column expressions, which Spark can optimize:

df = df.withColumn("Salary_COL", F.col("Salary") * (1 + 5 / 100))

0 / 1 checks: Create a new Salary_COL column by applying cost_of_living_adjust to Salary with inc_percent=5.

(Optional) DataFrame Aggregating

# Whole-DataFrame aggregate
df.agg(
    F.count("Salary").alias("count"),
    F.min("Salary").alias("min"),
    F.max("Salary").alias("max"),
    F.mean("Salary").alias("mean"),
    F.sum("Salary").alias("sum"),
).show()

# Grouped aggregate
df.groupBy("TShirtSize").agg(F.mean("Salary").alias("avg_salary")).show()

Task 2.0.6: Joining Data

Spark's .join() mirrors SQL joins. The join types are inner, left, right, outer, left_semi, and left_anti.

1. Observe both DataFrames (df and df2).
2. Use UserId as the on parameter.
3. Use left as the how parameter.
4. Assign the result to df_merged:

df_merged = df.join(df2, on="UserId", how="left")
df_merged.show()

Passing on="UserId" (a string) keeps a single UserId column. Joining on df.UserId == df2.UserId keeps both, which causes ambiguous-column errors later.

0 / 1 checks: Use .join to combine data in df and df2 into df_merged.

Task 2.0.7: Transforming with Spark SQL

Spark SQL lets you query DataFrames with plain SQL. You first register a DataFrame as a temporary view:

1. Register the view:

df_merged.createOrReplaceTempView("users")

2. Query it with spark.sql(), which returns a DataFrame:

result = spark.sql("""
    SELECT
        City,
        COUNT(*)                       AS user_count,
        ROUND(AVG(Salary_COL), 2)      AS avg_salary
    FROM users
    WHERE Salary_COL IS NOT NULL
    GROUP BY City
    HAVING COUNT(*) > 1
    ORDER BY avg_salary DESC
""")
result.show()

3. Because the result is a DataFrame, you can switch back to the DataFrame API at any point and mix the two styles freely.

0 / 1 checks: Register df_merged as a temp view named users and aggregate it with spark.sql().

---

Loading Data

Data writing in PySpark refers to saving a Spark DataFrame to files or databases through the df.write interface. Spark supports csv(), json(), parquet(), orc(), and jdbc(), allowing you to export data to files, data lakes, or relational databases.

There are three key differences from Pandas. Spark writes a directory of part files (one per partition) rather than a single file. You control overwrite behavior with save modes (overwrite, append, ignore, errorifexists). And Parquet is the preferred format for analytics because it is columnar, compressed, and keeps the schema.

To get started open the load.py file.

Task 3.0.1: Loading Data to a CSV File

1. Create a dictionary with the following python code:

data = {
    'Transaction_ID': ['T1001', 'T1002', 'T1003', 'T1004'],
    'Customer_Name': ['Alice Johnson', 'Bob Lee', 'Charlie Kim', 'David Clark'],
    'Product': ['Laptop', 'Smartphone', 'Tablet', 'Smartwatch'],
    'Quantity': [1, 2, 3, 1],
    'Price_Per_Unit': [1200, 800, 400, 200],
    'Total_Sale': [1200, 1600, 1200, 200],
    'Sale_Date': ['2025-03-20', '2025-03-21', '2025-03-21', '2025-03-22']
}

2. Convert the dictionary into a Spark DataFrame variable df. Spark expects a list of rows, so reshape the dictionary first:

import pandas as pd

df = spark.createDataFrame(pd.DataFrame(data))

3. Use df.write.csv to write the contents to data/Transactions. Include the following options:
   • header=True
   • mode="overwrite"

(
    df.coalesce(1)                 # collapse to one partition, so one part file
      .write
      .mode("overwrite")
      .option("header", True)
      .csv("data/Transactions")
)

Spark creates a folder named Transactions containing part-00000-....csv and a _SUCCESS marker. coalesce(1) is fine for small data, but avoid it on large datasets because it funnels everything through one task.

0 / 1 checks: Convert the dictionary values to a DataFrame and load it to Transactions in the data directory.

Task 3.0.2: Loading Data to Parquet

df.write.mode("overwrite").parquet("data/Transactions_parquet")

Partitioning by a column splits the output into subfolders, which speeds up later filtered reads:

df.write.mode("overwrite").partitionBy("Sale_Date").parquet("data/Transactions_by_date")

0 / 1 checks: Write the DataFrame to Parquet.

Task 3.0.3: Loading Data to a Database with JDBC

(
    df.write.format("jdbc")
    .option("url", "jdbc:mysql://localhost:3306/mydb")
    .option("dbtable", "transactions")
    .option("user", "root")
    .option("password", "your_password")
    .option("driver", "com.mysql.cj.jdbc.Driver")
    .mode("append")
    .save()
)

Use append to add rows. overwrite drops and recreates the table unless you also set .option("truncate", True).

0 / 1 checks: Load the DataFrame into a transactions table.

(Optional) Validating Data Operations

Use assert statements to validate that the pipeline finished successfully. Check that row counts, column types, and totals match the expected results.

csv_df = spark.read.csv("data/Transactions", header=True, inferSchema=True)

assert csv_df.count() == 4, "Row counts do not match expected values!"
print(f"Row counts match: {csv_df.count()} rows")

assert csv_df.agg(F.sum("Total_Sale")).first()[0] == 4200, "Total sales mismatch!"

expected_cols = {"Transaction_ID", "Customer_Name", "Product", "Quantity",
                 "Price_Per_Unit", "Total_Sale", "Sale_Date"}
assert set(csv_df.columns) == expected_cols, "Schema mismatch!"

To compare against the database table, read it back with spark.read.format("jdbc") and compare .count() to the source DataFrame.

---

Putting It All Together

Wrap the three phases in functions so the pipeline is testable and reusable:

def extract(spark):
    return spark.read.csv("data/example.csv", header=True, inferSchema=True)

def transform(df):
    return (
        df.withColumn("ZipCode", F.col("ZipCode").cast("string"))
          .withColumn("Birthdate", F.date_format(F.to_date("Birthdate"), "MM-dd-yyyy"))
          .withColumn("Salary_COL", F.col("Salary") * 1.05)
          .dropDuplicates()
    )

def load(df, path):
    df.write.mode("overwrite").parquet(path)

if __name__ == "__main__":
    spark = SparkSession.builder.appName("etl").master("local[*]").getOrCreate()
    raw = extract(spark)
    clean = transform(raw)
    load(clean, "data/output")
    assert spark.read.parquet("data/output").count() == clean.count()
    spark.stop()

I can also put this into a downloadable file or turn it into a Jupyter-style notebook.