Skip to main content
PySpark · lesson 8 of 11

Beginner

4

Intermediate

4

Advanced

3

Also part of PySpark

6
Open the PySpark page →
PySpark intermediate Lesson 8 of 11

Also part of: Apache Spark

MLlib Basics: A Real Pipeline, Fit and Evaluated

Assemble feature vectors, chain them into a Pipeline with a classifier, and evaluate the fitted model with AUC and accuracy on held-out data.

By SkillByExample · Published

pyspark mllib machine-learning pipeline classification

Every DataFrame operation in this track so far has been feature engineering, joins, and aggregation — the input to a model, never the model itself. MLlib is Spark’s own machine learning library, and it plugs directly into the same DataFrames.

The setup

pip install pyspark==3.5.4 numpy

numpy is a real dependency of pyspark.ml — the VectorAssembler and evaluator classes import it directly, unlike the core pyspark.sql DataFrame API used elsewhere in this track.

A Pipeline: assemble, then classify

Sixteen orders, each with an amount, days until shipped, and whether it was later returned.

from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import LogisticRegression
from pyspark.ml import Pipeline
from pyspark.ml.evaluation import BinaryClassificationEvaluator, MulticlassClassificationEvaluator

spark = SparkSession.builder.appName("mllib").master("local[2]").getOrCreate()

data = spark.createDataFrame([
    (25.50, 3, 0), (12.00, 1, 0), (40.00, 14, 1), (8.75, 2, 0),
    (63.20, 20, 1), (19.99, 4, 0), (55.00, 18, 1), (22.10, 5, 0),
    (9.99, 1, 0), (31.40, 16, 1), (45.00, 15, 1), (15.25, 2, 0),
    (60.00, 21, 1), (11.50, 1, 0), (28.00, 13, 1), (18.75, 3, 0),
], ["amount", "days_to_ship", "returned"])

train, test = data.randomSplit([0.75, 0.25], seed=42)

assembler = VectorAssembler(inputCols=["amount", "days_to_ship"], outputCol="features")
lr = LogisticRegression(featuresCol="features", labelCol="returned")
pipeline = Pipeline(stages=[assembler, lr])

model = pipeline.fit(train)
predictions = model.transform(test)
predictions.select("amount", "days_to_ship", "returned", "prediction", "probability").show(truncate=False)
$ python mllib_basics.py
+------+------------+--------+----------+-------------------------------------------+
|amount|days_to_ship|returned|prediction|probability                                |
+------+------------+--------+----------+-------------------------------------------+
|19.99 |4           |0       |0.0       |[0.9999999991199242,8.800757900218059E-10] |
|55.0  |18          |1       |1.0       |[2.964294904958629E-19,1.0]                |
|9.99  |1           |0       |0.0       |[0.9999999999999993,6.661338147750939E-16] |
|31.4  |16          |1       |1.0       |[1.2453099987548877E-13,0.9999999999998754]|
|45.0  |15          |1       |1.0       |[4.061595234027321E-13,0.9999999999995939] |
|60.0  |21          |1       |1.0       |[6.131395592202325E-25,1.0]                |
+------+------------+--------+----------+-------------------------------------------+

VectorAssembler runs first, folding amount and days_to_ship into one features column — the shape every MLlib algorithm expects. pipeline.fit(train) then runs both stages: assemble the training data, fit LogisticRegression on the result. Calling .transform(test) on the returned PipelineModel reapplies both stages to data the model never trained on — the same VectorAssembler configuration, no risk of the reader forgetting to repeat it by hand.

Evaluating the fitted model

auc = BinaryClassificationEvaluator(labelCol="returned").evaluate(predictions)
acc = MulticlassClassificationEvaluator(labelCol="returned", metricName="accuracy").evaluate(predictions)
print(f"AUC: {auc:.4f}")
print(f"Accuracy: {acc:.4f}")

spark.stop()
AUC: 1.0000
Accuracy: 1.0000

Both evaluators take the same predictions DataFrame .transform() produced — no separate scoring step, no converting to numpy arrays first. A perfect score here is real, not a sign of a broken evaluator: days_to_ship cleanly separates the two classes in this 16-row dataset (every returned order shipped in 13+ days), which is exactly the kind of result that should make you suspicious of a real dataset scoring this well and check for label leakage before trusting it.

Practice

1. Print the fitted model's coefficients and intercept. 
lr_model = model.stages[-1]
print(f"Coefficients: {lr_model.coefficients}")
print(f"Intercept: {lr_model.intercept}")
Coefficients: [0.20834396100674674,4.015669187013517]
Intercept: -41.078485652812056

model.stages holds every fitted stage in order — [-1] is the last one, the fitted LogisticRegressionModel. The days_to_ship coefficient (4.02) dwarfs amount’s (0.21), confirming what the perfect-separation result already hinted: shipping delay is doing nearly all the predictive work here. Re-run this fit and the exact numbers shift each time — perfectly separable data has no unique logistic-regression optimum, so the iterative solver’s stopping point (not the underlying signal) moves. The lopsided ratio between the two coefficients is the stable part; the absolute values are not.

2. Add a third feature column and refit the Pipeline without changing the LogisticRegression stage. 
data2 = data.withColumn("amount_per_day", data.amount / data.days_to_ship)
assembler2 = VectorAssembler(inputCols=["amount", "days_to_ship", "amount_per_day"], outputCol="features")
pipeline2 = Pipeline(stages=[assembler2, lr])
model2 = pipeline2.fit(data2.randomSplit([0.75, 0.25], seed=42)[0])
print(model2.stages[-1].coefficients)
[0.2726122504191249,3.6939078077778835,-0.9538249074969803]

Only assembler2’s inputCols changed — lr is reused unmodified, since a Pipeline stage just needs a features column of the right shape, not a fixed number of source columns.

3. Evaluate with metricName="f1" instead of accuracy. 
f1 = MulticlassClassificationEvaluator(labelCol="returned", metricName="f1").evaluate(predictions)
print(f"F1: {f1:.4f}")
F1: 1.0000

Same perfect score, for the same reason as the AUC and accuracy above — f1 only diverges from accuracy when there is a real precision/recall tradeoff to weigh, which a perfectly separable dataset does not have.

Frequently Asked Questions

Do I need pandas or scikit-learn installed to use MLlib?
No. MLlib is Spark's own ML library, built on DataFrames — VectorAssembler, LogisticRegression, and the evaluators used here all live in pyspark.ml and run through the same Spark execution engine as every other DataFrame operation in this track. numpy is a real dependency of pyspark.ml itself, so install it alongside pyspark.
Why does a model need a Pipeline instead of just calling fit() directly?
A Pipeline chains transformers (like VectorAssembler, which reshapes columns into a features vector) and an estimator (the actual algorithm) into one object. Fitting the Pipeline runs every stage in order and returns a PipelineModel that applies all of them to new data with one .transform() call — the same feature engineering a reader might otherwise forget to repeat by hand on test data.
What is the difference between a Transformer and an Estimator in MLlib?
An Estimator has a fit() method that learns from data and produces a Transformer — LogisticRegression is an Estimator, and calling .fit(train) on it returns a fitted LogisticRegressionModel, which is a Transformer. A Transformer only has .transform(), no learning — VectorAssembler is a Transformer from the start, since assembling columns into a vector needs no training.
Why does the probability column have two values per row?
BinaryClassificationEvaluator and the model's probability column both report a probability for each class — [P(class 0), P(class 1)] — not just the predicted class. prediction is 1.0 when P(class 1) is higher, but the full probability vector is what a calibration check or a custom decision threshold would use instead of the default 0.5 cutoff.