Skip to main content

Module 6 — Pipelines and distributed cross-validation

Doing the four preprocessing steps of module 5 by hand and then adding a classifier works exactly once. As soon as you touch cross-validation, hyperparameter search, or reloading the model for scoring, the manual chain becomes a source of subtle bugs — leakage first among them. Spark's Pipeline is the object that solves this.

Pipeline — one object, seven steps

A Pipeline is an Estimator that wraps an ordered list of Transformers and Estimators, called stages. Calling pipeline.fit(train) walks the list: each stage is either fitted on the current DataFrame and its output Transformer stored, or applied as-is; the fitted Transformer is used to transform the DataFrame for the next stage. The result is a PipelineModel, itself a Transformer.

from pyspark.ml import Pipeline
from pyspark.ml.classification import GBTClassifier

pipeline = Pipeline(stages=[
carrier_idx, # module 5: StringIndexer
carrier_oh, # OneHotEncoder
numeric_asm, # VectorAssembler on numeric columns
scaler, # StandardScaler on the numeric vector
features_asm, # VectorAssembler on everything -> "features"
GBTClassifier(labelCol="late", featuresCol="features", maxIter=50),
])

pipeline_model = pipeline.fit(train)
predictions = pipeline_model.transform(test)

Two properties of this object make it worth every line of ceremony. First, the whole thing serialises: pipeline_model.write().overwrite().save("s3://models/flights-v1") produces a folder that can be reloaded next month in another job. Second, and more importantly, the pipeline contains the leakage discipline: each stage's fit is called on the training portion of the current cross-validation fold, so the scaler never sees test rows before scoring them.

Preprocessing outside the pipeline leaks — silently

The single most common Spark ML bug is doing the preprocessing before splitting or before cross-validation, then feeding a "ready" DataFrame into the pipeline.

# WRONG - the scaler sees the whole dataset before we split
scaled = StandardScaler(...).fit(flights).transform(flights)
train, test = scaled.randomSplit([0.8, 0.2])
model = GBTClassifier(...).fit(train) # trained model looks great

The evaluate(test) numbers are optimistic because the mean and standard deviation used to standardise the test rows were computed with those test rows in the mix. On the flight-delays project this typically inflates AUC by one or two points — enough to promote a mediocre model. The fix is to put every learned transformation inside the pipeline, and only ever call pipeline.fit() on training data.

ParamGridBuilder builds a grid of hyperparameter combinations. Each .addGrid(param, [values]) multiplies the number of combinations by the number of values.

from pyspark.ml.tuning import ParamGridBuilder

gbt = pipeline.getStages()[-1] # the classifier stage

grid = (
ParamGridBuilder()
.addGrid(gbt.maxDepth, [4, 6, 8])
.addGrid(gbt.stepSize, [0.05, 0.1])
.addGrid(gbt.maxIter, [50, 100])
.build()
)
# 3 * 2 * 2 = 12 combinations

The size of the grid is a first-order cost driver. A five-fold cross-validation on twelve combinations trains sixty models. On a fifty-million-row dataset, this is measured in hours or in cloud dollars, not minutes. Always check the size before launching: print(len(grid)).

CrossValidator — distributed by design

from pyspark.ml.tuning import CrossValidator
from pyspark.ml.evaluation import BinaryClassificationEvaluator

cv = CrossValidator(
estimator=pipeline,
estimatorParamMaps=grid,
evaluator=BinaryClassificationEvaluator(labelCol="late", metricName="areaUnderROC"),
numFolds=5,
parallelism=4, # train up to 4 models in parallel per fold
seed=42,
)

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

Everything Spark can parallelise is parallelised here. Within one fold, the parallelism=4 setting trains four models at once, sharing the executors. Across folds, the driver schedules them sequentially by default but the computation of each fold uses the full cluster. On the flight-delays project this typically yields a five- to eight-times speed-up over a naive Python loop for the same total work.

cv_model.avgMetrics returns the mean metric across folds for each grid point, in the order of the grid; that array plus the grid itself is what you save to explain the choice of the winning combination to a stakeholder.

Reading bestModel

bestModel is a PipelineModel. Its stages contain the fitted parameters — the vocabularies, the scaler statistics, the trees of the gradient-boosted classifier. Anyone can reload it and score new data without knowing anything about the preprocessing that went into it. This is the payoff of the pipeline discipline: the deployable object is exactly the object you trained, not a slightly different Python script that "does roughly the same thing".

A cheaper alternative for exploration

When the dataset is large and the grid is small, TrainValidationSplit — a single 80/20 split — is often good enough and costs 1/numFolds of a full cross-validation. Use it while iterating on the pipeline structure, and switch to CrossValidator only when comparing final candidates.

Summary

  • A Pipeline chains Transformers and Estimators; the resulting PipelineModel is a single reloadable object.
  • Preprocessing must live inside the pipeline; preprocessing on the whole dataset before splitting leaks.
  • ParamGridBuilder describes the search; check len(grid) before you launch — cost is linear in it.
  • CrossValidator parallelises across grid points with parallelism; the bestModel is a PipelineModel ready for production.

Next module: what MLlib actually contains, and what it does not.