Book Case Study - US Flight Delays & UDFs

This guide covers the classic US Flight Delays case study from Chapter 4 of the official O'Reilly book Learning Spark (2nd Edition) by Jules S. Damji, Brooke Wenig, Tathagata Das, and Denny Lee. It demonstrates how to leverage Spark SQL by registering Temporary Views, running relational ANSI SQL queries, and extending the SQL engine with custom User-Defined Functions (UDFs).

The Scenario

We are working with a dataset of US domestic flights containing columns: date (represented as a string like 01010900), delay (in minutes), distance (in miles), origin (3-letter airport code), and destination (3-letter airport code).

By applying Spark SQL and custom functions, we will:

1. Load the flight delays dataset and register it as a temporary view.
2. Run standard ANSI SQL queries to discover high-distance delay patterns and SFO-to-ORD bottlenecks.
3. Write a custom Python UDF to categorize delays, register it within Spark's catalog, and call it dynamically from both Spark SQL and DataFrame DSL.
4. Learn the performance cost of Python UDFs and why they should be replaced by Spark's native functions wherever possible.

Programmatic Schema & Ingestion

We begin by defining the schema and loading the raw flight delay data:

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

flight_schema = StructType([
    StructField("date", StringType(), True),
    StructField("delay", IntegerType(), True),
    StructField("distance", IntegerType(), True),
    StructField("origin", StringType(), True),
    StructField("destination", StringType(), True)
])

PySpark Script: Flight Delays & UDF Analyzer

Below is the complete, self-contained Python script to load the dataset, register the view, run SQL commands, and implement the custom Python UDF.

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf
from pyspark.sql.types import StringType

# 1. Initialize SparkSession
spark = SparkSession.builder \
    .appName("USFlightDelaysUDFAnalyzer") \
    .master("local[*]") \
    .getOrCreate()

# 2. Ingest the dataset
# For testing, we mock data or load an actual file if present
raw_flights_df = spark.read.format("csv") \
    .option("header", "true") \
    .schema(flight_schema) \
    .load("departuredelays.csv")

# 3. Register as a Temporary View
# This registers the DataFrame as an in-memory SQL table scoped to the SparkSession
raw_flights_df.createOrReplaceTempView("us_delay_flights_tbl")

# ==========================================
# Part 1: Relational SQL Queries
# ==========================================

# Query 1: Find all flights with a distance greater than 1000 miles ordered by longest delay
print("Flights with distance > 1000 miles:")
spark.sql("""
    SELECT date, delay, distance, origin, destination
    FROM us_delay_flights_tbl
    WHERE distance > 1000
    ORDER BY delay DESC
""").show(5, False)

# Query 2: Find all flights between SFO and ORD with a delay greater than 120 minutes
print("Flights between SFO and ORD delayed > 2 hours:")
spark.sql("""
    SELECT date, delay, origin, destination
    FROM us_delay_flights_tbl
    WHERE origin = 'SFO' AND destination = 'ORD' AND delay > 120
    ORDER BY delay DESC
""").show(5, False)

# ==========================================
# Part 2: Custom User-Defined Functions (UDFs)
# ==========================================

# Define a Python function to categorize delay minutes
def categorize_delays(delay):
    if delay is None:
        return "Unknown"
    elif delay > 360:
        return "Very Long Delay (>6 Hours)"
    elif delay > 120:
        return "Long Delay (2-6 Hours)"
    elif delay > 0:
        return "Short Delay (<2 Hours)"
    elif delay == 0:
        return "On Time"
    else:
        return "Early Flight"

# 1. Register the UDF for Spark SQL
# This registers the UDF in Spark's catalog so it can be used inside SQL queries!
spark.udf.register("cat_delays_sql", categorize_delays, StringType())

# 2. Register the UDF for DataFrame DSL
# This wraps the Python function for use inside .withColumn() or .select()
cat_delays_dsl = udf(categorize_delays, StringType())

# Run UDF via Spark SQL
print("Using registered UDF in standard Spark SQL:")
spark.sql("""
    SELECT date, delay, origin, destination, cat_delays_sql(delay) AS delay_category
    FROM us_delay_flights_tbl
    WHERE origin = 'SFO'
    ORDER BY delay DESC
""").show(5, False)

# Run UDF via DataFrame DSL
print("Using UDF in DataFrame DSL:")
raw_flights_df.filter(col("origin") == "SFO") \
    .withColumn("delay_category", cat_delays_dsl(col("delay"))) \
    .select("date", "delay", "origin", "destination", "delay_category") \
    .orderBy(col("delay").desc()) \
    .show(5, False)

Spark SQL Plan and UDF Execution Bottleneck

Unlike native Spark operations, custom Python UDFs introduce massive performance overhead. This is because PySpark must run the Python function outside the JVM execution process.

graph TD
    A["Spark JVM (Executor Context)"] -->|"1. Serialize Data Rows"| B["Py4J / IPC socket channel"]
    B -->|"2. Send Rows to Python"| C["Python Subprocess (Worker)"]
    C -->|"3. Execute Python UDF Function"| D["categorize_delays(delay)"]
    D -->|"4. Serialize Results"| E["Py4J / IPC socket channel"]
    E -->|"5. Send Back to JVM"| A

    style A fill:#eff6ff,stroke:#2563eb,stroke-width:2px;
    style C fill:#fef2f2,stroke:#dc2626,stroke-width:2px;
    style D fill:#fef2f2,stroke:#dc2626,stroke-width:2px;

Why Custom Python UDFs are Slow:

1. Inter-Process Communication (IPC): Spark Executors run inside the Java Virtual Machine (JVM). When executing a Python UDF, Spark must serialize the data from the JVM, send it over an IPC socket to a spawned Python subprocess, run the Python code, serialize the output, and send it back to the JVM.
2. No Catalyst Optimizations: Spark's Catalyst Optimizer cannot "look inside" custom Python functions. It treats the UDF as a black box, which disables major optimizations like predicate pushdown, pipeline code generation, and memory management.

The High-Performance Alternative:

Instead of writing a custom Python UDF, use native Spark SQL expressions like when().otherwise() or standard SQL CASE WHEN constructs. These run completely inside the JVM at native speed and are fully optimized by Catalyst!

# High-performance native alternative:
from pyspark.sql.functions import when

optimized_df = raw_flights_df.withColumn(
    "delay_category",
    when(col("delay") > 360, "Very Long Delay (>6 Hours)")
    .when(col("delay") > 120, "Long Delay (2-6 Hours)")
    .when(col("delay") > 0, "Short Delay (<2 Hours)")
    .when(col("delay") == 0, "On Time")
    .otherwise("Early Flight")
)