Spark ML Regression, in depth: solvers, GLMs, boosted trees, a worked delivery-time model and the failures that matter

The regressor catalog

All of them accept weightCol, except where the documentation says otherwise, so you can reweight rows without duplicating them. Every fitted model is a Transformer that slots into a PipelineModel.

How training runs on a cluster

LinearRegression has two solvers, picked by solver (auto, normal or l-bfgs). The normal solver makes one pass over the data to build the d-by-d Gram matrix and the X-transpose-y vector, then solves the small system on the driver. It is fast and exact, but memory grows with d squared, so Spark restricts it to at most 4,096 features. The l-bfgs solver is iterative, as in the diagram. Each iteration broadcasts the current coefficients, each partition computes its loss and gradient, and treeAggregate sums the results. With an L1 component (elasticNetParam > 0) Spark uses the OWL-QN variant. auto chooses normal when it applies and L-BFGS otherwise.

Two consequences follow. First, standardization=True is the default, and Spark standardises features internally while reporting coefficients on the original scale. A StandardScaler in front of LinearRegression is therefore redundant. It still matters for models that lack this option. Second, an iterative fit over uncached input re-reads and re-parses the source on every iteration. Cache the assembled training set before fit() and unpersist it afterwards.

Tree ensembles train differently. Spark bins each continuous feature into at most maxBins candidate thresholds (default 32). Each tree level is one distributed pass that collects split statistics for the nodes being grown. Random forests grow trees in parallel within those passes. GBT grows trees one after another, so 200 boosting rounds at depth 5 means roughly a thousand passes. That is why GBT is the slowest regressor here and the one most sensitive to caching and checkpointInterval.

Worked example: predicting delivery minutes

The task is to predict delivery minutes for a food-delivery service from distance, hour of day, weekday, city, restaurant prep-time history and a rain flag. The table has 80 million rows, and the label is right-skewed: most deliveries take 25-40 minutes, with a long tail past 90.

from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler
from pyspark.ml.regression import LinearRegression, GBTRegressor
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.sql import functions as F

df = (spark.table("deliveries")
      .where("delivered_at IS NOT NULL AND minutes BETWEEN 1 AND 240")
      .withColumn("label", F.col("minutes").cast("double")))

# Split by time, not randomly: the model will be used on future orders.
train = df.where("order_date <  '2026-09-01'")
test  = df.where("order_date >= '2026-09-01'")

cats = ["city", "weekday"]
idx = StringIndexer(inputCols=cats, outputCols=[c + "_i" for c in cats], handleInvalid="keep")
ohe = OneHotEncoder(inputCols=[c + "_i" for c in cats], outputCols=[c + "_v" for c in cats])
vec = VectorAssembler(
    inputCols=["distance_km", "hour", "rain", "prep_p50_min", "city_v", "weekday_v"],
    outputCol="features", handleInvalid="error")

lr = LinearRegression(regParam=0.01, elasticNetParam=0.0, maxIter=100)
base = Pipeline(stages=[idx, ohe, vec, lr]).fit(train.cache())

ev = RegressionEvaluator(labelCol="label", predictionCol="prediction")
pred = base.transform(test)
print("linear RMSE", ev.evaluate(pred, {ev.metricName: "rmse"}),
      "MAE", ev.evaluate(pred, {ev.metricName: "mae"}))

The linear baseline treats hour as a straight line, so it misses the lunch and dinner peaks. You could one-hot encode hour, but the interaction between rain and distance is still missing. A boosted-tree model learns those shapes without hand-built features. Trees do not need one-hot encoding, so feed the indexed columns directly and tell the assembler nothing more:

vec_t = VectorAssembler(
    inputCols=["distance_km", "hour", "rain", "prep_p50_min", "city_i", "weekday_i"],
    outputCol="features")
gbt = GBTRegressor(maxDepth=6, maxIter=200, stepSize=0.1, subsamplingRate=0.8,
                   maxBins=64, lossType="squared",
                   validationIndicatorCol="is_val", validationTol=0.001)

train_v = train.withColumn("is_val", F.rand(seed=7) < 0.1)
model = Pipeline(stages=[idx, vec_t, gbt]).fit(train_v.cache())
pred = model.transform(test)
print("GBT RMSE", ev.evaluate(pred, {ev.metricName: "rmse"}))

validationIndicatorCol gives early stopping. Boosting stops once validation error stops improving by more than validationTol, which saves both time and overfitting. maxBins=64 must be at least the largest number of categories in any indexed column, otherwise the fit fails. Report RMSE and MAE together. RMSE punishes the 90-minute misses that generate complaints, while MAE tells operations the typical error. Compare both against a trivial baseline, such as the median minutes per city and hour. A model that cannot beat that table is not worth deploying.

Reading the training summary honestly

A fitted LinearRegressionModel exposes coefficients, intercept and a summary computed on the training data: rootMeanSquaredError, r2, meanAbsoluteError, residuals and objectiveHistory. The documentation states that coefficientStandardErrors, tValues and pValues are only available with the normal solver. Even then, they are textbook inference only for an unregularised fit with roughly independent, constant-variance errors. Delivery data violates both conditions, because orders from the same restaurant are correlated.

Read the summary for three things only. Does objectiveHistory flatten out, or did you hit maxIter while still descending? Is training R-squared far above test R-squared, which means the model is overfitting? And do residuals plotted against each feature, on a sample, show structure the model is missing? Any decision about whether the model is good comes from held-out data, never from the training summary.

GLMs for counts, skew and zeros

Least squares assumes errors of constant variance around the mean. Many business targets break that assumption. Counts (orders per store per hour) have variance that grows with the mean, positive amounts (claim size, spend) are skewed and never negative, and some targets are mostly zeros with a continuous positive tail. GeneralizedLinearRegression fits these with the right likelihood:

from pyspark.ml.regression import GeneralizedLinearRegression

# Orders per store-hour: Poisson with log link; exposure enters as an offset.
glr = GeneralizedLinearRegression(family="poisson", link="log", maxIter=50,
                                  regParam=0.0, offsetCol="log_open_minutes")
m = glr.fit(store_hours)
s = m.summary
print(s.aic, s.deviance, s.residualDegreeOfFreedom)

# Spend with many zeros and a skewed positive tail: Tweedie, 1 < p < 2.
tw = GeneralizedLinearRegression(family="tweedie", variancePower=1.5, linkPower=0.0)

GLM fits use iteratively reweighted least squares and, like the normal solver, are limited to 4,096 features. The common alternative is to log-transform the label and fit ordinary least squares. It is fine for ranking, but exp(prediction) estimates the median, not the mean, which under-forecasts totals. If you sum predictions into a budget, prefer a log-link GLM. For heavy outliers in an otherwise linear problem, LinearRegression(loss="huber", epsilon=1.35) down-weights large residuals. The documentation restricts Huber to L2 regularisation, so check the constraints for your Spark version before you combine it with L1.

Evaluation and tuning

Tune with CrossValidator or the cheaper TrainValidationSplit over a ParamGridBuilder grid, using a RegressionEvaluator whose metricName is one of rmse, mse, mae, r2 or var. Each grid point times each fold is a full fit, so a 3-by-3 grid with 3 folds over GBT is 27 boosted models. Set parallelism to fit several at once, and tune on a sample before you scale up. Random K-fold splitting leaks the future into the past for time-dependent targets. In that case, use a time-based validation split built by hand. Spark ML Cross-Validation covers the mechanics, and Spark ML Feature Engineering covers the encoders used above.

Scoring, persistence and uncertainty

Save the whole fitted pipeline, not just the regressor, so that scoring applies exactly the indexers and assembler that training used:

model.write().overwrite().save("gs://ml-models/delivery_minutes/2026-10-04")

from pyspark.ml import PipelineModel
scorer = PipelineModel.load("gs://ml-models/delivery_minutes/2026-10-04")
(scorer.transform(spark.table("open_orders"))
       .select("order_id", F.round("prediction", 1).alias("eta_min"))
       .write.mode("overwrite").saveAsTable("eta_predictions"))

Batch and Structured Streaming scoring reuse the same transform() call. Low-latency per-request serving is not what Spark is for, so export to another runtime or precompute predictions. Spark's regressors return a point prediction only. There is no built-in prediction interval and no quantile loss for GBT, with one exception: AFTSurvivalRegression can output quantiles through quantileProbabilities. If the product needs "arrives within 35 to 50 minutes", estimate intervals from held-out residuals grouped by city and hour, and validate that the stated coverage holds. Version model paths by date, and record the training data range and metrics next to each one, so a bad release can be rolled back by changing a path. Spark ML Model Persistence covers compatibility across Spark versions.

Failure modes

• NaN or null features. VectorAssembler errors on nulls by default. Setting handleInvalid="skip" silently drops rows and changes the population. Impute deliberately with Imputer and count what you changed.
• Unseen categories at scoring time. StringIndexer without handleInvalid="keep" fails on a new city. With keep, the city gets an extra bucket that the model learned nothing about, so monitor its share of traffic.
• Leakage. Features computed after the event, such as courier_wait_min recorded at delivery, give a superb RMSE offline and fail in production. Audit every feature for the time it becomes available.
• Slow or stuck fits. Uncached input, thousands of tiny partitions, or GBT lineage growing without checkpointInterval and a checkpoint directory.
• Driver out of memory. The normal solver near 4,096 features, or very wide one-hot vectors, aggregated on a small driver.
• Collinearity. Unregularised linear fits on correlated features give unstable, sign-flipping coefficients. Add a little regParam if coefficients are used for explanation.
• Drift. A model trained before a pricing change keeps predicting the old world. Track residuals by week and retrain on a schedule.

Trade-offs

Linear and GLM models train in minutes, have readable coefficients, extrapolate (sometimes dangerously), and need you to build non-linear features by hand. Random forests are hard to break, parallelise well, and give slightly worse accuracy than boosting. GBT usually wins on accuracy for tabular data, but trains sequentially and does not extrapolate beyond the label range it saw. When the data fits on one large machine after aggregation, single-node XGBoost or LightGBM often trains faster and offers more loss functions. In that case, use Spark for feature preparation and hand off a sampled Parquet file. Choose Spark ML when the training set truly needs the cluster, or when you want one PipelineModel that does both feature preparation and scoring inside existing Spark jobs.

What to do next

1. Build the trivial baseline (a grouped median) and record its RMSE and MAE on a time-based test split.
2. Fit LinearRegression in a pipeline with indexers and an assembler, on cached input, and read objectiveHistory to confirm convergence.
3. Check the label distribution: counts or positive skewed values point to GeneralizedLinearRegression with a log link.
4. Fit GBTRegressor with validationIndicatorCol for early stopping, maxBins above your largest category count, and a checkpoint directory.
5. Audit every feature for leakage by asking when it becomes known.
6. Tune a small grid on a sample, then refit the winner on full data and save the PipelineModel.
7. Schedule weekly residual monitoring and a retraining trigger.