PySpark MLlib
Train machine learning models on billions of rows using Spark's distributed ML library.
Start Learning PySpark MLlibWhat You'll Learn
- The MLlib Pipeline architecture — Transformers and Estimators
- Feature engineering at scale: StringIndexer, OneHotEncoder, VectorAssembler
- Training classification models: LogisticRegression, RandomForestClassifier
- Regression models: LinearRegression, GBTRegressor
- Model evaluation: BinaryClassificationEvaluator, RegressionEvaluator
- Cross-validation and hyperparameter tuning with ParamGridBuilder
- Clustering with KMeans and recommendation with ALS
- Saving and loading models for production deployment
Introduction to PySpark MLlib
PySpark MLlib is Apache Spark's built-in machine learning library. While scikit-learn is excellent for training models on data that fits in memory (millions of rows), MLlib is designed for data that doesn't fit on a single machine — hundreds of millions or billions of rows distributed across a cluster. It uses the same Spark DataFrame API you already know, making the transition from scikit-learn intuitive.
MLlib's core design is the Pipeline — a sequence of stages (feature transformers + a model estimator) that can be fit and then applied to new data in one object. This makes ML workflows reproducible and easy to deploy. A typical pipeline includes: StringIndexer (encode categorical columns), VectorAssembler (combine feature columns into one vector column), StandardScaler (normalise features), and a model like LogisticRegression or RandomForestClassifier.
MLlib supports classification, regression, clustering, recommendation (ALS), and feature extraction. In 2024-2025, many teams combine MLlib for feature engineering and preprocessing at scale with other tools (scikit-learn, XGBoost on Spark) for the final model training, taking the best of both worlds.
Video Tutorials
Handpicked free YouTube videos to accelerate your understanding
PySpark Machine Learning Full Tutorial
End-to-end PySpark MLlib tutorial — build classification pipelines with StringIndexer, VectorAssembler, and RandomForestClassifier.
Machine Learning with PySpark — Full Course
Complete PySpark ML course: classification, regression, clustering with KMeans, cross-validation, and ALS recommendation systems.
End-to-end classification pipeline with PySpark MLlib
Copy the code below and paste it into your Python environment or our free online compiler.
from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.feature import (
StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler
)
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
spark = SparkSession.builder.appName("MLlibDemo").getOrCreate()
# 1. Load data (predict if a customer will churn)
df = spark.read.csv("customer_churn.csv", header=True, inferSchema=True)
# 2. Encode categorical column "plan" → index → one-hot vector
indexer = StringIndexer(inputCol="plan", outputCol="plan_idx")
encoder = OneHotEncoder(inputCol="plan_idx", outputCol="plan_vec")
# 3. Combine all feature columns into one vector
assembler = VectorAssembler(
inputCols=["age", "tenure_months", "monthly_bill", "plan_vec"],
outputCol="features_raw"
)
# 4. Normalise features (zero mean, unit variance)
scaler = StandardScaler(
inputCol="features_raw",
outputCol="features"
)
# 5. Define the model (label = churn column: 1=churned, 0=stayed)
rf = RandomForestClassifier(
labelCol="churn",
featuresCol="features",
numTrees=100
)
# 6. Build the pipeline
pipeline = Pipeline(stages=[indexer, encoder, assembler, scaler, rf])
# 7. Split data
train, test = df.randomSplit([0.8, 0.2], seed=42)
# 8. Train the model
model = pipeline.fit(train)
# 9. Make predictions
predictions = model.transform(test)
predictions.select("churn", "prediction", "probability").show(5)
# 10. Evaluate
evaluator = BinaryClassificationEvaluator(
labelCol="churn", metricName="areaUnderROC"
)
auc = evaluator.evaluate(predictions)
print(f"AUC-ROC: {auc:.4f}")
# 11. Save the trained pipeline
model.save("/models/churn_rf_v1")Key Concepts Explained
Master these terms and you'll understand 80% of the conversations in this field.
Pipeline
An ordered sequence of stages (Transformers + one Estimator). Call pipeline.fit(train) to train, then model.transform(test) to predict. Ensures preprocessing is applied consistently to train and test data.
Transformer
A stage that transforms a DataFrame — adds new columns without learning from data. Examples: VectorAssembler, OneHotEncoder, StandardScaler, Tokenizer.
Estimator
A stage that learns from data during fit() and then transforms data during transform(). Examples: LogisticRegression, RandomForestClassifier, KMeans.
VectorAssembler
Combines multiple numeric columns into a single "features" vector column that ML algorithms expect. This is almost always the final transformation before the model.
StringIndexer
Converts a categorical string column to numeric indices (0, 1, 2...) sorted by frequency. "Apple"=0, "Banana"=1, "Cherry"=2. Required before OneHotEncoder.
CrossValidator
Performs k-fold cross-validation for hyperparameter tuning. Combines a pipeline, a ParamGrid of hyperparameter values to try, and an evaluator to pick the best model.
ALS (Alternating Least Squares)
Spark MLlib's collaborative filtering algorithm for building recommendation systems. Learns user and item factor matrices to predict ratings.
PipelineModel.save()
Persists the entire fitted pipeline (including learned parameters and weights) to disk or object storage. Load it later with PipelineModel.load() for serving.
Your PySpark MLlib Learning Path
Follow these steps in order — each one builds on the last. Designed for complete beginners.
- 1
PySpark DataFrames
Master reading data, transformations, and groupBy operations. MLlib works entirely with Spark DataFrames.
- 2
ML Fundamentals
Understand supervised learning, classification vs regression, train/test splits, and evaluation metrics (accuracy, AUC, RMSE).
- 3
Feature Engineering
Practice StringIndexer, OneHotEncoder, and VectorAssembler. Understanding these is 70% of MLlib work.
- 4
First Pipeline
Build a Logistic Regression pipeline to classify a dataset. Evaluate with BinaryClassificationEvaluator.
- 5
Tree-Based Models
Train RandomForest and Gradient Boosted Trees. Compare performance with Logistic Regression.
- 6
Hyperparameter Tuning
Use ParamGridBuilder and CrossValidator to systematically search for the best hyperparameters.
- 7
Deployment
Save your PipelineModel. Load it in a separate Spark job. Serve predictions as a batch job on a schedule.
Ready to master PySpark MLlib?
Explore our free tutorials, hands-on code examples, and interview questions. No sign-up. No paywalls. Forever free.