PracticeVideosBlogFlashcardsEspañol

Free exam prep

Databricks Certified Machine Learning Associate

Practice QuestionsStudy Guide

3.4 Developing a Training Pipeline

Key Takeaways

  • `Pipeline(stages=[...])` chains Transformers and Estimators into one Estimator; `pipeline.fit(train_df)` returns a `PipelineModel` in which every stage is a fitted Transformer.
  • During `fit`, each Estimator stage is fitted and immediately used to transform the data before the next stage sees it, so stage order matters.
  • Encapsulating preprocessing inside the pipeline is what makes cross-validation leak-free: each fold refits the imputers, indexers, and scalers on its own training rows.
  • The fitted `PipelineModel` guarantees identical preprocessing at training and inference, which is the structural cure for train/serve skew.
  • Persist with `pipeline_model.write().overwrite().save(path)` and reload with `PipelineModel.load(path)`, or log the whole pipeline with `mlflow.spark.log_model`.
Last updated: August 2026

3.4 Developing a Training Pipeline

A training pipeline is more than tidy code. In Spark ML the Pipeline is the object that makes correctness structural: because preprocessing lives inside the same Estimator as the model, there is no code path in which a scaler is fitted on data the model was not trained on, and no way for production scoring to apply different transformations than training did.

Pipeline Execution Mechanics & DAG Data Flow

A Pipeline chains multiple Estimators and Transformers into an ordered sequence of stages: Pipeline(stages=[stage1, stage2, ..., stageK]).

+-----------------------------------------------------------------------------+
|                   PIPELINE EXECUTION & DATA FLOW (FIT PHASE)                |
|                                                                             |
|  Input Train DataFrame                                                      |
|         |                                                                   |
|         v                                                                   |
|  +---------------+  .fit()    +--------------------+                        |
|  | StringIndexer | ---------> | StringIndexerModel | (Transformer 1)        |
|  +---------------+            +--------------------+                        |
|         |                               | .transform()                      |
|         +-------------------------------+                                   |
|         v                                                                   |
|  Intermediate DF 1                                                          |
|         |                                                                   |
|         v                                                                   |
|  +-----------------+          +--------------------+                        |
|  | VectorAssembler | -------> |  (Is Transformer)  | (Transformer 2)        |
|  +-----------------+          +--------------------+                        |
|         |                               | .transform()                      |
|         +-------------------------------+                                   |
|         v                                                                   |
|  Intermediate DF 2                                                          |
|         |                                                                   |
|         v                                                                   |
|  +--------------------+ .fit() +--------------------+                       |
|  | LogisticRegression | -----> | LogisticRegrModel  | (Transformer 3)       |
|  +--------------------+        +--------------------+                       |
|                                         |                                   |
|                                         v                                   |
|                              =======================                        |
|                              FITTED PIPELINEMODEL                           |
|                              [T1, T2, T3]                                   |
|                              =======================                        |
+-----------------------------------------------------------------------------+

Step-by-Step Execution Lifecycle

  1. Pipeline Fitting (pipeline.fit(train_df)):

    • The training DataFrame enters Stage 0 (StringIndexer). Because Stage 0 is an Estimator, Spark invokes stage0.fit(train_df) to generate StringIndexerModel (a Transformer). Spark immediately calls StringIndexerModel.transform(train_df) to produce Intermediate_DF_1.
    • Intermediate_DF_1 passes to Stage 1 (VectorAssembler). Because Stage 1 is already a Transformer, Spark directly calls stage1.transform(Intermediate_DF_1) to produce Intermediate_DF_2.
    • Intermediate_DF_2 passes to Stage 2 (LogisticRegression). As an Estimator, Spark calls stage2.fit(Intermediate_DF_2) to produce LogisticRegressionModel.
    • All resulting fitted Transformers are bundled into a PipelineModel.

  2. Pipeline Inference (pipeline_model.transform(test_df)):

    • When transform() is called on a PipelineModel, all stages are already Transformers.
    • The test DataFrame flows through StringIndexerModel.transform(), then VectorAssembler.transform(), then LogisticRegressionModel.transform(), outputting the final scored DataFrame containing predicted probabilities and class labels.

Eliminating Train-Serve Skew & Preventing Data Leakage

A critical failure mode in machine learning is data leakage, where information from the validation or test set inadvertently influences training parameters (e.g., computing global mean/std over the entire dataset before splitting).

+-----------------------------------------------------------------------------+
|                      DATA LEAKAGE PREVENTION COMPARISON                     |
|                                                                             |
|   INCORRECT (MANUAL PREPROCESSING - DATA LEAKAGE):                          |
|   Raw Data ---> [ Global Scaler Fit (Train + Test) ] ---> Train / Test Split|
|                 * Leakage: Mean & Std include Test Set distribution! *      |
|                                                                             |
|   CORRECT (PIPELINE ENCAPSULATION - ZERO LEAKAGE):                          |
|   Raw Data ---> Train / Test Split                                          |
|                   |                                                         |
|                   +---> Train Set ---> [ Pipeline.fit() ]                   |
|                                                | (Scaler learns ONLY Train) |
|                                                v                            |
|                   +---> Test Set  ---> [ PipelineModel.transform() ]        |
+-----------------------------------------------------------------------------+

Benefits of Spark ML Pipeline Encapsulation

  1. Zero Leakage during Cross-Validation: When tuning hyperparameter grids with CrossValidator, the Pipeline Estimator is passed directly into the evaluator. CrossValidator automatically fits all preprocessing stages (e.g., StandardScaler, StringIndexer) exclusively on the training folds, evaluating strictly on holdout folds.
  2. Identical Preprocessing in Production: A saved PipelineModel encapsulates the exact indexing mappings, imputation medians, and scaling factors learned during training. Real-time or batch inference calls pipeline_model.transform(raw_df) without requiring external lookup tables or manual preprocessing scripts.

Ordering the Stages

Stages execute in the order given, and each stage's output columns must exist by the time the next stage runs. The canonical tabular ordering is:

  1. Impute numeric nulls (Imputer) — before anything computes statistics on them.
  2. Index categorical strings (StringIndexer, handleInvalid="keep").
  3. Encode indices to indicator vectors (OneHotEncoder) — only if the model benefits (see Section 2.7).
  4. Assemble every numeric and encoded column into one vector (VectorAssembler).
  5. Scale the assembled vector (StandardScaler) — only for scale-sensitive models.
  6. Estimate — the learning algorithm, reading featuresCol and labelCol.

Getting the order wrong produces immediate, recognisable failures: assembling before indexing raises a type error on the string column; scaling before assembling scales the wrong object; imputing after assembling cannot reach the individual columns at all.

End-to-End Production Code Example

from pyspark.ml import Pipeline, PipelineModel
from pyspark.ml.feature import (
    StringIndexer,
    OneHotEncoder,
    VectorAssembler,
    StandardScaler,
    Imputer
)
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator

# 1. Load Data
df = spark.table("lakehouse_gold.customer_churn_raw")
train_df, test_df = df.randomSplit([0.8, 0.2], seed=42)

# 2. Define Preprocessing Stages
# Handle missing numerical values
imputer = Imputer(
    inputCols=["age", "monthly_charges"],
    outputCols=["age_imputed", "monthly_charges_imputed"]
).setStrategy("median")

# Index categorical strings
indexer = StringIndexer(
    inputCols=["contract_type", "payment_method"],
    outputCols=["contract_idx", "payment_idx"],
    handleInvalid="keep"  # Assigns unseen categories to a dedicated index
)

# Encode indices to one-hot vectors
encoder = OneHotEncoder(
    inputCols=["contract_idx", "payment_idx"],
    outputCols=["contract_ohe", "payment_ohe"]
)

# Assemble all features into a unified vector
assembler = VectorAssembler(
    inputCols=["age_imputed", "monthly_charges_imputed", "contract_ohe", "payment_ohe"],
    outputCol="unscaled_features"
)

# Scale numerical features
scaler = StandardScaler(
    inputCol="unscaled_features",
    outputCol="features",
    withStd=True,
    withMean=False
)

# 3. Define Model Estimator
rf = RandomForestClassifier(
    featuresCol="features",
    labelCol="churned",
    numTrees=100,
    maxDepth=6,
    seed=42
)

# 4. Construct Pipeline
pipeline = Pipeline(stages=[imputer, indexer, encoder, assembler, scaler, rf])

# 5. Fit Pipeline (Produces PipelineModel)
pipeline_model = pipeline.fit(train_df)

# 6. Transform and Evaluate on Test Set
predictions = pipeline_model.transform(test_df)

evaluator = BinaryClassificationEvaluator(
    labelCol="churned",
    rawPredictionCol="rawPrediction",
    metricName="areaUnderROC"
)

roc_auc = evaluator.evaluate(predictions)
print(f"Test Set ROC-AUC: {roc_auc:.4f}")

Pipeline Serialization & MLflow Logging

Spark ML provides native parquet-based serialization for both unfitted Pipeline DAGs and fitted PipelineModel instances:

# Save an unfitted Pipeline definition
pipeline.write().overwrite().save("dbfs:/models/churn_pipeline_def")
loaded_pipeline = Pipeline.load("dbfs:/models/churn_pipeline_def")

# Save a fitted PipelineModel
pipeline_model.write().overwrite().save("dbfs:/models/churn_pipeline_model")
loaded_model = PipelineModel.load("dbfs:/models/churn_pipeline_model")

Logging to MLflow

In enterprise Databricks environments, pipelines should be logged to MLflow to preserve schema metadata and enable one-click serving:

import mlflow
import mlflow.spark

with mlflow.start_run(run_name="spark_ml_pipeline_run"):
    # Fit pipeline
    pipeline_model = pipeline.fit(train_df)
    
    # Log parameters and metrics
    mlflow.log_param("max_depth", 5)
    mlflow.log_metric("roc_auc", auc_score)
    
    # Log the complete fitted pipeline
    mlflow.spark.log_model(
        spark_model=pipeline_model,
        artifact_path="spark_churn_model",
        registered_model_name="catalog.schema.churn_classifier"
    )
Loading diagram...
Spark ML Pipeline Stages and Transformation Graph
Test Your Knowledge

When pipeline.fit(train_df) is called on a Pipeline containing [StringIndexer, VectorAssembler, LogisticRegression], what object is returned and what are its properties?

A
B
C
D
Test Your Knowledge

Why is it considered best practice to include feature preprocessing stages (like StandardScaler and Imputer) inside a PySpark Pipeline rather than applying them to the full dataset prior to splitting?

A
B
C
D
Test Your Knowledge

A data engineer wants to save a trained PipelineModel to cloud object storage so that a separate batch inference pipeline can load and score new data. Which code snippet correctly persists and reloads the fitted pipeline in PySpark ML?

A
B
C
D
Up Next

3.5 Tuning Hyperparameters with Hyperopt's fmin

Continue learning