Study guide
3.4 Developing a Training Pipeline
Databricks Certified Machine Learning Associate Study Guide Exam
Track your progress as you read
Finished sections get a check mark.
1 Chapter 1: Databricks Machine Learning (38% of the Exam)
1.1 MLOps Best Practices on Databricks 1.2 Advantages of Databricks ML Runtimes and Compute Selection 1.3 Promoting Code vs. Promoting Models 1.4 How AutoML Facilitates Model and Feature Selection 1.5 AutoML Advantages in the Model Development Process 1.6 Feature Tables in Unity Catalog vs. the Workspace Feature Store 1.7 Creating a Feature Table and Writing Data to It 1.8 Training a Model with Features from a Feature Table 1.9 Scoring a Model with Features from a Feature Table 1.10 Online vs. Offline Feature Tables 1.11 Manually Logging Metrics, Artifacts, and Models in an MLflow Run 1.12 Identifying the Best Run with the MLflow Client API 1.13 Information Available in the MLflow UI 1.14 Registering a Model in the Unity Catalog Registry 1.15 Model Tags, Aliases, and Champion/Challenger Promotion
2 Chapter 2: Data Processing and ML Workflows (19% of the Exam)
2.1 Computing Summary Statistics on a Spark DataFrame 2.2 Creating Visualizations for Categorical and Continuous Features 2.3 Comparing Two Categorical or Two Continuous Features 2.4 Removing Outliers with Standard Deviation and IQR 2.5 When Log-Scale Transformation Is Appropriate 2.6 Comparing and Applying Mean, Median, and Mode Imputation 2.7 One-Hot Encoding and When It Is or Is Not Appropriate
3 Chapter 3: Model Development (31% of the Exam)
3.1 Selecting the Appropriate Algorithm for a Scenario 3.2 Methods to Mitigate Data Imbalance in Training Data 3.3 Comparing Estimators and Transformers 3.4 Developing a Training Pipeline 3.5 Tuning Hyperparameters with Hyperopt's fmin 3.6 Grid, Random, and Bayesian Search Strategies 3.7 Parallelizing Single-Node Models for Hyperparameter Tuning 3.8 Cross-Validation vs. a Train-Validation Split 3.9 Counting the Models Trained by Grid Search and Cross-Validation 3.10 Using Common Classification Metrics 3.11 Using Common Regression Metrics 3.12 Choosing the Most Appropriate Metric for the Objective 3.13 Exponentiating Log-Transformed Predictions Before Evaluation 3.14 Model Complexity and the Bias-Variance Tradeoff
4 Chapter 4: Model Deployment (12% of the Exam)
4.1 Differences and Advantages of Batch, Streaming, and Real-Time Serving 4.2 Performing Batch Inference with pandas and Spark UDFs 4.3 Streaming Inference with Delta Live Tables 4.4 Deploying a Custom Model to a Model Endpoint 4.5 Deploying, Querying, and Splitting Traffic Across Real-Time Endpoints
3.4 Developing a Training Pipeline
Key Takeaways Explain This Topic AI
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
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 .
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
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.
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:
Impute numeric nulls ( Imputer ) — before anything computes statistics on them.
Index categorical strings ( StringIndexer , handleInvalid="keep" ).
Encode indices to indicator vectors ( OneHotEncoder ) — only if the model
benefits (see Section 2.7).
Assemble every numeric and encoded column into one vector ( VectorAssembler ).
Scale the assembled vector ( StandardScaler ) — only for scale-sensitive models.
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"
- )