A week after I finished writing k-means by hand as a map and a reduceByKey call, I opened MLlib’s Pipeline API for the first time properly, rather than as a single method call I trusted and moved past. KMeans.train() had stopped being a mystery. What I had not looked at closely was everything MLlib does before and after the training step itself: turning a categorical column into numbers, assembling several feature columns into one vector, running the estimator, and checking the result against held-out data, all chained together as one object. Having written k-means and linear regression by hand on RDDs a week earlier changed how I read that chain. It stopped looking like a framework doing something clever and started looking like the exact sequence of steps I would have written myself, given a decent afternoon and a reason to.

The problem a pipeline actually solves

Training a model is rarely the first thing that happens to raw data. A categorical column has to become numeric. Several numeric columns have to become one feature vector, because most Spark ML estimators expect a single vector column, not a scattering of individual ones. The trained model then needs applying to new data using exactly the same encoding and assembly steps used during training, not a slightly different version written from memory a week later. Doing each of those steps as a separate, manually-tracked transformation on a DataFrame works, but it means remembering the exact order and exact parameters every time the model gets applied to anything new. A Pipeline bundles that sequence into one object that remembers its own steps.

import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.feature.{StringIndexer, VectorAssembler}
import org.apache.spark.ml.classification.DecisionTreeClassifier
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator

val data = spark.read
  .option("header", "true")
  .option("inferSchema", "true")
  .csv("loan_applications.csv")

val Array(trainData, testData) = data.randomSplit(Array(0.8, 0.2), seed = 42)

val labelIndexer = new StringIndexer()
  .setInputCol("outcome")
  .setOutputCol("label")

val assembler = new VectorAssembler()
  .setInputCols(Array("income", "existing_debt", "credit_history_months"))
  .setOutputCol("features")

val tree = new DecisionTreeClassifier()
  .setLabelCol("label")
  .setFeaturesCol("features")

val pipeline = new Pipeline().setStages(Array(labelIndexer, assembler, tree))

Nothing has run yet. Building a Pipeline from an array of stages is a declaration, not an execution, the same laziness that makes sc.textFile() return instantly rather than reading a file. Each stage in that array is one of two kinds of object, and the difference between the two kinds turned out to be the part worth understanding properly.

Pipeline stages running left to right from raw data through StringIndexer, VectorAssembler, and DecisionTreeClassifier to a MulticlassClassificationEvaluator, with fit() producing a PipelineModel and transform() applying that fitted model to test data

Transformers do not learn anything, estimators do

VectorAssembler takes columns in, produces one vector column out, and needs no training data to do it. Given income, existing_debt, and credit_history_months, it just concatenates the three values into a vector, row by row, with nothing to learn from the data first. That is a Transformer, in MLlib’s terms: something with a .transform() method and no .fit().

StringIndexer looks similar on the surface, another feature-shaping step, but it has to look at the actual outcome column first to decide which string maps to which numeric index, approved might become 0 on one run and 1 on another, depending purely on which value shows up more often in the training data it saw. That means StringIndexer needs a .fit() step before it can transform anything, and calling .fit() on it does not produce transformed data directly, it produces a StringIndexerModel, a fitted object that itself has a .transform() method carrying the exact string-to-index mapping it learned.

val indexerModel = labelIndexer.fit(trainData)
val indexed = indexerModel.transform(trainData)

DecisionTreeClassifier works the identical way, .fit() on training data produces a DecisionTreeClassificationModel, which is itself a Transformer with a .transform() method that adds prediction columns to whatever DataFrame it is handed. This is the pattern that made the whole API click for me: an Estimator’s .fit() always hands back a Model, and every Model is also a Transformer. The type hierarchy is not incidental, it is what lets a Pipeline, itself an Estimator, treat “fit each stage in order, feeding fitted transformers forward” as one uniform operation regardless of whether a given stage needed fitting at all.

Two API patterns side by side: a Transformer takes a DataFrame and returns a DataFrame with no fitting step, exemplified by VectorAssembler; an Estimator takes a DataFrame, calls fit(), and produces a Model, which is itself a Transformer, exemplified by DecisionTreeClassifier producing a DecisionTreeClassificationModel

Fitting the whole pipeline at once

val pipelineModel = pipeline.fit(trainData)
val predictions = pipelineModel.transform(testData)

predictions.select("outcome", "label", "prediction", "features").show(5, truncate = false)

Calling .fit() on the Pipeline runs every stage in order against trainData: labelIndexer.fit() produces its StringIndexerModel, that model transforms the label column, assembler.transform() builds the feature vector, and tree.fit() trains against the resulting features and labels. The object that comes back, pipelineModel, is a PipelineModel, which bundles the fitted StringIndexerModel and the fitted DecisionTreeClassificationModel together as one chain. Calling .transform() on testData runs that exact same chain again, same string-to-index mapping, same tree, this time producing predictions rather than training one.

That second detail is the one that would have bitten me if I had written the manual version carelessly. Applying a fitted model to new data has to use the encoding learned on training data, not a fresh encoding fitted separately on the test set, or a category present only in test data would get an index that means something different to the model, or no index at all. PipelineModel carries the training-time StringIndexerModel forward automatically. A manual version has to remember to do that on purpose.

Evaluating on the distributed result

val evaluator = new MulticlassClassificationEvaluator()
  .setLabelCol("label")
  .setPredictionCol("prediction")
  .setMetricName("f1")

val f1Score = evaluator.evaluate(predictions)
println(s"F1 score: $f1Score")

MulticlassClassificationEvaluator runs as a distributed aggregation over predictions, which is itself a DataFrame spread across partitions the same way any Spark DataFrame is. Computing an F1 score needs true positives, false positives, and false negatives counted across every row, which is a reduction no different in kind from reduceByKey summing partial gradients across partitions in the manual gradient descent version. The evaluator hides that reduction behind .evaluate(), but it is still a shuffle-bearing aggregation running underneath, not a free operation just because it is a one-line call.

What the pipeline abstraction buys, specifically

The value is not that a Pipeline does something a manual sequence of transforms cannot do. It is that the manual sequence has several places to get subtly wrong, fitting the label indexer on the wrong dataset, forgetting to reuse a fitted transformer rather than refitting it on test data, losing track of stage order when the pipeline grows past three or four steps, and Pipeline removes every one of those failure modes by construction rather than by discipline. pipelineModel.write.save("loan_model") persists the whole fitted chain, indexer mapping and tree structure together, and PipelineModel.load() on the other side gets back an object that transforms new data identically to how the original did, without me having to serialise three separate pieces of state and reassemble them correctly by hand.

I would not have appreciated any of that a month earlier. Reading Pipeline documentation before writing k-means manually, the Transformer and Estimator distinction read as API ceremony, two names for things that both take a DataFrame and give one back. Having actually written a .fit()-shaped loop myself, one that computes something from training data and then has to apply that exact computed thing consistently to new data afterwards, the distinction reads as the entire point. Estimator and Transformer are not two arbitrary categories MLlib invented. They are the two things every training pipeline has to do, named honestly, and enforced by the type system rather than left to whoever writes the pipeline to remember correctly every time.