Skip to main content

Module 5 — Spark ML: transformers and estimators

MLlib, the machine-learning library that ships with Spark, is smaller than scikit-learn and organised around a stricter idea: everything is either a Transformer or an Estimator. Learn the two contracts and every algorithm in the library uses them the same way. In this module we turn the flight-delays DataFrame of module 4 into the numerical feature vector the next module will train a model on.

The two contracts​

A Transformer has one method that matters: transform(df) -> df. It takes a DataFrame in, returns a DataFrame out, with a new or altered column. A StringIndexer already fitted on the training data is a Transformer; so is a VectorAssembler.

An Estimator has one method: fit(df) -> Transformer. It learns something from the data — the vocabulary of a categorical column, the mean and standard deviation of a numeric column, the parameters of a model — and returns a Transformer that carries those learned quantities. Every trainable object in MLlib is an Estimator.

These two contracts are what make the pipeline of module 6 possible: a pipeline is nothing but a list of Transformers and Estimators, and Spark knows how to fit and transform a list.

Encoding a categorical column​

The carrier column of the flights dataset is a two-letter code — DL, AA, WN, and so on. A model needs numbers, so we index the strings and then one-hot encode them.

from pyspark.ml.feature import StringIndexer, OneHotEncoder

carrier_idx = StringIndexer(
    inputCol="carrier", outputCol="carrier_idx",
    handleInvalid="keep",           # unseen codes at scoring time go to a bucket
)
carrier_oh = OneHotEncoder(inputCol="carrier_idx", outputCol="carrier_vec")

StringIndexer is an Estimator: fit produces the vocabulary, transform applies it. OneHotEncoder is also an Estimator in Spark 3+ because it needs to know how many distinct indices exist to size the sparse vector. Set handleInvalid="keep" on the indexer to route unseen values (the airline that starts flying next year) to a dedicated bucket instead of throwing at scoring time.

Scaling a numeric column​

Distance and departure delay live on very different scales. For anything gradient-based — logistic regression, the linear part of a gradient-boosted tree's leaves — they should be standardised.

from pyspark.ml.feature import StandardScaler, VectorAssembler

numeric_asm = VectorAssembler(
    inputCols=["distance", "departure_delay"],
    outputCol="numeric_vec",
)
scaler = StandardScaler(
    inputCol="numeric_vec", outputCol="numeric_vec_scaled",
    withMean=True, withStd=True,
)

Notice the small ceremony that has no equivalent in scikit-learn: StandardScaler works on a vector column, not on individual numeric columns. That is why we assemble the numeric columns into numeric_vec first, then scale the result.

VectorAssembler — the mandatory final step​

Every MLlib algorithm expects one input column of type vector. Not many. Not one per feature. One. VectorAssembler is the transformer that concatenates whatever columns you point it at into that single vector.

features_asm = VectorAssembler(
    inputCols=["numeric_vec_scaled", "carrier_vec"],
    outputCol="features",
)

Two rules govern this final assembly and they trip up every newcomer.

First, all input columns must be numeric or already a vector. A raw string column will make VectorAssembler throw; a raw integer works fine. That is why the flow is always StringIndexer -> OneHotEncoder -> VectorAssembler, in that order.

Second, the output column name is always features by convention. Every algorithm's featuresCol parameter defaults to "features", and every downstream tutorial assumes it. Follow the convention.

Putting the four steps together, by hand​

train = flights.filter(F.col("year") < 2024)
test  = flights.filter(F.col("year") == 2024)

indexer_model = carrier_idx.fit(train)
train_idx = indexer_model.transform(train)

oh_model = carrier_oh.fit(train_idx)
train_oh = oh_model.transform(train_idx)

train_num = numeric_asm.transform(train_oh)
scaler_model = scaler.fit(train_num)
train_scaled = scaler_model.transform(train_num)

train_ready = features_asm.transform(train_scaled)
train_ready.select("features").show(3, truncate=False)

Read that block carefully: every fit is called on the training set only, and every model produced by a fit is then used to transform both the training and the eventual test set. Do it in the other order — fit the scaler on the whole dataset before splitting — and you have leaked test information into the training pipeline. Module 6 makes this discipline automatic by wrapping the whole sequence in a Pipeline.

A vector is not a list of numbers

Spark's Vector is a distinct type, backed by a specialised binary layout that keeps sparse encodings compact — a hundred-thousand-carrier one-hot vector is stored as a handful of indices, not a hundred-thousand zeros. Turning a Vector back into a Python list with .toArray() in the middle of a Spark job defeats every optimization and shuffles Python objects between the JVM and the workers. Never do it inside a pipeline.

Summary​

  • Every feature-engineering step in MLlib is either a Transformer (transform) or an Estimator (fit -> Transformer).
  • Categorical columns go through StringIndexer (with handleInvalid="keep") then OneHotEncoder.
  • Numeric columns are assembled and scaled as a vector, because StandardScaler expects a vector input.
  • Every algorithm needs a single features column of type vector, built by VectorAssembler as the last feature step.

Next module: chaining these steps into a Pipeline and doing distributed cross-validation without leakage.