Skip to main content

Module 7 — Available algorithms and their constraints

MLlib is not scikit-learn scaled up: it is a deliberately smaller library that ships the algorithms Spark's engineers were confident they could implement in a distributed way without silent quality loss. Knowing what is in the box, and — more importantly — what is not, is what stops a project halfway when a stakeholder asks for "the same model our data science team uses on their laptops".

Regression and classification

Linear models are complete and well tuned. LinearRegression, LogisticRegression, and their multinomial variants live in pyspark.ml.regression and pyspark.ml.classification. Both support L1 and L2 penalties through the elasticNetParam and regParam parameters, and both scale linearly with the number of executors.

Tree-based methods dominate the rest of the library. DecisionTreeClassifier, RandomForestClassifier and GBTClassifier — the gradient-boosted tree — are the workhorses of tabular models on Spark. The gradient-boosted tree in particular is well suited to the flight-delays project: it handles mixed feature types, ignores the scale of numeric features, and produces well-calibrated probabilities out of the box.

from pyspark.ml.classification import GBTClassifier

gbt = GBTClassifier(
labelCol="late", featuresCol="features",
maxIter=100, maxDepth=6, stepSize=0.1, seed=42,
)

Two limits to know about MLlib trees. First, they only handle binary classification with GBTClassifier; multiclass gradient boosting exists in the ecosystem but not inside MLlib itself. Second, MLlib's implementation is honest but not fast compared with modern libraries: on the same data, a well-tuned XGBoost typically trains in half the time and reaches slightly higher AUC.

Clustering and recommendation

KMeans and BisectingKMeans cover clustering. Both scale well because the k-means iteration is naturally parallel — assign points to centroids, then average — and both accept the same features vector column as every other MLlib algorithm.

ALS in pyspark.ml.recommendation deserves a paragraph on its own. Alternating Least Squares is the algorithm behind matrix factorisation for recommendation, and it is the one place where MLlib is objectively the best choice available: it was co-designed with Spark's shuffle model, it handles hundreds of millions of user-item interactions, and no single-node library gets close on that regime. If the problem is "users and items", start here.

What MLlib does not contain

The absences are as informative as the presences. Multiclass gradient boosting is missing. Neural networks beyond a very basic MultilayerPerceptronClassifier are missing — Spark is not a deep-learning framework. SVMs with non-linear kernels are missing; only linear SVM is available. DBSCAN, HDBSCAN and other density-based clustering methods are missing. UMAP and t-SNE are missing. Isolation Forest for anomaly detection is missing.

For each of these, the honest answer is the same: train on a sample small enough to fit on one machine, and use scikit-learn, XGBoost, PyTorch or whatever is standard in the ecosystem. Spark handles the sampling and the scoring; the actual training happens where the specialised library lives.

# Sample down to 5 million rows, bring them local, train on one machine
pdf = flights.sample(0.05, seed=42).toPandas()
from xgboost import XGBClassifier
model = XGBClassifier(n_estimators=500, max_depth=6).fit(pdf[FEATURES], pdf["late"])

Distributed XGBoost — the middle path

XGBoost provides a Spark integration (xgboost.spark) that plugs directly into the pipeline API. It is not part of MLlib, so you install it separately, but the ergonomics are identical:

from xgboost.spark import SparkXGBClassifier

xgb = SparkXGBClassifier(
label_col="late", features_col="features",
num_workers=8, use_gpu=False,
)
pipeline = Pipeline(stages=[..., xgb])

num_workers matches the number of executors. On the flight-delays project this typically doubles the training speed of MLlib's GBTClassifier at equal AUC and gives access to XGBoost's better regularization options. It is the pragmatic default for tabular models on Spark in 2026.

The choice, made honestly

The right question is never "which algorithm does MLlib have?" — it is "which algorithm does this problem need, and where should it train?".

Data sizeModel neededWhere to train
Fits in RAM (a few GB)Anythingscikit-learn, XGBoost, PyTorch on one machine
Larger than RAM, fits on diskLinear, tree ensemblesSample and train locally, or MLlib / SparkXGBoost
Truly distributed (100 GB+)Linear, trees, ALS, k-meansMLlib or SparkXGBoost, no other option scales
Truly distributedNeural networks, non-linear SVM, DBSCANSample down or use a specialised distributed framework
An algorithm missing from MLlib is not a Spark bug

Reaching for a random forest that supports multiclass gradient boosting because "MLlib should have it" is a losing move. MLlib is deliberately minimal. Combine it with the rest of the Python ecosystem — sample for training, Spark for scoring — and you get both scale and algorithm choice.

Summary

  • MLlib covers linear models, trees, gradient-boosted trees, k-means and ALS well; ALS in particular is best-in-class.
  • Multiclass gradient boosting, non-linear SVM, deep networks and density-based clustering are missing.
  • When an algorithm is missing, sample down and train on one machine, then use Spark to score at scale.
  • SparkXGBoost fills the tabular gap: same pipeline API, faster training, better regularization than MLlib's GBT.

Next module: making the whole apparatus fast — partitions, memory, and the shuffle you cannot avoid.