490 lines
15 KiB
Markdown
490 lines
15 KiB
Markdown
# 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
|
|
|
|
```python
|
|
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`.
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
df4 = spark.read.csv("data/dask_data_*.csv", header=True, inferSchema=True)
|
|
```
|
|
|
|
3. Optionally inspect metadata about the DataFrame:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
query = "SELECT TABLE_NAME FROM information_schema.TABLES"
|
|
```
|
|
|
|
2. Read the results into `df5`:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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`:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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*
|
|
|
|
```python
|
|
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]
|
|
}
|
|
```
|
|
|
|
```python
|
|
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`:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
df = df.withColumn("CharacterSeq", F.regexp_replace("CharacterSeq", r"\*{4}", "____"))
|
|
```
|
|
|
|
0 / 1 checks: Use a regex on `CharacterSeq` to replace `****` with `____`.
|
|
|
|
### (Optional) DataFrame Filtering
|
|
|
|
```python
|
|
# 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`:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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
|
|
|
|
```python
|
|
# 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`:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
df_merged.createOrReplaceTempView("users")
|
|
```
|
|
|
|
2. Query it with `spark.sql()`, which returns a DataFrame:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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"`
|
|
|
|
```python
|
|
(
|
|
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
|
|
|
|
```python
|
|
df.write.mode("overwrite").parquet("data/Transactions_parquet")
|
|
```
|
|
|
|
Partitioning by a column splits the output into subfolders, which speeds up later filtered reads:
|
|
|
|
```python
|
|
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
|
|
|
|
```python
|
|
(
|
|
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.
|
|
|
|
```python
|
|
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:
|
|
|
|
```python
|
|
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. |