Spark ML feature pipelines: VectorAssembler, StringIndexer, OneHotEncoder, and the Pipeline API

LEARN · DISTRIBUTED ML WITH PYSPARK

Machine learning models rarely consume raw tables directly. Before training, data usually needs cleaning, categorical encoding, feature assembly, and a consistent transformation path for future predictions.

Apache Spark’s spark.ml library provides the Pipeline API to package these steps into a reusable workflow. This approach is central to practical SparkML projects because it keeps preprocessing and modeling together, reduces duplicated code, and makes FeatureEngineering easier to test and deploy.

This lesson focuses on four core components:

StringIndexer for converting categories into numeric indexes
OneHotEncoder for creating model-friendly categorical vectors
VectorAssembler for combining features into a single vector column
Pipeline for chaining transformations and models

The key production rule throughout this workflow is simple:

Only learn preprocessing parameters from training data.

The Spark ML pipeline mental model

A Spark ML Pipeline is a sequence of stages. Each stage is either:

A Transformer: applies existing rules to a DataFrame.
An Estimator: learns from data and creates a Transformer.

The typical machine learning lifecycle is:

Split raw data into training and test sets.
Fit the pipeline only on training data.
Transform training and test data using the fitted pipeline.
Evaluate the model.
Save the fitted pipeline for future predictions.

The important boundary is the fit() operation.

Anything that learns from data, such as category mappings or model weights, belongs inside the training phase.

Transformers: applying learned or fixed logic

A Transformer takes a DataFrame and returns a new DataFrame.

Its main operation is:

transformed_df = transformer.transform(input_df)

Transformers do not learn new information from the input data. They already contain the logic required to modify columns.

Common Spark ML Transformers include:

VectorAssembler
OneHotEncoderModel
A fitted StringIndexerModel
A trained classification model

For example, VectorAssembler combines multiple feature columns into the vector format expected by most Spark ML algorithms.

from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(
    inputCols=["age", "income"],
    outputCol="features"
)

result_df = assembler.transform(df)

result_df.select("features").show()

A resulting vector may look like:

[25.0,50000.0]

Most Spark ML algorithms expect this single features column rather than many separate input columns.

Estimators: learning from data

An Estimator learns parameters from a DataFrame.

Its main operation is:

model = estimator.fit(training_df)

The result is a Transformer.

Examples of Estimators:

StringIndexer
LogisticRegression
RandomForestClassifier
DecisionTreeClassifier

A StringIndexer learns how string categories map to numeric values.

For example, training data might produce:

country	countryIndex
Germany	0
France	1
Spain	2

The exact ordering is learned during fitting. It should not be manually recreated in another script.

from pyspark.ml.feature import StringIndexer

country_indexer = StringIndexer(
    inputCol="country",
    outputCol="countryIndex"
)

country_model = country_indexer.fit(training_df)

indexed_df = country_model.transform(training_df)

After fitting, country_model remembers the category mapping and can apply it consistently to new data.

StringIndexer and OneHotEncoder together

Many machine learning algorithms cannot use raw string values. A common categorical feature workflow is:

String category
      |
      v
StringIndexer
      |
      v
Numeric category index
      |
      v
OneHotEncoder
      |
      v
Sparse feature vector

Example:

from pyspark.ml.feature import OneHotEncoder, StringIndexer

country_indexer = StringIndexer(
    inputCol="country",
    outputCol="countryIndex",
    handleInvalid="keep"
)

country_encoder = OneHotEncoder(
    inputCols=["countryIndex"],
    outputCols=["countryVector"]
)

The handleInvalid="keep" option is useful when production data contains categories that were not present during training.

For example:

Training data:

Germany
France
Spain

Production data:

Italy

Instead of failing because the category is unknown, Spark can place the unexpected value into a separate handling bucket.

VectorAssembler: creating the final feature vector

After preprocessing, Spark ML models usually expect a single column named features.

A dataset might start as:

age	income	country
32	72000	Germany

After transformation, the model receives something like:

features = [32,72000,1,0]

VectorAssembler combines numeric columns and encoded vectors.

from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(
    inputCols=[
        "age",
        "income",
        "countryVector"
    ],
    outputCol="features"
)

assembled_df = assembler.transform(encoded_df)

The resulting features column becomes the input for the machine learning algorithm.

Building a complete Spark ML Pipeline

The Pipeline API connects preprocessing and modeling into one object.

The following example:

creates sample data
splits training and test datasets
indexes a categorical feature
one-hot encodes it
assembles features
trains a logistic regression model
generates predictions

from pyspark.sql import SparkSession

from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import (
    StringIndexer,
    OneHotEncoder,
    VectorAssembler
)

spark = SparkSession.builder \
    .appName("SparkMLPipelineExample") \
    .getOrCreate()

data = [
    (25, 50000, "Germany", 0),
    (35, 70000, "France", 1),
    (45, 90000, "Germany", 1),
    (23, 40000, "Spain", 0),
    (52, 110000, "France", 1)
]

columns = [
    "age",
    "income",
    "country",
    "label"
]

df = spark.createDataFrame(data, columns)

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

country_indexer = StringIndexer(
    inputCol="country",
    outputCol="countryIndex",
    handleInvalid="keep"
)

country_encoder = OneHotEncoder(
    inputCols=["countryIndex"],
    outputCols=["countryVector"]
)

assembler = VectorAssembler(
    inputCols=[
        "age",
        "income",
        "countryVector"
    ],
    outputCol="features"
)

classifier = LogisticRegression(
    featuresCol="features",
    labelCol="label"
)

pipeline = Pipeline(
    stages=[
        country_indexer,
        country_encoder,
        assembler,
        classifier
    ]
)

pipeline_model = pipeline.fit(train_df)

predictions = pipeline_model.transform(test_df)

predictions.select(
    "country",
    "features",
    "prediction",
    "probability"
).show()

The important line is:

pipeline_model = pipeline.fit(train_df)

During fitting, Spark learns:

category mappings from StringIndexer
encoding metadata
logistic regression parameters

Then:

predictions = pipeline_model.transform(test_df)

applies the learned transformations to unseen data.

Avoiding feature leakage

A common machine learning mistake is fitting preprocessing on the full dataset before splitting.

Risky example:

from pyspark.ml.feature import StringIndexer

indexer = StringIndexer(
    inputCol="country",
    outputCol="countryIndex"
)

full_model = indexer.fit(all_data)

train_data = full_model.transform(train_data)
test_data = full_model.transform(test_data)

The problem is that information from the test dataset influenced the learned category mapping.

For a small example this may appear harmless, but in real systems learned preprocessing can affect evaluation results and make metrics less trustworthy.

The safer approach is:

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

pipeline_model = pipeline.fit(train_df)

train_predictions = pipeline_model.transform(train_df)

test_predictions = pipeline_model.transform(test_df)

The correct order is:

Split.
Fit.
Transform.

Preventing training-serving skew

Without pipelines, teams often duplicate preprocessing logic:

SQL transformations for training
separate scripts for batch scoring
custom inference code in applications

Small differences can change predictions.

Common problems include:

different category mappings
different missing-value handling
incorrect feature ordering
outdated preprocessing code in production

A saved PipelineModel keeps feature preparation and prediction together.

Save a trained pipeline:

pipeline_model.write() \
    .overwrite() \
    .save("models/customer_pipeline")

Load it later:

from pyspark.ml import PipelineModel

loaded_model = PipelineModel.load(
    "models/customer_pipeline"
)

new_predictions = loaded_model.transform(new_data)

The same transformation rules can now be reused for future inference.

Cherry on the cake: a real Spark security lesson

Feature leakage is a machine learning reliability problem, but production ML systems also depend on secure infrastructure.

A notable Apache Spark example is CVE-2022-33891. This vulnerability affected certain Spark versions where the Spark UI access control mechanism could be bypassed under specific configurations, potentially allowing unauthorized command execution as the Spark process user.

The issue was fixed through Spark security releases, including updates in the affected Spark 3.2.x and 3.3.x release lines. Organizations running affected versions were advised to upgrade to patched versions and review Spark UI access controls.

Why does this matter for ML pipelines?

A production model is not only a file containing weights. It depends on:

training jobs
feature pipelines
stored artifacts
cluster configuration
permissions
dependency management

Good FeatureEngineering practices and good security practices follow the same principle: make the workflow reproducible, controlled, and observable.

Spark ML Pipeline checklist

Before deploying a Pipeline-based workflow, verify:

Training and test data are split before fitting learned transformations.
StringIndexer handles unexpected categories correctly.
Categorical encoding is identical during training and inference.
VectorAssembler receives the expected columns.
The fitted PipelineModel is saved and versioned.
Test data does not influence preprocessing parameters.
Spark dependencies and security updates are reviewed regularly.

Summary

Spark ML pipelines provide a structured way to build reliable machine learning workflows.

The core ideas are:

Transformers apply existing transformations.
Estimators learn parameters through fit().
StringIndexer converts categories into numeric indexes.
OneHotEncoder converts categorical indexes into feature vectors.
VectorAssembler creates the final model input.
Pipeline combines preprocessing and modeling into one reusable object.
Learned transformations must only use training data.

The next step is practical: take one of your existing SparkML notebooks, replace manual FeatureEngineering steps with a Pipeline, save the resulting model, and test it on new data.

Try building your own production-ready Spark ML Pipeline today, and share your results with your team.