15 KiB
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:
- Read
data/example.csvintodf1withspark.read.csv. Setheader=TrueandinferSchema=True, otherwise every column is read as a string. - Read
data/example.jsonintodf2withspark.read.json. - Read
data/example.parquetintodf3withspark.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.
- Observe the
dask_data_*.csvfiles in thedatafolder. - Assign
df4the result of:
df4 = spark.read.csv("data/dask_data_*.csv", header=True, inferSchema=True)
- 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").
- Assign the query to the
queryvariable:
query = "SELECT TABLE_NAME FROM information_schema.TABLES"
- 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.
- Open a new terminal tab and run
python3 scripts/app.pyto start the Flask app. - 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.
- In the
headersdictionary, add the keyAuthorization. - Set the value to
Bearer mysecrettoken123. - 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
- Observe the
replace_mapdictionary. The value being replaced comes first, followed by the replacement. - Apply it to the
Citycolumn:
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.
- Observe the
CharacterSeqcolumn. - Write a regex pattern that matches four asterisks:
\*{4}(or\*\*\*\*). - 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
- Observe the
cost_of_living_adjustfunction, which takes an original value and aninc_percent. - 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.
- Observe both DataFrames (
dfanddf2). - Use
UserIdas theonparameter. - Use
leftas thehowparameter. - 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:
- Register the view:
df_merged.createOrReplaceTempView("users")
- 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()
- 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
- 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']
}
- 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))
- Use
df.write.csvto write the contents todata/Transactions. Include the following options:header=Truemode="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.