🤖 ML Pipeline Integration in PySpark

Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────┐
│                    PYSPARK ML PIPELINE ARCHITECTURE                     │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   ┌──────────────┐     ┌──────────────┐     ┌──────────────┐          │
│   │   Raw Data   │────▶│  Feature     │────▶│  Model       │          │
│   │  (DataFrame) │     │  Engineering │     │  Training    │          │
│   └──────────────┘     └──────┬───────┘     └──────┬───────┘          │
│                               │                     │                   │
│                               ▼                     ▼                   │
│                    ┌──────────────────┐   ┌──────────────────┐         │
│                    │  Transformers    │   │  Estimators      │         │
│                    │  ────────────    │   │  ────────────    │         │
│                    │  StringIndexer   │   │  LogisticReg     │         │
│                    │  OneHotEncoder   │   │  RandomForest    │         │
│                    │  VectorAssembler │   │  GBTClassifier   │         │
│                    │  StandardScaler  │   │  CrossValidator  │         │
│                    └────────┬─────────┘   └────────┬─────────┘         │
│                             │                      │                    │
│                             ▼                      ▼                    │
│                    ┌──────────────────┐   ┌──────────────────┐         │
│                    │  Feature         │   │  Model           │         │
│                    │  Pipeline        │   │  Persistence     │         │
│                    │  (Chained        │   │  (Save/Load)     │         │
│                    │   stages)        │   │                  │         │
│                    └────────┬─────────┘   └────────┬─────────┘         │
│                             │                      │                    │
│                             ▼                      ▼                    │
│                    ┌──────────────────┐   ┌──────────────────┐         │
│                    │  Evaluation      │   │  Deployment      │         │
│                    │  ─────────────   │   │  ─────────────   │         │
│                    │  BinaryClassif   │   │  MLflow          │         │
│                    │  Multiclassif    │   │  Airflow         │         │
│                    │  Regression      │   │  Kubernetes      │         │
│                    └──────────────────┘   └──────────────────┘         │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────┐
│                  TRANSFORMER vs ESTIMATOR PATTERN                       │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   TRANSFORMER (DataFrame ──▶ DataFrame)                                │
│   ┌─────────────────────────────────────────────────────────────────┐  │
│   │  Input DataFrame                                                │  │
│   │  ┌─────┬─────┬─────┬─────┐                                    │  │
│   │  │ A   │ B   │ C   │ D   │                                    │  │
│   │  ├─────┼─────┼─────┼─────┤                                    │  │
│   │  │ 1.0 │ 2.0 │ 3.0 │ 4.0 │                                    │  │
│   │  └─────┴─────┴─────┴─────┘                                    │  │
│   │       │                                                        │  │
│   │       ▼  StandardScaler.transform()                            │  │
│   │       │                                                        │  │
│   │  ┌─────┬─────┬─────┬─────┬─────┐                             │  │
│   │  │ A   │ B   │ C   │ D   │ features_scaled │                 │  │
│   │  ├─────┼─────┼─────┼─────┼─────┤                             │  │
│   │  │ 1.0 │ 2.0 │ 3.0 │ 4.0 │ [0.27,0.53,   │                 │  │
│   │  │     │     │     │     │  0.80,1.07]    │                 │  │
│   │  └─────┴─────┴─────┴─────┴─────┘                             │  │
│   └─────────────────────────────────────────────────────────────────┘  │
│                                                                         │
│   ESTIMATOR (DataFrame ──▶ Model ──▶ DataFrame)                        │
│   ┌─────────────────────────────────────────────────────────────────┐  │
│   │  Training Data                                                  │  │
│   │  ┌─────┬─────────────────┬──────┐                              │  │
│   │  │ id  │ features        │label │                              │  │
│   │  ├─────┼─────────────────┼──────┤                              │  │
│   │  │ 1   │ [1.0, 2.0, 3.0] │  0   │                              │  │
│   │  │ 2   │ [4.0, 5.0, 6.0] │  1   │                              │  │
│   │  └─────┴─────────────────┴──────┘                              │  │
│   │       │                                                        │  │
│   │       ▼  LogisticRegression.fit()                              │  │
│   │       │                                                        │  │
│   │  ┌─────────────────────────────────────┐                      │  │
│   │  │  Model (LogisticRegressionModel)     │                      │  │
│   │  │  ──────────────────────────────────  │                      │  │
│   │  │  coefficients: [0.1, 0.2, 0.3]      │                      │  │
│   │  │  intercept: -0.5                     │                      │  │
│   │  └─────────────────────────────────────┘                      │  │
│   │       │                                                        │  │
│   │       ▼  model.transform(testData)                             │  │
│   │       │                                                        │  │
│   │  ┌─────┬────────────┬──────┬──────────┐                      │  │
│   │  │ id  │ features   │label │ prediction│                      │  │
│   │  ├─────┼────────────┼──────┼──────────┤                      │  │
│   │  │ 1   │ [1,2,3]    │  0   │    0      │                      │  │
│   │  │ 2   │ [4,5,6]    │  1   │    1      │                      │  │
│   │  └─────┴────────────┴──────┴──────────┘                      │  │
│   └─────────────────────────────────────────────────────────────────┘  │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────┐
│                 CROSS-VALIDATION & HYPERPARAMETER TUNING                │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   Dataset                                                               │
│   ┌─────────────────────────────────────────────────────────────────┐  │
│   │ Fold 1    │ Fold 2    │ Fold 3    │ Fold 4    │ Fold 5         │  │
│   │ ┌───────┐ │ ┌───────┐ │ ┌───────┐ │ ┌───────┐ │ ┌───────┐     │  │
│   │ │ Train │ │ │ Train │ │ │ Train │ │ │ Train │ │ │ Train │     │  │
│   │ │  80%  │ │ │  80%  │ │ │  80%  │ │ │  80%  │ │ │  80%  │     │  │
│   │ ├───────┤ │ ├───────┤ │ ├───────┤ │ ├───────┤ │ ├───────┤     │  │
│   │ │  Test │ │ │  Test │ │ │  Test │ │ │  Test │ │ │  Test │     │  │
│   │ │  20%  │ │ │  20%  │ │ │  20%  │ │ │  20%  │ │ │  20%  │     │  │
│   │ └───────┘ │ └───────┘ │ └───────┘ │ └───────┘ │ └───────┘     │  │
│   └─────────────────────────────────────────────────────────────────┘  │
│                                                                         │
│   For each fold:                                                       │
│   ┌─────────────────────────────────────────────────────────────────┐  │
│   │  Hyperparameter Grid:                                           │  │
│   │  ┌─────────────┬─────────────┬─────────────┐                  │  │
│   │  │ regParam    │ elasticNet  │ maxIter     │                  │  │
│   │  ├─────────────┼─────────────┼─────────────┤                  │  │
│   │  │ 0.01        │ 0.0         │ 100         │  ← Config 1     │  │
│   │  │ 0.1         │ 0.5         │ 100         │  ← Config 2     │  │
│   │  │ 1.0         │ 1.0         │ 200         │  ← Config 3     │  │
│   │  └─────────────┴─────────────┴─────────────┘                  │  │
│   │                                                                │  │
│   │  For each config × fold:                                       │  │
│   │    Train on (K-1) folds → Evaluate on held-out fold           │  │
│   │                                                                │  │
│   │  Select config with best average metric across all folds      │  │
│   └─────────────────────────────────────────────────────────────────┘  │
│                                                                         │
│   ┌─────────────────────────────────────────────────────────────────┐  │
│   │  Results:                                                       │  │
│   │  Config 1: Avg AUC = 0.892                                     │  │
│   │  Config 2: Avg AUC = 0.934  ← Best                             │  │
│   │  Config 3: Avg AUC = 0.921                                     │  │
│   └─────────────────────────────────────────────────────────────────┘  │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

Detailed Explanation

PySpark MLlib provides a high-level API built on DataFrames that enables the construction of machine learning pipelines. The fundamental design pattern distinguishes between Transformers (which transform DataFrames) and Estimators (which learn from DataFrames to produce Transformers). This abstraction enables the composition of complex workflows through the Pipeline class, which chains stages together in a linear sequence where the output of one stage feeds into the input of the next.

Feature engineering in PySpark is primarily accomplished through Transformers. The StringIndexer converts categorical string columns into numerical indices, and the OneHotEncoder transforms those indices into binary vectors. The VectorAssembler combines multiple columns into a single feature vector, which is the required input format for all MLlib algorithms. The StandardScaler normalizes features to zero mean and unit variance, which is essential for algorithms sensitive to feature scales like logistic regression and SVM.

The Estimator pattern implements the fit-transform paradigm. An Estimator (like LogisticRegression) accepts a training DataFrame and produces a Model (like LogisticRegressionModel), which is itself a Transformer. This model can then be applied to new data via the transform method. The model persists the learned parameters (coefficients, intercepts, etc.) and encapsulates the complete inference logic.

Pipeline composition enables reproducible workflows. When you call fit() on a Pipeline, it executes each stage sequentially: first fitting all Estimators to produce Models, then applying all Transformers (including the produced Models) to the data. The entire Pipeline is serializable, meaning the fitted Pipeline (including all intermediate Transformers and Models) can be persisted to disk and loaded for inference without retraining.

Hyperparameter tuning is accomplished through CrossValidator and TrainValidationSplit. CrossValidator performs K-fold cross-validation, training K models on different data splits and evaluating on the held-out folds. The ParamGridBuilder defines the search space over hyperparameters. The evaluator (BinaryClassificationEvaluator, RegressionEvaluator, etc.) measures model quality. The best model is selected based on average metric across folds, and this model is then refit on the entire training dataset.

The evaluation metrics are crucial for model selection. For binary classification, PySpark provides BinaryClassificationEvaluator with metrics like AUC-ROC and AUC-PR. For multiclass classification, MulticlassClassificationEvaluator provides accuracy, precision, recall, F1 score, and log-loss. For regression, RegressionEvaluator provides MSE, RMSE, MAE, and R-squared. Understanding which metric aligns with your business objective is essential—imbalanced datasets require precision/recall over accuracy, while cost-sensitive applications need custom weight configurations.

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: Snowflake Time Travel (snowflake/02), Kafka CDC (kafka/04), Airflow DAGs (airflow/02)