Model Development · Topic 2 of 5
Estimators, Transformers and Training Pipelines
Databricks ML Associate: tell Spark ML estimators from transformers, know what fit and transform return, and chain the stages in a training pipeline.
- 5 min
- Includes a diagram
- Transformers have
transform()and learn nothing; estimators havefit()and return a model. - A fitted model is a transformer, so it is applied with
transform(). - A Pipeline is an estimator; its
fit()returns a PipelineModel.
What it is
Spark ML builds every workflow from two kinds of stage. A transformer turns one DataFrame into another with transform(). An estimator learns something from data with fit() and returns a model, which is itself a transformer.
A Pipeline chains these stages in order. It is an estimator too: calling fit() on it runs every stage on the training data and returns a PipelineModel you apply to new data with transform().
How it works
Estimators versus transformers
| Transformer | Estimator | |
|---|---|---|
| Main method | transform(df) | fit(df) |
| Returns | A new DataFrame with added columns | A fitted model (a transformer) |
| Learns from data | No | Yes |
| Examples | VectorAssembler, fitted models such as LogisticRegressionModel | StringIndexer, OneHotEncoder, Imputer, StandardScaler, LogisticRegression, RandomForestClassifier |
The rule of thumb: if a stage must look at the data to know what to do (category list, mean, scale, weights), it is an estimator. VectorAssembler only packs columns into one vector, so it needs nothing from the data and is a transformer.
| Estimator | fit() returns | What it learned |
|---|---|---|
| StringIndexer | StringIndexerModel | The category to index mapping |
| Imputer | ImputerModel | The mean, median or mode per column |
| StandardScaler | StandardScalerModel | Column standard deviations (and means if centering) |
| LogisticRegression | LogisticRegressionModel | Coefficients and intercept |
| Pipeline | PipelineModel | Every fitted stage in order |
Building a training pipeline
A typical pipeline splits the data, fits preprocessing and the model on the training part only, and scores the test part with the same fitted stages.
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.regression import RandomForestRegressor
train_df, test_df = df.randomSplit([0.8, 0.2], seed=42)
idx = StringIndexer(inputCol="city", outputCol="city_idx", handleInvalid="keep")
vec = VectorAssembler(inputCols=["city_idx", "rooms"], outputCol="features")
rf = RandomForestRegressor(featuresCol="features", labelCol="price")
model = Pipeline(stages=[idx, vec, rf]).fit(train_df)Then model.transform(test_df) adds a prediction column. The StringIndexer inside the PipelineModel uses the mapping learned from training data, so test rows are encoded the same way.
handleInvalid="keep" gives unseen categories their own index. The default is "error", which fails on a city that was not in the training data; "skip" drops those rows.
Stage order and common preprocessing
| Order | Stage | Why |
|---|---|---|
| 1 | Imputer | Fill numeric nulls (mean, median or mode) |
| 2 | StringIndexer | Turn string categories into indexes |
| 3 | OneHotEncoder | Turn indexes into vectors, for linear models |
| 4 | VectorAssembler | Pack all features into one vector column |
| 5 | StandardScaler | Scale the vector, for linear models with regularization |
| 6 | Estimator (the model) | Learn from the features column |
Tree models in Spark can use the StringIndexer output directly, because the indexer marks the column as categorical. The tree's maxBins (default 32) must then be at least the number of categories in any categorical feature, or training fails.
Why a pipeline matters
Fitting each stage by hand on the full dataset leaks test information into training, for example an imputer mean that includes test rows. A pipeline fitted on train_df learns every statistic from training data only.
The fitted PipelineModel is one object, so you can log it with mlflow.spark.log_model or save it with model.write().overwrite().save(path). Scoring later reuses exactly the same preprocessing, with no copy of the feature code to keep in step.
The same idea exists in scikit-learn: sklearn.pipeline.Pipeline with fit() and predict(), used for single node models.
How the exam asks about this
- "Which stage returns a model when fit is called?" Any estimator, such as StringIndexer or LogisticRegression.
- "What does Pipeline.fit(train_df) return?" A PipelineModel, which is a transformer.
- "Scoring fails because the test data contains a new category." Set handleInvalid to keep or skip on the StringIndexer.
What to remember for the exam
- Transformers have
transform()and learn nothing; estimators havefit()and return a model. - A fitted model is a transformer, so it is applied with
transform(). - A Pipeline is an estimator; its
fit()returns a PipelineModel. - VectorAssembler is a transformer; StringIndexer, OneHotEncoder, Imputer and StandardScaler are estimators.
- Fit the pipeline on training data only, then transform the test data.
- StringIndexer handleInvalid defaults to error; keep and skip are the alternatives.
Do not call fit() on a PipelineModel or transform() on an unfitted Pipeline. The Pipeline is fitted once; the PipelineModel it returns is what scores data.
Common mistakes
Calling VectorAssembler an estimator. It only combines columns into a vector and learns nothing from the data. It is a transformer and has no fit() step.
Fitting preprocessing on the whole dataset. An imputer or scaler fitted before the split learns from test rows. Split first, then fit the pipeline on the training part.
Forgetting new categories at scoring time. A StringIndexer with the default handleInvalid raises an error on unseen labels. Choose keep or skip when production data may contain new values.
Check yourself
Is StringIndexer an estimator or a transformer?
An estimator, because it must read the data to learn the category to index mapping; fit() returns a StringIndexerModel.
What does calling fit on a Pipeline return?
A PipelineModel, a transformer that holds every fitted stage and scores new data with transform().
Why fit the pipeline on training data only?
So statistics such as imputer means and category mappings come only from training rows, which avoids leaking test information into the model.
Topic 2 of 5 in Model Development
Test yourself on Model Development
The full Databricks Machine Learning Associate question bank has 600 questions, with timed practice exams and an explanation for every answer.
This study guide is independent and is not affiliated with, endorsed by, or sponsored by Databricks. Vendor names, certification names and exam codes are trademarks of their respective owners and are used for identification only.