🤖 ML Pipeline Integration in PySpark

Pipeline Architecture

Architecture Diagram

Detailed Explanation

Core Design Pattern: Transformers and Estimators

PySpark MLlib distinguishes between two fundamental abstractions:

Key Takeaway: The Pipeline class chains stages together where output of one stage feeds into the input of the next.

Feature Engineering Transformers

Estimator Pattern: Fit-Transform Paradigm

1. Estimator accepts training DataFrame
2. Produces Model (which is itself a Transformer)
3. Model applies to new data via transform() method
4. Persists learned parameters (coefficients, intercepts)

Pipeline Composition

When fit() is called on a Pipeline:

1. Executes each stage sequentially
2. Fits all Estimators to produce Models
3. Applies all Transformers to the data
4. Entire Pipeline is serializable for persistence

Best Practice: Persist fitted Pipelines to disk for inference without retraining.

Hyperparameter Tuning

Evaluation Metrics by Task

Important: Match metric to business objective—imbalanced datasets require precision/recall over accuracy.

Mathematical Foundations

Definition: ML Pipeline

An ML pipeline $P$ is a directed acyclic graph of stages $P = (V, E)$ where vertices $V = \{S_1, S_2, \ldots, S_n\}$ are transformation stages and edges $E$ define data dependencies. Each stage $S_i$ implements a function $f_i: X_i \rightarrow Y_i$.

Cross-Validation Error

For $k$-fold cross-validation with dataset $D$ split into $k$ equal partitions $\{D_1, \ldots, D_k\}$:

where $\mathcal{L}$ is the loss function and $f(D \setminus D_i)$ is the model trained on all folds except $i$.

Bias-Variance Decomposition

The expected prediction error of model $\hat{f}$ decomposes as:

Regularization Strength

For L2-regularized regression with parameter $\lambda$:

The effective degrees of freedom reduce as $\lambda$ increases: $\text{df}(\lambda) = \sum_{j=1}^{p} \frac{d_j^2}{d_j^2 + \lambda}$

Pipeline Complexity

Total pipeline computation cost with $n$ stages, each stage $i$ processing $m_i$ records:

where $c_i$ is the per-record cost of stage $i$.

Key Insight

Spark ML pipelines distribute each stage independently, but data shuffle between stages creates synchronization barriers. Optimal pipeline design minimizes shuffles by combining compatible transformations within stages.

Summary

ML pipelines formalize ML workflows as DAGs of transformations. Cross-validation estimates generalization error, bias-variance decomposition guides model selection, and regularization trades off fit complexity against generalization. Pipeline cost scales linearly with stages and data volume.

Key Concepts Table

Code Examples

Complete Classification Pipeline

from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.feature import (
    StringIndexer, OneHotEncoder, VectorAssembler,
    StandardScaler, Imputer
)
from pyspark.ml.classification import (
    LogisticRegression, RandomForestClassifier,
    GBTClassifier, LinearSVC
)
from pyspark.ml.evaluation import (
    BinaryClassificationEvaluator,
    MulticlassClassificationEvaluator
)
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder
from pyspark.sql.functions import *

spark = SparkSession.builder \
    .appName("MLPipeline") \
    .config("spark.jars.packages", "org.apache.spark:spark-mllib_2.12:3.5.0") \
    .getOrCreate()

# Load and prepare data
df = spark.read.parquet("/data/customer_churn")

# Define feature columns
categorical_cols = ["gender", "contract_type", "payment_method"]
numerical_cols = ["tenure_months", "monthly_charges", "total_charges"]
label_col = "churn"

# Handle missing values
imputer = Imputer(
    inputCols=numerical_cols,
    outputCols=[f"{col}_imputed" for col in numerical_cols],
    strategy="median"
)

# Index categorical columns
indexers = [
    StringIndexer(
        inputCol=col,
        outputCol=f"{col}_index",
        handleInvalid="keep"
    )
    for col in categorical_cols
]

# One-hot encode indexed columns
encoders = [
    OneHotEncoder(
        inputCol=f"{col}_index",
        outputCol=f"{col}_encoded",
        dropLast=True
    )
    for col in categorical_cols
]

# Assemble all features into a single vector
assembler = VectorAssembler(
    inputCols=[f"{col}_imputed" for col in numerical_cols] +
              [f"{col}_encoded" for col in categorical_cols],
    outputCol="features_raw"
)

# Scale features
scaler = StandardScaler(
    inputCol="features_raw",
    outputCol="features",
    withStd=True,
    withMean=True
)

# Define classifiers
lr = LogisticRegression(
    featuresCol="features",
    labelCol=label_col,
    maxIter=100,
    regParam=0.01,
    elasticNetParam=0.5
)

rf = RandomForestClassifier(
    featuresCol="features",
    labelCol=label_col,
    numTrees=100,
    maxDepth=10,
    featureSubsetStrategy="sqrt"
)

gbt = GBTClassifier(
    featuresCol="features",
    labelCol=label_col,
    maxIter=100,
    maxDepth=5,
    stepSize=0.1
)

# Build pipeline
pipeline = Pipeline(stages=[
    imputer,
    *indexers,
    *encoders,
    assembler,
    scaler,
    lr
])

# Hyperparameter tuning
paramGrid = ParamGridBuilder() \
    .addGrid(lr.regParam, [0.01, 0.1, 1.0]) \
    .addGrid(lr.elasticNetParam, [0.0, 0.5, 1.0]) \
    .addGrid(lr.maxIter, [50, 100, 200]) \
    .build()

evaluator = BinaryClassificationEvaluator(
    labelCol=label_col,
    metricName="areaUnderROC"
)

crossval = CrossValidator(
    estimator=pipeline,
    estimatorParamMaps=paramGrid,
    evaluator=evaluator,
    numFolds=5,
    parallelism=4,
    seed=42
)

# Train with cross-validation
cv_model = crossval.fit(df)

# Evaluate on test set
predictions = cv_model.transform(test_df)
auc = evaluator.evaluate(predictions)
print(f"AUC-ROC: {auc:.4f}")

# Detailed metrics
accuracy_evaluator = MulticlassClassificationEvaluator(
    labelCol=label_col, metricName="accuracy"
)
f1_evaluator = MulticlassClassificationEvaluator(
    labelCol=label_col, metricName="f1"
)

print(f"Accuracy: {accuracy_evaluator.evaluate(predictions):.4f}")
print(f"F1 Score: {f1_evaluator.evaluate(predictions):.4f}")

Feature Engineering Pipeline

from pyspark.ml.feature import (
    Bucketizer, QuantileDiscretizer, SQLTransformer,
    Interaction, ChiSqSelector, PCA
)

# Bucket continuous features
bucketizer = Bucketizer(
    splits=[0, 18, 30, 45, 60, 100],
    inputCol="age",
    outputCol="age_bucket"
)

# SQL-based feature engineering
sql_transformer = SQLTransformer(
    statement="""
        SELECT
            *,
            monthly_charges / NULLIF(tenure_months, 0) AS avg_monthly_spend,
            CASE
                WHEN total_charges > 1000 THEN 'high_value'
                WHEN total_charges > 500 THEN 'medium_value'
                ELSE 'low_value'
            END AS value_segment
        FROM __THIS__
    """
)

# Create interaction features
interaction = Interaction(
    inputCols=["gender_encoded", "contract_type_encoded"],
    outputCol="gender_contract_interaction"
)

# Feature selection using Chi-Squared
selector = ChiSqSelector(
    featuresCol="features",
    labelCol=label_col,
    numTopFeatures=20,
    outputCol="selected_features"
)

# Dimensionality reduction with PCA
pca = PCA(
    k=10,
    inputCol="features",
    outputCol="pca_features"
)

# Complete feature pipeline
feature_pipeline = Pipeline(stages=[
    sql_transformer,
    imputer,
    *indexers,
    *encoders,
    interaction,
    assembler,
    scaler,
    selector,
    pca
])

feature_model = feature_pipeline.fit(df)
feature_df = feature_model.transform(df)
feature_df.select("features", "pca_features", "selected_features").show(5, truncate=False)

Performance Metrics

Best Practices

1. Always use Pipeline to ensure reproducibility and prevent data leakage from test set into training transformations
2. Fit all preprocessors on training data only and apply the fitted preprocessors to test data using the same fitted model
3. Use handleInvalid="keep" in StringIndexer to gracefully handle unseen categories at inference time
4. Set a random seed in all stochastic algorithms and cross-validation to ensure reproducible results
5. Use CrossValidator with parallelism > 1 to leverage Spark's distributed computing for faster tuning
6. Scale features before training models sensitive to feature magnitudes (logistic regression, SVM, neural networks)
7. Monitor overfitting by comparing training and validation metrics—if gap exceeds 5%, reduce model complexity
8. Use featureSubsetStrategy="sqrt" in RandomForest to decorrelate trees and reduce overfitting
9. Persist intermediate DataFrames when building complex pipelines to avoid recomputation
10. Log all experiments with MLflow or similar tracking system, including data versions, hyperparameters, and metrics
11. Use VectorAssembler before StandardScaler to ensure proper feature normalization
12. Implement stratified sampling for imbalanced datasets using df.sampleBy() to ensure balanced class representation in training

See also: ML Feature Engineering (32), Model Deployment (33), Production Hardening (40)

See Also

• ML Feature Engineering — Feature engineering with Spark ML
• Model Deployment — Model deployment patterns
• Structured Streaming — Real-time ML scoring
• Data Quality — Data quality validation