Spark MLlib basics: a practical Spark & Data Eng guide

Spark MLlib basics: a practical Spark & Data Eng guide

When your dataset grows beyond what a single machine can hold in memory, scikit-learn and pandas stop being viable options. Machine learning at scale requires a different architecture, and that's where Spark MLlib comes in. We'll walk through the practical patterns that make distributed ML feasible on production data pipelines.

Why MLlib, and when does it matter

MLlib is Spark's machine learning library, built on top of RDDs and DataFrames. It handles the distribution headache for you—training models on terabytes of data across a cluster without writing explicit map-reduce logic. That said, it's not always the right tool. We use it when our training dataset is genuinely too large for one machine, when we need to retrain frequently as new data arrives, or when feature engineering happens naturally in a Spark pipeline before model training.

If your data fits on a single node, scikit-learn or XGBoost will often be faster and more mature. If you need cutting-edge models, TensorFlow or PyTorch with Spark distribution (via Horovod or distributed training frameworks) might be better. MLlib shines in the middle ground: structured tabular data at scale, ensemble models, and tight integration with your data processing DAG.

Pipelines: Transformers and Estimators

MLlib pipelines mirror scikit-learn's design. A pipeline chains Transformers (which take a DataFrame and return a modified DataFrame) and Estimators (which learn from data and produce Transformers). Think of a Transformer as a stateless function; an Estimator as something that needs to see data first to build its state.

Here's a concrete example: we're building a churn prediction model for a subscription service. Raw events come in with user IDs, event timestamps, and feature counts. We need to assemble them into vectors, scale them, and fit a logistic regression model.

from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.feature import VectorAssembler, StandardScaler
from pyspark.ml.classification import LogisticRegression

spark = SparkSession.builder.appName("churn_ml").getOrCreate()

# Sample data: user events with numerical features
df = spark.createDataFrame([
    (1, 45.0, 12.0, 1),
    (2, 23.0, 5.0, 0),
    (3, 78.0, 25.0, 1),
], ["user_id", "days_active", "support_tickets", "churned"])

# Pipeline stages
assembler = VectorAssembler(
    inputCols=["days_active", "support_tickets"],
    outputCol="features"
)

scaler = StandardScaler(
    inputCol="features",
    outputCol="scaled_features"
)

lr = LogisticRegression(
    featuresCol="scaled_features",
    labelCol="churned",
    maxIter=10
)

# Build and fit the pipeline
pipeline = Pipeline(stages=[assembler, scaler, lr])
model = pipeline.fit(df)

# Transform new data through the entire pipeline
predictions = model.transform(df)
predictions.select("user_id", "probability", "prediction").show()

The fit() call trains all Estimators in the pipeline (scaler learns mean and std; logistic regression learns coefficients). transform() applies the entire fitted pipeline to new data. This clean separation matters in production: you save the model once, load it on serving infrastructure, and apply it consistently without recalculating statistics.

Features: VectorAssembler and Encoders

MLlib models work with dense or sparse vectors. You can't pass raw DataFrames to estimators; you need a features vector column. That's where VectorAssembler and encoders come in.

VectorAssembler concatenates numerical columns into a single vector. If you have categorical features, you first encode them (one-hot or label encoding), then assemble.

from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler

# Sample data with categorical feature
df = spark.createDataFrame([
    (1, "free", 45.0),
    (2, "pro", 120.0),
    (3, "free", 23.0),
], ["user_id", "plan_type", "days_active"])

# Encode categorical: plan_type -> planIndex (0, 1, ...)
indexer = StringIndexer(inputCol="plan_type", outputCol="planIndex")
indexed = indexer.fit(df).transform(df)

# One-hot encode: planIndex -> planVec (one-hot sparse vector)
encoder = OneHotEncoder(inputCol="planIndex", outputCol="planVec")
encoded = encoder.fit(indexed).transform(indexed)

# Assemble numeric and encoded features
assembler = VectorAssembler(
    inputCols=["planVec", "days_active"],
    outputCol="features"
)
final = assembler.transform(encoded)

final.select("user_id", "features").show(truncate=False)

Key point: Estimators like StringIndexer and OneHotEncoder must be fit on training data, then applied to test data. Always include them in your pipeline to avoid data leakage and ensure consistency.

Models: Regression, Classification, and Clustering

MLlib offers standard algorithms across three problem types:

• Regression: LinearRegression, IsotonicRegression, GeneralizedLinearRegression
• Classification: LogisticRegression, RandomForest, GradientBoostedTrees, SVM, Naive Bayes
• Clustering: K-Means, Bisecting K-Means, Gaussian Mixture Model

All follow the same pattern: instantiate, fit on training data, evaluate, predict. Here's a RandomForest classifier for fraud detection:

from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.sql.functions import col

# Assume df has "features" and "is_fraud" columns
train, test = df.randomSplit([0.8, 0.2], seed=42)

rf = RandomForestClassifier(
    featuresCol="features",
    labelCol="is_fraud",
    numTrees=50,
    maxDepth=10,
    seed=42
)

model = rf.fit(train)
predictions = model.transform(test)

# Evaluate
evaluator = BinaryClassificationEvaluator(
    labelCol="is_fraud",
    rawPredictionCol="rawPrediction"
)
auc = evaluator.evaluate(predictions)
print(f"AUC: {auc}")

# Feature importance
importance_df = spark.createDataFrame(
    [(i, model.featureImportances[i]) for i in range(len(model.featureImportances))],
    ["feature_idx", "importance"]
).orderBy(col("importance").desc())
importance_df.show()

For clustering, the pattern is similar but without a label column:

from pyspark.ml.clustering import KMeans

kmeans = KMeans(k=5, seed=42, featuresCol="features")
km_model = kmeans.fit(df)

clustered = km_model.transform(df)
clustered.select("user_id", "prediction").show()

# Centers are in km_model.clusterCenters()
print(km_model.clusterCenters())

Common mistakes

Not fitting encoders on training data only. If you fit a StringIndexer on your entire dataset before splitting train/test, category indices might differ between splits. Always fit on training data, then transform both train and test.

Ignoring feature scaling. Tree-based models (RandomForest, GBT) are scale-invariant, but linear models (LinearRegression, LogisticRegression, SVM, KMeans) are not. If you skip StandardScaler, features with larger ranges will dominate. Include it in your pipeline.

Treating model serialization casually. Spark models are large and distributed. Saving to a single file won't work. Use model.save("s3://bucket/path") and PipelineModel.load(). Always version your models and include preprocessing in the pipeline before saving.

Running pipelines on small local data without cluster testing. MLlib can be slow on small datasets due to serialization overhead. Test your pipeline locally, then scale to a real cluster. Measure end-to-end timing on representative data volumes.

Hyperparameter tuning without cross-validation. Spark provides CrossValidator and TrainValidationSplit. Use them:

from pyspark.ml.tuning import CrossValidator, ParamGridBuilder
from pyspark.ml.evaluation import RegressionEvaluator

lr = LinearRegression(featuresCol="features", labelCol="target")

param_grid = ParamGridBuilder() \
    .addGrid(lr.regParam, [0.01, 0.1, 1.0]) \
    .addGrid(lr.elasticNetParam, [0.0, 0.5, 1.0]) \
    .build()

cv = CrossValidator(
    estimator=lr,
    estimatorParamMaps=param_grid,
    evaluator=RegressionEvaluator(labelCol="target"),
    numFolds=5
)

cv_model = cv.fit(train)
best_predictions = cv_model.transform(test)

When to use Spark MLlib

Use MLlib when:

• Your training dataset is multi-terabyte and won't fit on a single machine
• You're already processing data in Spark and want to keep the pipeline unified
• You need to retrain frequently as new data streams in
• Ensemble models and tree-based algorithms meet your accuracy needs
• Latency for model training is acceptable (minutes to hours, not seconds)

Don't use MLlib if:

• Your training data fits comfortably on one machine (scikit-learn will be faster)
• You need state-of-the-art deep learning or transformer architectures
• You need sub-second model serving (use deployed models on CPUs/GPUs, not Spark)
• Your data is unstructured (images, text) without heavy feature engineering first

Key takeaways

• MLlib pipelines chain Transformers (stateless) and Estimators (learn from data), mirroring scikit-learn but for distributed data
• Always assemble and encode features into vectors before passing to models; include scaling for linear models
• Fit encoders on training data only to avoid leakage; save entire pipelines, not isolated models
• Use RandomForest, GBT, or LogisticRegression for standard supervised learning; K-Means for clustering
• Hyperparameter tune with CrossValidator on a representative data sample before production deployment
• MLlib excels at training on multi-terabyte datasets; for smaller data or cutting-edge models, other frameworks are often better