Advanced

Spark ML Pipelines

Build reproducible, end-to-end ML workflows by chaining transformers and estimators into Pipelines with cross-validation and hyperparameter tuning.

What is a Pipeline?

A Pipeline chains multiple Transformers and Estimators together into a single workflow. When you call fit() on a Pipeline, it runs each stage in order - fitting estimators and transforming data - and produces a PipelineModel that can be used for prediction.

Building a Pipeline

Python - Complete ML Pipeline
from pyspark.ml import Pipeline
from pyspark.ml.feature import (
    StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler
)
from pyspark.ml.classification import LogisticRegression

# Stage 1: Index categorical columns
category_indexer = StringIndexer(inputCol="category", outputCol="category_idx")
city_indexer = StringIndexer(inputCol="city", outputCol="city_idx")

# Stage 2: One-hot encode
encoder = OneHotEncoder(
    inputCols=["category_idx", "city_idx"],
    outputCols=["category_vec", "city_vec"]
)

# Stage 3: Assemble features
assembler = VectorAssembler(
    inputCols=["age", "income", "category_vec", "city_vec"],
    outputCol="raw_features"
)

# Stage 4: Scale features
scaler = StandardScaler(inputCol="raw_features", outputCol="features")

# Stage 5: Train classifier
lr = LogisticRegression(maxIter=100, regParam=0.01)

# Build the pipeline
pipeline = Pipeline(stages=[
    category_indexer, city_indexer, encoder,
    assembler, scaler, lr
])

# Fit the entire pipeline
model = pipeline.fit(train_df)

# Predict with all stages applied automatically
predictions = model.transform(test_df)

Cross-Validation

Python - CrossValidator with Pipeline
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder
from pyspark.ml.evaluation import BinaryClassificationEvaluator

# Define parameter grid
paramGrid = ParamGridBuilder() \
    .addGrid(lr.regParam, [0.001, 0.01, 0.1]) \
    .addGrid(lr.elasticNetParam, [0.0, 0.5, 1.0]) \
    .addGrid(lr.maxIter, [50, 100]) \
    .build()

# Set up cross-validator
cv = CrossValidator(
    estimator=pipeline,
    estimatorParamMaps=paramGrid,
    evaluator=BinaryClassificationEvaluator(metricName="areaUnderROC"),
    numFolds=5,
    parallelism=4,
    seed=42
)

# Run cross-validation
cv_model = cv.fit(train_df)

# Best model performance
best_auc = cv_model.avgMetrics[
    cv_model.avgMetrics.index(max(cv_model.avgMetrics))
]
print(f"Best AUC: {best_auc:.4f}")

# Use the best model for predictions
predictions = cv_model.transform(test_df)

Train-Validation Split

Python - TrainValidationSplit (faster than CV)
from pyspark.ml.tuning import TrainValidationSplit

tvs = TrainValidationSplit(
    estimator=pipeline,
    estimatorParamMaps=paramGrid,
    evaluator=BinaryClassificationEvaluator(),
    trainRatio=0.8,
    parallelism=4,
    seed=42
)

tvs_model = tvs.fit(train_df)
predictions = tvs_model.transform(test_df)

Saving and Loading Pipelines

Python - Pipeline Persistence
# Save the fitted pipeline model
model.write().overwrite().save("hdfs:///models/my_pipeline_v1")

# Load it back
from pyspark.ml import PipelineModel
loaded_model = PipelineModel.load("hdfs:///models/my_pipeline_v1")

# Use for predictions
new_predictions = loaded_model.transform(new_data)

# Save the unfitted pipeline (for retraining)
pipeline.write().overwrite().save("hdfs:///pipelines/my_pipeline")

# Load unfitted pipeline
from pyspark.ml import Pipeline
loaded_pipeline = Pipeline.load("hdfs:///pipelines/my_pipeline")
new_model = loaded_pipeline.fit(updated_train_df)
✅
Use TrainValidationSplit for large datasets: Cross-validation trains k models per parameter combination. For large datasets, TrainValidationSplit is faster since it only trains once per combination. Use CV for smaller datasets where variance estimation matters.
💡
Pipeline + MLflow: Spark ML pipelines integrate seamlessly with MLflow. Use mlflow.spark.log_model(model, "spark-model") to log your fitted PipelineModel for versioning and deployment.

Ready to Go Deeper?

Live instructor-led courses from our partners. Affiliate disclosure.

🧠

AI & ML Courses - 30% Off

Live instructor-led AI, machine learning, data science, and cloud courses for working professionals. Use code Limited30 at checkout.

Edureka
📊

DataCamp - AI & Data Science

Hands-on Python, machine learning, and AI courses with interactive exercises and real projects.

DataCamp
🎓

edX - Top AI Courses

University-level AI courses from MIT, Harvard, Stanford. Earn certificates that employers recognize.

edX