Spark ML Pipelines

A machine learning project repeats the same steps: clean the data, convert text to numbers, combine features, and train a model. A Spark ML Pipeline chains those steps into one object. You train the pipeline once and reuse it on new data with the same preparation every time.

The Layman View

A factory assembly line moves a product through fixed stations. Each station performs one job, and the product leaves as a finished item. A pipeline moves a DataFrame through stations and returns predictions.

Two Building Blocks

  • Transformer converts one DataFrame into another. It applies a rule and needs no learning. Example: VectorAssembler.
  • Estimator learns from data and produces a Transformer. Example: LogisticRegression produces a trained model.
Transformer :  DataFrame --transform()--> DataFrame
Estimator   :  DataFrame --fit()-------> Model (a Transformer)

Sample Data

data = [
    ("Delhi",  25, 40000.0, 0), ("Pune",   45, 90000.0, 1),
    ("Delhi",  35, 60000.0, 1), ("Mumbai", 22, 30000.0, 0),
    ("Pune",   50, 85000.0, 1), ("Mumbai", 28, 35000.0, 0),
    ("Delhi",  41, 78000.0, 1), ("Pune",   23, 32000.0, 0),
]
df = spark.createDataFrame(data, ["city", "age", "income", "label"])
train, test = df.randomSplit([0.75, 0.25], seed=42)

The label column holds the answer that the model learns to predict, such as whether a customer buys a product.

The Stations

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler
from pyspark.ml.classification import LogisticRegression

indexer   = StringIndexer(inputCol="city", outputCol="city_idx", handleInvalid="keep")
encoder   = OneHotEncoder(inputCols=["city_idx"], outputCols=["city_vec"])
assembler = VectorAssembler(inputCols=["city_vec", "age", "income"], outputCol="raw_features")
scaler    = StandardScaler(inputCol="raw_features", outputCol="features")
lr        = LogisticRegression(featuresCol="features", labelCol="label")
  • StringIndexer turns city names into numbers.
  • OneHotEncoder turns those numbers into separate yes-or-no columns.
  • VectorAssembler packs all inputs into one feature vector.
  • StandardScaler puts features on a similar scale.
  • LogisticRegression learns the prediction rule.

Building and Training

pipeline = Pipeline(stages=[indexer, encoder, assembler, scaler, lr])
model = pipeline.fit(train)

train data
   |
StringIndexer --> OneHotEncoder --> VectorAssembler --> StandardScaler --> LogisticRegression
                                                                            |
                                                                    PipelineModel

The fit call runs every Estimator in order and returns a PipelineModel that holds the learned stations.

Making Predictions

predictions = model.transform(test)
predictions.select("city", "age", "income", "label", "prediction").show()

The same fitted model prepares and scores brand-new data. No step gets repeated by hand, so training and prediction treat the data identically.

Measuring Quality

from pyspark.ml.evaluation import BinaryClassificationEvaluator

evaluator = BinaryClassificationEvaluator(labelCol="label", metricName="areaUnderROC")
print("AUC:", evaluator.evaluate(predictions))

A score near 1.0 shows strong separation between the two classes. A score near 0.5 shows guessing.

Tuning with Cross-Validation

A grid of settings lets Spark test several model configurations and keep the best.

from pyspark.ml.tuning import ParamGridBuilder, CrossValidator

grid = (ParamGridBuilder()
          .addGrid(lr.regParam, [0.01, 0.1, 1.0])
          .addGrid(lr.maxIter, [10, 50])
          .build())

cv = CrossValidator(estimator=pipeline, estimatorParamMaps=grid,
                    evaluator=evaluator, numFolds=3)
best_model = cv.fit(train).bestModel

Saving and Loading

from pyspark.ml import PipelineModel

model.write().overwrite().save("/models/purchase_model")
loaded = PipelineModel.load("/models/purchase_model")
loaded.transform(test).show()

Good Practices

  • Fit the pipeline on training data only. Test data stays unseen until evaluation.
  • Use handleInvalid="keep" so unknown categories at prediction time do not crash the job.
  • Save the full PipelineModel, never only the final model, so preparation steps travel with it.
  • Version saved models with dates or numbers.

Inspecting the Learned Model

Every stage of a PipelineModel stays accessible. Pull out the last stage to read the learned weights.

lr_model = model.stages[-1]
print("Coefficients:", lr_model.coefficients)
print("Intercept:", lr_model.intercept)

Each coefficient shows how strongly one feature pushes the prediction toward the positive class. A positive value raises the chance of a positive result, and a negative value lowers it.

Train and Test Split Reminder

All data (100%)
   |
   |-- Train (75%) : fit the pipeline
   |-- Test  (25%) : measure quality on unseen rows
Previous lesson
Back to course
Next lesson