Module 10 — Project: a model trained on a large dataset

The nine previous modules assembled themselves for this one. Here is the entire flight-delays project as a single, runnable file, followed by the same pipeline in pandas on a five-per-cent sample. The comparison at the end is what tells you, honestly, whether Spark was worth the effort — the answer is not always yes.

The complete Spark job

from pyspark.sql import SparkSession, functions as F
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler
from pyspark.ml.classification import GBTClassifier
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
from pyspark.ml.evaluation import BinaryClassificationEvaluator

spark = (
    SparkSession.builder
    .appName("flight-delays-training")
    .config("spark.sql.shuffle.partitions", 400)
    .config("spark.sql.adaptive.enabled", True)
    .config("spark.sql.sources.partitionOverwriteMode", "dynamic")
    .getOrCreate()
)

# --- 1. Read with an explicit schema, partitioned by year on disk (module 4). ---
schema = ("carrier STRING, origin STRING, dest STRING, distance INT, "
          "departure_delay INT, arrival_delay INT, cancelled BOOLEAN, "
          "year INT, month INT, day INT")

flights = (
    spark.read.schema(schema)
    .parquet("s3://open-data/flights-parquet/")
    .filter(F.col("cancelled") == False)
    .withColumn("late", (F.col("arrival_delay") > 15).cast("int"))
)

train = flights.filter(F.col("year") < 2024)
test  = flights.filter(F.col("year") == 2024)

# --- 2. Feature pipeline (module 5). ---
carrier_idx = StringIndexer(inputCol="carrier", outputCol="carrier_idx", handleInvalid="keep")
origin_idx  = StringIndexer(inputCol="origin",  outputCol="origin_idx",  handleInvalid="keep")
dest_idx    = StringIndexer(inputCol="dest",    outputCol="dest_idx",    handleInvalid="keep")

oh = OneHotEncoder(
    inputCols=["carrier_idx", "origin_idx", "dest_idx"],
    outputCols=["carrier_vec", "origin_vec", "dest_vec"],
)

numeric_asm = VectorAssembler(
    inputCols=["distance", "departure_delay", "month"],
    outputCol="numeric_vec",
)
scaler = StandardScaler(inputCol="numeric_vec", outputCol="numeric_scaled",
                        withMean=True, withStd=True)

features_asm = VectorAssembler(
    inputCols=["numeric_scaled", "carrier_vec", "origin_vec", "dest_vec"],
    outputCol="features",
)

gbt = GBTClassifier(labelCol="late", featuresCol="features", seed=42)

pipeline = Pipeline(stages=[
    carrier_idx, origin_idx, dest_idx, oh,
    numeric_asm, scaler, features_asm, gbt,
])

# --- 3. Cross-validated tuning (module 6). ---
grid = (ParamGridBuilder()
        .addGrid(gbt.maxDepth, [5, 7])
        .addGrid(gbt.maxIter, [80, 120])
        .build())    # 4 combinations

evaluator = BinaryClassificationEvaluator(labelCol="late", metricName="areaUnderROC")
cv = CrossValidator(
    estimator=pipeline, estimatorParamMaps=grid,
    evaluator=evaluator, numFolds=3, parallelism=4, seed=42,
)

cv_model = cv.fit(train)
best_pipeline = cv_model.bestModel

# --- 4. Evaluate on the held-out year. ---
predictions = best_pipeline.transform(test)
auc = evaluator.evaluate(predictions)
print(f"Test AUC on 2024 flights: {auc:.4f}")

# --- 5. Persist model and scored predictions (module 9). ---
best_pipeline.write().overwrite().save("s3://models/flights-delay-gbt-v1/")

(predictions
    .select("carrier", "origin", "dest", "year", "month", "day", "prediction", "probability")
    .coalesce(8)
    .write.mode("overwrite").partitionBy("year", "month")
    .parquet("s3://predictions/flights/"))

On a 12-worker cluster with 4 cores each and 60 million training rows, this job runs in roughly 14 minutes end to end, most of which is spent in the three cross-validation folds. On a laptop in local[*] mode with 8 cores, on a five-million-row sample, it runs in about 9 minutes.

The pandas control

import pandas as pd
from sklearn.pipeline import Pipeline as SkPipeline
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import roc_auc_score
from xgboost import XGBClassifier

df = pd.read_parquet("s3://open-data/flights-parquet/", filters=[("cancelled", "==", False)])
df["late"] = (df["arrival_delay"] > 15).astype(int)

train_df = df[df["year"] < 2024].sample(frac=0.05, random_state=42)   # 5 % sample
test_df  = df[df["year"] == 2024]

numeric = ["distance", "departure_delay", "month"]
categorical = ["carrier", "origin", "dest"]

pre = ColumnTransformer([
    ("num", StandardScaler(), numeric),
    ("cat", OneHotEncoder(handle_unknown="ignore"), categorical),
])

pipe = SkPipeline([("pre", pre), ("clf", XGBClassifier(n_estimators=100, random_state=42))])

grid = GridSearchCV(pipe,
                    {"clf__max_depth": [5, 7], "clf__n_estimators": [80, 120]},
                    cv=3, scoring="roc_auc", n_jobs=-1)
grid.fit(train_df[numeric + categorical], train_df["late"])
print("Test AUC:", roc_auc_score(test_df["late"], grid.predict_proba(test_df[numeric + categorical])[:, 1]))

Same features, same grid, same metric. On a five-per-cent sample of the training set — three million rows — this pipeline trains in about 8 minutes on a laptop and reaches an AUC within 1 percentage point of the Spark run on the full data.

What the volume actually bought

Three lessons come out of running both jobs side by side.

The Spark job wins on the tail. The full-data model's advantage over the sampled one is not in overall AUC — it is in the rare combinations: the small regional carrier flying an unusual route in a specific month. Those combinations barely exist in a five-per-cent sample, and the pandas model does poorly on them. If the downstream question is "give me the best model on average", the sample is often good enough. If the question is "predict every specific route accurately", the volume matters.

Spark's fixed cost is real. For anything below tens of millions of rows, the driver startup, the JVM warmup and the plan compilation add a two-to-five-minute overhead that pandas simply does not pay. This is the module 1 warning made concrete.

The pipelines look the same. The Spark code and the pandas code are structurally identical: split, preprocess, assemble, tune, evaluate, save. The vocabulary of module 6 — Transformer, Estimator, Pipeline — is what makes that possible. Any team fluent in scikit-learn can read the Spark version, and vice versa.

The honest decision

The right process is: build the pipeline in pandas on a sample first, verify it produces a reasonable model, then translate it to Spark only if the volume you actually need in production justifies the fixed cost. The reverse — starting in Spark and hoping the extra rows will help — is how projects burn a quarter on infrastructure and deliver the same model six months late.

Summary

The whole course fits in one Python file: partitioned Parquet read, feature pipeline, cross-validated GBT, persisted model, partitioned write.
On the same data, a pandas + XGBoost pipeline on a 5 % sample reaches an AUC within 1 point of the full Spark run.
Spark's payoff is in the tail of the distribution — rare feature combinations that need every row to be learnable.
Prototype in pandas, translate to Spark only when the production volume actually forces it.

Next: the recap and the final exam.

The complete Spark job​

The pandas control​

What the volume actually bought​

Summary​

The complete Spark job

The pandas control

What the volume actually bought

Summary