Databricks ML — Study Notes
Module 6: Model Training
1. sklearn vs Spark ML — When to Use Which
What is sklearn
scikit-learn is a Python library for ML on a single machine. All data lives in one machine's RAM as a
pandas DataFrame.
→ Data size: GBs (fits in RAM)
→ Execution: single machine, one CPU
→ Data format: pandas DataFrame
→ Best for: prototyping, small datasets, rich algorithm ecosystem
What is Spark ML
Spark ML is a distributed ML library built on Apache Spark. Data is split across multiple machines
and processed simultaneously.
→ Data size: TBs (distributed across cluster)
→ Execution: multiple machines, all cores simultaneously
→ Data format: Spark DataFrame
→ Best for: production training, large datasets, already using Spark
When to use which — exam critical
Situation Use
Data fits in RAM (GBs) sklearn — faster, simpler
Data too large for RAM (TBs) Spark ML — distributed
Quick prototyping sklearn — richer ecosystem
Production on large data Spark ML — scales linearly
Already using Spark pipeline Spark ML — no conversion needed
Complex preprocessing + model Either — both have Pipeline
Typical Databricks workflow — uses BOTH
Step 1: Data processing (Spark) → large data, use Spark
Step 2: Prototyping (sklearn) → [Link](), try algorithms
Step 3: Production (Spark ML) → train on full data distributed
Step 4: Serving (either) → depends on model size
2. sklearn vs Spark ML — Code Comparison
Data split
# sklearn
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
) # returns 4 objects, X and y separated
# Spark ML
train_df, test_df = df_spark.randomSplit([0.8, 0.2], seed=42)
# returns 2 DataFrames, features and label stay together
Feature preparation
# sklearn — multiple columns OK
X = [Link](columns=['quality']) # pandas DataFrame
# Spark ML — must combine into ONE vector column
assembler = VectorAssembler(
inputCols=feature_cols,
outputCol='features' # required single vector column
)
Model training
# sklearn — pass X and y separately
model = RandomForestRegressor(n_estimators=100, max_depth=6)
[Link](X_train, y_train)
# Spark ML — pass DataFrame, specify column names
rf = RandomForestRegressor(
featuresCol='features', # which column has features
labelCol='quality', # which column to predict
numTrees=100, maxDepth=6
)
pipeline = Pipeline(stages=[assembler, rf])
pipeline_model = [Link](train_df)
Making predictions
# sklearn — returns numpy array
predictions = [Link](X_test)
# Spark ML — returns Spark DataFrame with 'prediction' column
predictions = pipeline_model.transform(test_df)
# predictions DataFrame has all original columns + 'features' + 'prediction'
Evaluation
# sklearn — pass arrays
from [Link] import r2_score
r2 = r2_score(y_test, predictions)
# Spark ML — pass DataFrame, specify column names
evaluator = RegressionEvaluator(
labelCol='quality',
predictionCol='prediction',
metricName='r2'
)
r2 = [Link](predictions_df)
MLflow logging
# sklearn
[Link].log_model(model, artifact_path='model')
[Link].load_model(f'runs:/{run_id}/model')
# Spark ML — needs UC volume path on serverless
[Link].log_model(
spark_model=pipeline_model,
artifact_path='model',
dfs_tmpdir='/Volumes/catalog/schema/volume'
)
[Link].load_model(
f'runs:/{run_id}/model',
dfs_tmpdir='/Volumes/catalog/schema/volume'
)
Full comparison table
sklearn Spark ML
Data format pandas DataFrame Spark DataFrame
Data location single machine RAM distributed across cluster
Data size GBs TBs
Train/test split train_test_split() randomSplit([0.8, 0.2])
Feature prep manual or sklearn Pipeline VectorAssembler required
Features column multiple columns OK ONE vector column required
X and y separated stay in same DataFrame
Model training [Link](X, y) [Link](df)
Predictions [Link](X) [Link](df)
Prediction output numpy array Spark DataFrame + prediction col
Evaluation sklearn metrics functions RegressionEvaluator
MLflow flavor [Link] [Link]
Speed on small data faster slower (overhead)
Speed on big data crashes scales linearly
3. VectorAssembler
VectorAssembler combines multiple feature columns into ONE vector column. Required because
Spark ML models only accept a single vector column as input.
assembler = VectorAssembler(
inputCols=['col1', 'col2', 'col3'], # columns to combine
outputCol='features' # name of new vector column
)
What it does to your DataFrame:
Before VectorAssembler:
fixed_acidity | volatile_acidity | alcohol | quality
7.4 | 0.7 | 9.4 | 5
After VectorAssembler:
fixed_acidity | volatile_acidity | alcohol | quality | features
7.4 | 0.7 | 9.4 | 5 | [7.4, 0.7, 9.4]
↑
all values packed into ONE vector
→ inputCols must NOT include the label column (quality)
→ outputCol = 'features' is the convention (can be any name)
→ featuresCol in the model must match outputCol in assembler
→ VectorAssembler is a Transformer — no .fit() needed
⚠ EXAM TRAP: If featuresCol in the model does not match outputCol in VectorAssembler,
Spark will throw a column not found error.
4. Estimator vs Transformer — Most Tested Concept
The core difference
Transformer Estimator
What it does Applies fixed logic to data Learns from data first, then
applies
Methods .transform() only .fit() → returns Transformer,
then .transform()
Needs training data? No Yes — must fit on training data
Example VectorAssembler RandomForestRegressor
After .fit() Unchanged (already a Becomes a fitted Transformer
Transformer)
Estimators — need .fit()
→ RandomForestRegressor → learns tree splits from data
→ GBTRegressor → learns gradient boosted trees from data
→ LinearRegression → learns coefficients from data
→ LogisticRegression → learns weights from data
→ StringIndexer → learns category-to-number mapping from data
→ OneHotEncoder → learns number of categories from data
→ StandardScaler → learns mean and std from data
→ MinMaxScaler → learns min and max from data
→ Pipeline itself → fits all Estimator stages
Transformers — just .transform()
→ VectorAssembler → combines columns, no learning needed
→ Tokenizer → splits text, fixed logic
→ Bucketizer → applies fixed boundaries
→ SQLTransformer → applies SQL query
→ Fitted models → RandomForestRegressionModel etc.
What happens during [Link]()
pipeline = Pipeline(stages=[assembler, rf])
pipeline_model = [Link](train_df)
Internally:
Stage 1: [Link](train_df)
→ VectorAssembler is Transformer, just applies logic
→ adds 'features' column
→ passes result to stage 2
Stage 2: [Link](transformed_train_df)
→ RandomForestRegressor is Estimator, LEARNS from data
→ builds 100 decision trees
→ produces RandomForestRegressionModel
Returns: PipelineModel containing:
Stage 1: VectorAssembler ← unchanged Transformer
Stage 2: RandomForestRegressionModel ← new fitted Transformer
What happens during pipeline_model.transform()
predictions = pipeline_model.transform(test_df)
Internally:
Stage 1: [Link](test_df)
→ adds 'features' column to test data
Stage 2: rf_model.transform(result)
→ applies learned trees
→ adds 'prediction' column
Result: test_df + 'features' + 'prediction' columns
💡 KEY INSIGHT: Pipeline guarantees Estimators only learn from training data. No leakage into
test data possible.
5. Pipeline
Pipeline chains multiple stages in the correct order. fit() and transform() happen automatically across
all stages.
Why Pipeline matters
→ Ensures correct stage order — assembler before model
→ Prevents data leakage — Estimators only fit on training data
→ Single object for entire workflow — save, load, deploy together
→ Clean code — one fit() and one transform() call
Pipeline vs manual approach
# Manual approach — error prone
transformed_train = [Link](train_df)
rf_model = [Link](transformed_train)
transformed_test = [Link](test_df)
predictions = rf_model.transform(transformed_test)
# Pipeline approach — clean and safe
pipeline = Pipeline(stages=[assembler, rf])
pipeline_model = [Link](train_df)
predictions = pipeline_model.transform(test_df)
Stage order matters
# CORRECT — assembler creates 'features' before rf reads it
Pipeline(stages=[assembler, rf])
# WRONG — rf tries to read 'features' before it exists
Pipeline(stages=[rf, assembler]) # ERROR: column 'features' not found
6. Evaluators
Spark ML evaluators compare actual vs predicted values inside a DataFrame.
RegressionEvaluator — for predicting numbers
from [Link] import RegressionEvaluator
evaluator = RegressionEvaluator(
labelCol='quality', # actual values column
predictionCol='prediction' # predicted values column
)
r2 = [Link](preds, {[Link]: 'r2'})
rmse = [Link](preds, {[Link]: 'rmse'})
mae = [Link](preds, {[Link]: 'mae'})
BinaryClassificationEvaluator — for 0/1 prediction
from [Link] import BinaryClassificationEvaluator
evaluator = BinaryClassificationEvaluator(
labelCol='churned',
rawPredictionCol='rawPrediction' # not 'prediction'!
)
auc = [Link](predictions) # default metric: areaUnderROC
MulticlassClassificationEvaluator — for multiple classes
from [Link] import MulticlassClassificationEvaluator
evaluator = MulticlassClassificationEvaluator(
labelCol='quality',
predictionCol='prediction',
metricName='accuracy' # or 'f1', 'weightedPrecision'
)
⚠ EXAM TRAP: BinaryClassificationEvaluator uses rawPredictionCol not predictionCol. This is
a common exam trap.
7. Cross Validation + ParamGridBuilder
What is Cross Validation
Instead of one train/test split, CV does multiple splits and averages results for more reliable
evaluation.
3-Fold Cross Validation:
All data split into 3 folds:
Fold 1: rows 1-533
Fold 2: rows 534-1066
Fold 3: rows 1067-1599
Round 1: Train on Fold 2+3 → test on Fold 1 → R2=0.46
Round 2: Train on Fold 1+3 → test on Fold 2 → R2=0.48
Round 3: Train on Fold 1+2 → test on Fold 3 → R2=0.51
Average R2 = (0.46+0.48+0.51)/3 = 0.483 ← more reliable
ParamGridBuilder — define combinations to try
from [Link] import ParamGridBuilder
param_grid = ParamGridBuilder()\
.addGrid([Link], [50, 100]) # try 2 values
.addGrid([Link], [3, 6]) # try 2 values
.build()
# Creates ALL combinations:
# Combo 1: numTrees=50, maxDepth=3
# Combo 2: numTrees=50, maxDepth=6
# Combo 3: numTrees=100, maxDepth=3
# Combo 4: numTrees=100, maxDepth=6
# Total: 2 x 2 = 4 combinations
CrossValidator — combines CV + param search
from [Link] import CrossValidator
cv = CrossValidator(
estimator=pipeline, # WHAT to train
estimatorParamMaps=param_grid, # WHICH params to try
evaluator=evaluator, # HOW to measure best
numFolds=3, # HOW MANY folds
seed=42
)
cv_model = [Link](train_df)
# Runs: 4 combinations x 3 folds = 12 total fits
# Picks best combination automatically
# Retrains on FULL training data with best params
best_model = cv_model.bestModel # Pipeline with best params
predictions = cv_model.transform(test_df)
CrossValidator vs Hyperopt
CrossValidator Hyperopt
Search strategy Grid search — tries ALL Smart search — learns from trials
combinations
Built into Spark ML Separate library
Works with Spark ML Pipeline directly Any Python code
Speed Slower — exhaustive Faster — skips bad combinations
Best for Small search spaces Large search spaces
Total fits combos × folds always max_evals (can stop early)
8. Exam Questions & Answers
Q: When should you use Spark ML instead of sklearn?
A: When data is too large to fit in a single machine's RAM (TBs of data), when you need
distributed training across a cluster, or when already using Spark for data processing.
Q: What does VectorAssembler do?
A: Combines multiple feature columns into a single vector column (usually called 'features').
Required because Spark ML models only accept ONE vector column as input — not multiple
separate columns.
Q: What is the difference between an Estimator and a Transformer in Spark ML?
A: Estimator has .fit() — it learns from data and produces a fitted Transformer. Transformer
has .transform() — it applies fixed logic without learning. Examples: RandomForestRegressor is
Estimator, VectorAssembler is Transformer.
Q: What does [Link](train_df) return?
A: A PipelineModel where all Estimator stages have been fitted and converted to their Transformer
versions. Example: RandomForestRegressor becomes RandomForestRegressionModel.
Q: What does pipeline_model.transform(test_df) do?
A: Applies all fitted stages in order to test_df. Adds 'features' column (VectorAssembler) and
'prediction' column (model). Nothing is re-fitted on test data.
Q: Why must VectorAssembler come before the model in Pipeline stages?
A: Because the model reads the 'features' column that VectorAssembler creates. If model runs first,
'features' column doesn't exist yet — Spark throws a column not found error.
Q: What is randomSplit() in Spark ML?
A: Equivalent of train_test_split() for Spark DataFrames. Works distributed across cluster. Returns
2 Spark DataFrames (not 4 like sklearn). Features and label stay in same DataFrame.
Q: What MLflow flavor do you use for Spark ML models?
A: [Link] — specifically [Link].log_model() and [Link].load_model(). On
serverless/UC clusters, dfs_tmpdir must be set to a /Volumes/... path.
Q: What does ParamGridBuilder do?
A: Creates all combinations of hyperparameter values to try during cross validation. Example: 2
values for numTrees × 2 values for maxDepth = 4 combinations.
Q: How many model fits happen with 4 param combinations and 3 folds in CrossValidator?
A: 4 × 3 = 12 total fits. Each combination is evaluated across all folds, then the best is retrained on
full training data.
Q: What does cv_model.bestModel return?
A: The Pipeline retrained on full training data using the best hyperparameter combination found by
CrossValidator.
Q: What is the difference between [Link]() and [Link]()?
A: [Link]() is sklearn — returns numpy array. [Link]() is Spark ML — returns
Spark DataFrame with prediction column added.
Q: Which evaluator do you use for binary classification in Spark ML?
A: BinaryClassificationEvaluator. Note: it uses rawPredictionCol not predictionCol — a common
exam trap.
Q: List 3 Estimators and 3 Transformers in Spark ML.
A: Estimators: RandomForestRegressor, StringIndexer, StandardScaler. Transformers:
VectorAssembler, Tokenizer, Bucketizer.
9. Quick Reference — All Spark ML Commands
Command Purpose
train_df, test_df = [Link]([0.8, 0.2]) Split Spark DataFrame
VectorAssembler(inputCols=[...], Combine columns into vector
outputCol='features')
RandomForestRegressor(featuresCol='features', Define RF regressor
labelCol='y')
Pipeline(stages=[assembler, model]) Chain stages in order
[Link](train_df) Train all stages, returns PipelineModel
pipeline_model.transform(test_df) Apply all stages, adds prediction column
RegressionEvaluator(labelCol=..., predictionCol=...) Evaluate regression model
BinaryClassificationEvaluator(labelCol=..., Evaluate binary classifier
rawPredictionCol=...)
ParamGridBuilder().addGrid([Link], [v1, Create param combinations
v2]).build()
CrossValidator(estimator=pipeline, numFolds=3) Cross validation with param search
cv_model.bestModel Get best fitted Pipeline
[Link].log_model(model, Log Spark ML model
dfs_tmpdir='/Volumes/...')
[Link].load_model(uri, Load Spark ML model
dfs_tmpdir='/Volumes/...')
10. Estimator vs Transformer — Quick Reference
Estimator → .fit() learns from data → produces Transformer → .transform() applies logic
ESTIMATORS (need .fit()):
RandomForestRegressor/Classifier, GBTRegressor/Classifier
LinearRegression, LogisticRegression
StringIndexer, OneHotEncoder, StandardScaler, MinMaxScaler
Pipeline (fits all Estimator stages inside it)
TRANSFORMERS (just .transform()):
VectorAssembler, Tokenizer, Bucketizer, SQLTransformer
All FITTED models: RandomForestRegressionModel, etc.
Databricks ML Associate Exam Prep — Module 6: Model Training