PART 6 · PYTHON FOR DATA ENGINEERING · CHAPTER 35

From pandas to PySpark, scheduled with Airflow

FreeLesson 72 of 78

When the data no longer fits on one machine, the same ideas move to PySpark. When it has to run every night, Airflow runs it.

EACH CUSTOMER'S BIGGEST DELIVERED ORDERPYSPARK 4 · LOCAL

from pyspark.sql import SparkSession, Window, functions as F

spark = (SparkSession.builder.master("local[2]")
         .appName("orders").config("spark.ui.enabled", "false")
         .getOrCreate())
spark.sparkContext.setLogLevel("ERROR")

orders = spark.read.csv("orders.csv", header=True, inferSchema=True)
w = Window.partitionBy("customer_id").orderBy(F.col("amount").desc())

(orders
    .filter(F.col("status") == "delivered")
    .withColumn("rnk", F.dense_rank().over(w))
    .filter("rnk = 1")
    .select("customer_id", "order_id", "amount")
    .orderBy("customer_id")
    .show())
spark.stop()

Output

+-----------+--------+------+
|customer_id|order_id|amount|
+-----------+--------+------+
|          1|     114|2397.0|
|          2|     111|1999.0|
|          3|     109| 747.0|
|          4|     106|1048.0|
|          5|     108|3498.0|
|          6|     110|1998.0|
|          7|     112| 846.0|
+-----------+--------+------+

Task	pandas	PySpark
Filter	`df[df.status == 'delivered']`	`df.filter(F.col('status') == 'delivered')`
New column	`df.assign(rev=df.qty * df.price)`	`df.withColumn('rev', F.col('qty') * F.col('price'))`
Group and sum	`df.groupby('c').agg(t=('amt', 'sum'))`	`df.groupBy('c').agg(F.sum('amt').alias('t'))`
Join	`a.merge(b, on='id', how='left')`	`a.join(b, 'id', 'left')`
Window rank	`df.groupby('c')['amt'].rank(ascending=False)`	`F.dense_rank().over(Window.partitionBy('c')...)`
When it runs	Immediately, in memory	Lazily, on an action (show, count, write)
Scale	One machine, data must fit in RAM	A cluster, data split into partitions

THE SAME PIPELINE AS AN AIRFLOW DAGAIRFLOW 3 · NOT RUN

from airflow.sdk import dag, task
import pendulum

@dag(schedule="@daily", catchup=False, tags=["orders"],
     start_date=pendulum.datetime(2025, 4, 1, tz="Asia/Kolkata"))
def orders_daily():

    @task(retries=3)
    def extract(ds=None) -> str:     # ds = logical date
        return f"s3://raw/orders/{ds}.csv"

    @task
    def transform(path: str) -> str:
        return path.replace("raw", "clean")

    @task
    def load(path: str) -> int:
        ...                              # MERGE into the warehouse

    load(transform(extract()))           # sets the task order

orders_daily()

Practice before you move on

Rewrite the PySpark query in pandas and check that the rows match.
List the tasks of the Airflow DAG in the order they would run.