Training models with Spark MLlib: LogisticRegression, RandomForest, and GBTClassifier

LEARN · DISTRIBUTED ML WITH PYSPARK

Why Spark MLlib Classification Still Matters

When your data grows beyond the memory of a single machine, moving it into pandas or NumPy just to train a model becomes expensive. Apache Spark’s spark.ml library keeps feature engineering, model training, and evaluation inside the same distributed DataFrame pipeline.

In this lesson, you’ll build a complete Spark MLlib classification workflow using three popular binary classifiers:

LogisticRegression
RandomForestClassifier
GBTClassifier

You’ll also:

Build a reusable ML Pipeline
Evaluate models with BinaryClassificationEvaluator
Compare MLlib with scikit-learn
See an empirical benchmark showing when distributed training actually wins
Learn a real Spark security lesson from a recent CVE

All examples target Apache Spark 4.0.x (PySpark 4.0.x) and Python 3.12+, which are current and fully supported.

Prerequisites

Install the latest PySpark release:

python -m venv .venv
source .venv/bin/activate

pip install pyspark

This tutorial assumes:

Python 3.12+
Apache Spark / PySpark 4.0.x
A local installation (local[*]) or a Spark cluster
A CSV containing numeric features and a binary label column

The Spark MLlib Classification Pipeline

Unlike many machine learning libraries, Spark encourages reusable pipelines.

A typical workflow is:

CSV / Parquet
      │
DataFrame
      │
Feature Engineering
      │
VectorAssembler
      │
Pipeline
      │
Classifier
      │
Predictions
      │
BinaryClassificationEvaluator

Every stage works on Spark DataFrames, making the same code portable from your laptop to a distributed cluster.

Create a Spark Session

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder
    .appName("SparkMLlibClassification")
    .master("local[*]")
    .getOrCreate()
)

Load the Dataset

df = (
    spark.read
    .option("header", True)
    .option("inferSchema", True)
    .csv("data.csv")
)

df.printSchema()
df.show(5)

Example schema:

Column	Type
feature1	double
feature2	double
feature3	double
label	integer

Prepare Features

Spark ML algorithms expect a single vector column named features.

from pyspark.ml.feature import VectorAssembler

feature_columns = [
    "feature1",
    "feature2",
    "feature3",
]

assembler = VectorAssembler(
    inputCols=feature_columns,
    outputCol="features",
)

dataset = (
    assembler
    .transform(df)
    .select("features", "label")
)

Split the data:

train_df, test_df = dataset.randomSplit(
    [0.8, 0.2],
    seed=42,
)

Cache the datasets because they’ll be reused several times:

train_df.cache()
test_df.cache()

train_df.count()
test_df.count()

Materializing the cache avoids recomputing the feature engineering stage during every model fit.

Build a Reusable ML Pipeline

Spark Pipelines make training and deployment significantly cleaner.

from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression

lr = LogisticRegression(
    featuresCol="features",
    labelCol="label",
    maxIter=100,
    regParam=0.01,
)

pipeline = Pipeline(
    stages=[
        assembler,
        lr,
    ]
)

model = pipeline.fit(df)

In production, this approach ensures identical preprocessing during both training and inference.

Train Logistic Regression

Logistic Regression remains one of the strongest baseline classifiers.

from pyspark.ml.classification import LogisticRegression

lr = LogisticRegression(
    featuresCol="features",
    labelCol="label",
    maxIter=100,
    regParam=0.01,
    elasticNetParam=0.0,
)

lr_model = lr.fit(train_df)

lr_predictions = lr_model.transform(test_df)

lr_predictions.select(
    "label",
    "prediction",
    "probability",
).show(10, truncate=False)

Logistic Regression works especially well when:

Relationships are approximately linear
Fast training matters
Model interpretability is important

Evaluate with BinaryClassificationEvaluator

For binary classification, AUC is usually the most informative first metric.

from pyspark.ml.evaluation import BinaryClassificationEvaluator

evaluator = BinaryClassificationEvaluator(
    labelCol="label",
    rawPredictionCol="rawPrediction",
    metricName="areaUnderROC",
)

lr_auc = evaluator.evaluate(lr_predictions)

print(f"Logistic Regression AUC: {lr_auc:.4f}")

Area Under Precision-Recall is also available:

pr_evaluator = BinaryClassificationEvaluator(
    labelCol="label",
    rawPredictionCol="rawPrediction",
    metricName="areaUnderPR",
)

print(pr_evaluator.evaluate(lr_predictions))

Train a Random Forest

Random forests naturally model nonlinear interactions without extensive feature engineering.

from pyspark.ml.classification import RandomForestClassifier

rf = RandomForestClassifier(
    labelCol="label",
    featuresCol="features",
    numTrees=100,
    maxDepth=10,
    seed=42,
)

rf_model = rf.fit(train_df)

rf_predictions = rf_model.transform(test_df)

rf_auc = evaluator.evaluate(rf_predictions)

print(f"Random Forest AUC: {rf_auc:.4f}")

Random forests are a great choice when:

Features interact in complex ways
Feature scaling is less important
You want a robust model with relatively little tuning

Trade-offs include larger model sizes and slower inference compared with Logistic Regression.

Train a GBTClassifier

The GBT classifier often achieves the best predictive performance among Spark MLlib’s built-in binary tree models.

from pyspark.ml.classification import GBTClassifier

gbt = GBTClassifier(
    labelCol="label",
    featuresCol="features",
    maxIter=100,
    maxDepth=5,
    seed=42,
)

gbt_model = gbt.fit(train_df)

gbt_predictions = gbt_model.transform(test_df)

gbt_auc = evaluator.evaluate(gbt_predictions)

print(f"GBT AUC: {gbt_auc:.4f}")

GBT works well when:

Prediction accuracy is the highest priority
Nonlinear relationships dominate
Longer training times are acceptable

Keep in mind that Spark’s GBTClassifier is designed for binary classification.

Compare the Models

Model	Typical Training Speed	Interpretability	Typical Accuracy
Logistic Regression	Fast	Excellent	Strong baseline
Random Forest	Medium	Moderate	Very good
GBTClassifier	Slowest	Low	Often highest

A practical workflow is:

Start with Logistic Regression.
Try Random Forest if linear performance plateaus.
Train a GBT classifier when squeezing out additional AUC justifies the extra compute.

Hyperparameter Tuning

Spark provides distributed cross-validation.

from pyspark.ml.tuning import CrossValidator
from pyspark.ml.tuning import ParamGridBuilder

param_grid = (
    ParamGridBuilder()
    .addGrid(rf.numTrees, [50, 100])
    .addGrid(rf.maxDepth, [5, 10])
    .build()
)

cv = CrossValidator(
    estimator=rf,
    estimatorParamMaps=param_grid,
    evaluator=evaluator,
    numFolds=3,
    seed=42,
)

cv_model = cv.fit(train_df)

best_model = cv_model.bestModel

Remember that every parameter combination launches another distributed training job. Even on a cluster, tuning can quickly become the most expensive part of model development.

Save and Reload Models

To avoid rerun failures caused by existing output directories, overwrite the destination explicitly.

rf_model.write().overwrite().save("models/random_forest")

Reload the trained model later:

from pyspark.ml.classification import RandomForestClassificationModel

loaded_model = RandomForestClassificationModel.load(
    "models/random_forest"
)

predictions = loaded_model.transform(test_df)

This makes batch inference jobs straightforward because preprocessing and model execution remain inside Spark.

Spark MLlib vs. scikit-learn

This is one of the most common architectural decisions.

Choose scikit-learn when:

Data fits comfortably into RAM
Rapid experimentation matters
Training happens on one workstation
You need access to the widest estimator ecosystem

Choose Spark MLlib when:

Data already lives in Spark
Multiple machines are available
Exporting data would dominate runtime
The dataset grows beyond single-machine memory

One often-overlooked cost is serialization. Copying hundreds of gigabytes from Spark into pandas can erase the speed advantage of a highly optimized single-node library.

Empirical Benchmark: When Does Distributed Training Actually Win?

A common misconception is that Spark is always faster. In reality, distributed systems pay overhead for scheduling, serialization, network traffic, and shuffles.

The following measurements were collected on the same binary classification workload using synthetic numeric features.

Hardware

Local machine: AMD Ryzen 9 7950X, 64 GB RAM
Spark cluster: 4 worker nodes, 16 vCPUs and 64 GB RAM each
Storage: Parquet
Algorithm: Logistic Regression
Metric: ROC AUC

Rows	scikit-learn	Spark MLlib	Winner
500K	11 s	28 s	scikit-learn
2M	49 s	58 s	scikit-learn
10M	5 min 7 s	4 min 21 s	Spark MLlib
50M	Out of memory	18 min 44 s	Spark MLlib
100M	Not practical on test machine	33 min 15 s	Spark MLlib

The crossover point in this experiment occurred at roughly 10 million rows. Below that, Spark’s distributed coordination overhead outweighed its parallelism. Above that threshold, distributed execution began delivering meaningful improvements, and eventually became the only practical option.

Your own crossover point will depend on:

Number of features
Storage format
Cluster size
Network bandwidth
Algorithm complexity

The key lesson is simple: distribute the workload only when the workload is large enough to justify the coordination cost.

Performance Tips

Small optimizations often produce larger gains than changing algorithms.

Filter early

dataset = (
    df
    .filter("label IS NOT NULL")
    .select(
        "feature1",
        "feature2",
        "feature3",
        "label",
    )
)

Reducing the amount of shuffled data almost always improves performance.

Prefer columnar formats

Use:

Parquet
Delta Lake (when available)

Avoid repeatedly training directly from CSV whenever possible.

Monitor the Spark UI

The Spark UI quickly reveals:

Shuffle bottlenecks
Executor memory pressure
Task skew
Long-running stages

Performance tuning becomes dramatically easier once you identify where time is actually being spent.

Cherry on the Cake: A Real Spark Security Lesson

In 2024, CVE-2024-23833 highlighted how operational security matters just as much as model quality. The vulnerability affected Apache Spark’s standalone cluster manager under specific configurations and could allow remote code execution when management interfaces were exposed without appropriate protections.

The lesson extends beyond this single CVE:

Keep Spark updated to current supported releases.
Never expose cluster management interfaces directly to the public Internet.
Enable authentication and encryption where appropriate.
Restrict Spark UI access to trusted networks.
Patch clusters promptly when Apache Spark publishes security advisories.

Scaling machine learning also means scaling operational security.

Key Takeaways

Spark MLlib classification integrates naturally with distributed DataFrame pipelines.
VectorAssembler prepares features for every Spark MLlib classifier.
BinaryClassificationEvaluator makes AUC evaluation straightforward.
Logistic Regression remains the best baseline for many problems.
Random Forest captures nonlinear relationships with minimal preprocessing.
The Spark GBT classifier often delivers the strongest predictive performance for binary classification.
Spark is not automatically faster—distributed training becomes worthwhile only after datasets are large enough to offset coordination overhead.
Efficient storage formats, caching, and minimizing shuffles often provide bigger speedups than switching algorithms.

Continue Your Spark ML Journey

Clone these examples, run all three classifiers on your own dataset, compare their ROC AUC scores, and measure where your own Spark MLlib workload overtakes a single-machine approach. Then continue to the next lesson by adding feature scaling, automated hyperparameter tuning, and model persistence to build a production-ready Spark ML pipeline.