Spark ML

Spark includes a machine-learning library named MLlib. The phrase “Spark ML” is commonly used for MLlib’s DataFrame-based API because its Java and Scala classes live under org.apache.spark.ml and its Python classes under pyspark.ml.

These names are easy to confuse:

spark.mllib is not removed, but new application code should normally use spark.ml. The DataFrame API integrates with Spark SQL, has a consistent pipeline model, and provides similar APIs across Java, Scala, and Python.

Spark ML supplies distributed feature preparation, model fitting, evaluation, hyperparameter tuning, and batch inference. It is most useful when training or scoring data is already large enough to benefit from Spark. It is not a general deep-learning framework and does not automatically make every algorithm efficient merely by running it on a cluster.

What Spark ML Provides #

The library covers several categories:

• Classification: Logistic regression, decision trees, random forests, gradient-boosted trees, naive Bayes, linear support vector machines, and multilayer perceptrons
• Regression: Linear and generalized linear regression, decision trees, random forests, gradient-boosted trees, survival regression, isotonic regression, and factorization machines
• Clustering: K-means, bisecting k-means, Gaussian mixtures, latent Dirichlet allocation, and power iteration clustering
• Recommendation: Alternating least squares
• Frequent patterns: FP-growth and prefix-span
• Feature engineering: Tokenization, categorical encoding, feature assembly, scaling, hashing, TF-IDF, discretization, imputation, polynomial expansion, dimensionality reduction, and feature selection
• Statistics: Correlation, hypothesis tests, and vector summaries
• Workflow: Pipelines, evaluators, cross-validation, train-validation splitting, parameter grids, and model persistence

Algorithm support and parameters vary among Java, Scala, Python, and R. Check the API documentation for the Spark version used by the cluster rather than assuming every binding has identical coverage.

DataFrames as ML Datasets #

Spark ML consumes and produces DataFrames. A supervised-learning dataset conventionally contains:

+-----+-------------------+
|label|features           |
+-----+-------------------+
|1.0  |[2.1, 0.0, 3.4]    |
|0.0  |(3,[0,2],[1.2,4.1])|
+-----+-------------------+

label is usually a Double, and features is an org.apache.spark.ml.linalg.Vector. A vector can be:

• Dense: Stores every value.
• Sparse: Stores the vector size plus only nonzero indices and values.

Sparse vectors are essential for high-dimensional data such as one-hot encoded categories and hashed text. Expanding them into dense arrays can multiply memory and network use.

The default column names are only conventions. Most algorithms let callers change them:

new LogisticRegression()
    .setLabelCol("isFraud")
    .setFeaturesCol("modelFeatures")
    .setPredictionCol("predictedFraud");

Other common output columns are:

• rawPrediction: Unnormalized model scores
• probability: Class probabilities, when the model provides them
• prediction: The selected class or regression value

DataFrame schemas make pipeline stages composable, but checks occur at runtime. Spark validates that required columns exist and have compatible types when a stage is fitted or transformed.

Transformers and Estimators #

Spark ML standardizes algorithms around two abstractions.

A Transformer converts one DataFrame into another:

input DataFrame --transform()--> output DataFrame

It usually appends columns while retaining the input columns. Feature transformers and trained models are both Transformers. For example, Tokenizer turns a text column into an array of words, while LogisticRegressionModel turns a feature vector into prediction columns.

An Estimator learns from a DataFrame and produces a model:

training DataFrame --fit()--> Model

A Model is a Transformer:

test DataFrame --model.transform()--> predictions

For example:

LogisticRegression trainer = new LogisticRegression()
    .setMaxIter(50)
    .setRegParam(0.01);

LogisticRegressionModel model = trainer.fit(training);
Dataset<Row> predictions = model.transform(test);

fit() and transform() return objects that describe distributed DataFrame work. Spark SQL remains lazy where possible; an action such as write, evaluate, count, or collect causes jobs to run.

Parameters #

Every Estimator and Transformer implements a common parameter API. A parameter belongs to a particular stage instance:

LogisticRegression lr = new LogisticRegression()
    .setMaxIter(50)
    .setRegParam(0.01)
    .setElasticNetParam(0.5);

Inspect supported parameters and defaults with:

System.out.println(lr.explainParams());
System.out.println(lr.extractParamMap());

A ParamMap supplied to fit() or transform() overrides values previously set on the instance:

ParamMap override = new ParamMap()
    .put(lr.maxIter(), 100)
    .put(lr.regParam(), 0.05);

LogisticRegressionModel model = lr.fit(training, override);

Because parameters are bound to stage instances, two logistic-regression stages can have independent values in one parameter grid.

Feature Engineering #

Most algorithms expect one numeric vector column, but source data commonly has numeric, categorical, boolean, and text columns. Feature transformers convert those columns into a usable vector.

Numeric Features #

VectorAssembler combines scalar and vector columns:

VectorAssembler assembler = new VectorAssembler()
    .setInputCols(new String[] {"age", "income", "accountAgeDays"})
    .setOutputCol("rawFeatures");

Missing values must be handled deliberately. Imputer can replace missing numeric values using a mean, median, or mode. It does not make missingness semantically harmless; the strategy must fit the data.

StandardScaler, MinMaxScaler, and MaxAbsScaler rescale vectors. Scaling is important for distance-based and gradient-based algorithms, while decision trees normally do not require it.

Be careful with centering sparse vectors. Subtracting a nonzero mean turns implicit zeros into nonzero values and can make a sparse dataset dense.

Categorical Features #

StringIndexer converts category strings to numeric indices:

StringIndexer countryIndexer = new StringIndexer()
    .setInputCol("country")
    .setOutputCol("countryIndex")
    .setHandleInvalid("keep");

OneHotEncoder converts indices into sparse indicator vectors:

OneHotEncoder countryEncoder = new OneHotEncoder()
    .setInputCol("countryIndex")
    .setOutputCol("countryVector");

The indexer must be fitted only on training data. Fitting it independently on test or production data can assign different indices to the same categories. Putting both stages in a fitted Pipeline preserves the training mapping.

handleInvalid controls what happens when null or unseen values arrive. Silently choosing the wrong policy can either fail production scoring or hide data-quality problems.

Text Features #

A basic text pipeline may use:

text
  -> RegexTokenizer
  -> StopWordsRemover
  -> HashingTF or CountVectorizer
  -> IDF
  -> feature vector

HashingTF avoids learning a vocabulary but permits hash collisions. CountVectorizer learns an explicit vocabulary and can be easier to inspect, at the cost of storing and distributing that vocabulary.

Spark ML’s built-in text features are useful for classical models. They are not a replacement for modern language-model tokenizers, embeddings, or GPU-oriented deep-learning frameworks.

Pipelines #

A Pipeline records feature preparation and training as one ordered workflow:

raw data
  -> StringIndexer
  -> OneHotEncoder
  -> VectorAssembler
  -> StandardScaler
  -> LogisticRegression

Transformer stages call transform(). Estimator stages call fit() to produce a Model, then that Model transforms data for subsequent stages.

StringIndexer indexer = new StringIndexer()
    .setInputCol("country")
    .setOutputCol("countryIndex")
    .setHandleInvalid("keep");

OneHotEncoder encoder = new OneHotEncoder()
    .setInputCol("countryIndex")
    .setOutputCol("countryVector");

VectorAssembler assembler = new VectorAssembler()
    .setInputCols(new String[] {"age", "income", "countryVector"})
    .setOutputCol("features");

LogisticRegression classifier = new LogisticRegression()
    .setLabelCol("label")
    .setFeaturesCol("features");

Pipeline pipeline = new Pipeline().setStages(new PipelineStage[] {
    indexer,
    encoder,
    assembler,
    classifier
});

PipelineModel model = pipeline.fit(training);
Dataset<Row> predictions = model.transform(test);

Pipeline.fit() is itself an Estimator operation and returns a PipelineModel, which is a Transformer. Every Estimator in the original pipeline has become a fitted Transformer in the model.

This is more than syntactic convenience. It guarantees that validation and production data use feature mappings learned from training data. Performing feature fitting before splitting a dataset leaks information from validation data into training and produces an optimistic evaluation.

Pipeline stages are listed linearly but may form a DAG through their input and output columns. Stages must be in topological order and should be unique instances because each stage has a unique ID used by the parameter system.

A Complete Training Flow #

A reliable supervised-learning job follows a sequence such as:

1. Read raw data with an explicit schema.
2. Validate labels, ranges, null handling, and row identity.
3. Split data into training, validation, and final test sets.
4. Fit the entire feature-and-model Pipeline on training data only.
5. Select parameters using training and validation data.
6. Evaluate the selected model once on untouched test data.
7. Save the fitted PipelineModel and its operational metadata.

For example:

Dataset<Row>[] split = input.randomSplit(
    new double[] {0.8, 0.2},
    12345L
);

Dataset<Row> training = split[0].cache();
Dataset<Row> test = split[1];

PipelineModel model = pipeline.fit(training);
Dataset<Row> predictions = model.transform(test);

BinaryClassificationEvaluator evaluator =
    new BinaryClassificationEvaluator()
        .setLabelCol("label")
        .setRawPredictionCol("rawPrediction")
        .setMetricName("areaUnderROC");

double areaUnderRoc = evaluator.evaluate(predictions);

Set seeds where APIs support them, but do not assume a distributed job is bit-for-bit deterministic. Partitioning, floating-point reduction order, input file layout, and algorithm implementations can affect results.

For time-dependent data, random splitting is usually wrong. Train on earlier events and validate on later events to avoid learning from the future. Grouped entities may also require a group-aware split so one user’s or device’s events do not appear on both sides.

Evaluation #

Choose an Evaluator that matches the task:

Accuracy alone can be misleading for imbalanced classification. Depending on the business cost, inspect precision, recall, F-measure, area under the precision-recall curve, calibration, and confusion matrices. A model that ranks well may still need a separately selected decision threshold.

Regression metrics such as RMSE, MAE, and \(R^2\) answer different questions. Select a metric based on the cost of errors rather than whichever number looks best.

Evaluation is a distributed Spark job, but some summaries eventually return small matrices, histories, or scalar metrics to the driver. Avoid collecting the entire prediction DataFrame merely to compute a metric that Spark can aggregate in the cluster.

Hyperparameter Tuning #

Spark ML provides two model-selection Estimators:

• CrossValidator evaluates each parameter map over multiple folds.
• TrainValidationSplit evaluates each parameter map with one train/validation split and is cheaper but less statistically robust.

Both require:

1. An Estimator, often a complete Pipeline
2. An array of parameter maps
3. An Evaluator

ParamMap[] grid = new ParamGridBuilder()
    .addGrid(classifier.regParam(), new double[] {0.001, 0.01, 0.1})
    .addGrid(classifier.elasticNetParam(), new double[] {0.0, 0.5, 1.0})
    .build();

CrossValidator cv = new CrossValidator()
    .setEstimator(pipeline)
    .setEvaluator(new BinaryClassificationEvaluator())
    .setEstimatorParamMaps(grid)
    .setNumFolds(3)
    .setParallelism(3)
    .setSeed(12345L);

CrossValidatorModel selectedModel = cv.fit(training);

The example trains:

3 regParam values
x 3 elasticNetParam values
x 3 folds
= 27 model fits

After selecting the best parameters, cross-validation fits that configuration again on the complete input dataset. Large grids therefore consume substantial CPU, executor memory, storage I/O, and time.

parallelism allows parameter configurations to be evaluated concurrently. It does not create free capacity: excessive parallelism can make several training jobs compete for the same executors and perform worse.

Tune the full Pipeline when feature-stage parameters are part of the question. This also fits stateful feature transformers independently within each fold, avoiding validation leakage.

Distributed Execution #

The driver creates Pipeline stages, analyzes DataFrame plans, submits fitting and transformation jobs, and holds the resulting model metadata. Executors scan partitions, calculate feature transformations, compute distributed statistics and gradients, build partial tree statistics, and score rows.

The exact execution pattern depends on the algorithm:

• Feature transforms often use ordinary DataFrame operations over partitions.
• Iterative algorithms repeatedly run jobs and aggregate updates.
• Tree algorithms exchange histograms and split statistics.
• ALS partitions user and item factors and repeatedly exchanges required factor blocks.
• Some learned model state is collected on the driver and later distributed for prediction.

“Distributed” does not mean that one model object is independently trained in each executor. The driver coordinates one fitting operation, while executors perform its partition-level work.

Model size still matters. Coefficients, trees, vocabularies, category maps, and other fitted state may need to reside on the driver, be serialized, or be sent to executors. A dataset that fits in the cluster does not imply that an arbitrarily large model fits.

Cache reused training data when repeated model fits would otherwise reread and recompute it:

Dataset<Row> training = preparedTraining.cache();
training.count();

The action materializes the cache before tuning starts. Unpersist it after the last fit:

training.unpersist();

Caching every intermediate DataFrame wastes memory and can increase spills. Persist only expensive data reused by multiple actions.

Saving and Loading Models #

Most Spark ML Estimators, Models, Pipelines, and PipelineModels implement Spark’s ML persistence API:

model.write()
    .overwrite()
    .save("hdfs:///models/fraud/2026-07-20");

Load it later:

PipelineModel model = PipelineModel.load(
    "hdfs:///models/fraud/2026-07-20"
);

Dataset<Row> predictions = model.transform(input);

A saved PipelineModel includes fitted feature stages and the trained model. It does not automatically include every operational fact needed to reproduce the training run. Store additional metadata such as:

• Spark and application versions
• Source-code revision
• Training-data snapshot or version
• Feature and label definitions
• Split method and random seeds
• Evaluation metrics and selected threshold
• Parameter maps
• Expected input and output schemas

Spark aims for persistence compatibility across minor and patch versions. Compatibility across major Spark versions is best-effort, not guaranteed. Validate both loading and prediction behavior before upgrading production inference jobs.

Batch Inference #

The simplest inference job loads a model and transforms a DataFrame:

PipelineModel model = PipelineModel.load(modelPath);

Dataset<Row> predictions = model.transform(input)
    .select("recordId", "probability", "prediction");

predictions.write()
    .mode(SaveMode.Overwrite)
    .parquet(outputPath);

The transformation is lazy. Executors perform scoring when the output action runs.

Spark ML is naturally suited to batch inference. Low-latency online serving is a different workload. Starting a Spark job for each request is inappropriate, and embedding a Spark driver in every web-service instance is usually operationally expensive.

For online serving, possible approaches include:

• reproducing a compatible model in a serving-oriented runtime;
• exporting supported model types to an interoperable format;
• using a dedicated model server; or
• scoring requests in micro-batches when latency requirements allow it.

Not every Pipeline stage has a portable representation. Test parity between offline Spark features and online features before deploying separate serving code.

Custom Transformers and Estimators #

A custom Transformer can package application-specific feature logic as a pipeline stage. It must define parameters, input/output schema behavior, transform(), and copying/persistence semantics.

Prefer native Spark SQL expressions over row-by-row UDFs:

Dataset<Row> prepared = input.withColumn(
    "amountLog",
    functions.log1p(functions.col("amount"))
);

Native expressions participate in Catalyst optimization and code generation. Java or Scala UDFs introduce an opaque call boundary; Python UDFs additionally cross the JVM/Python boundary unless a vectorized execution path is available.

If custom state must be learned from data, implement an Estimator that returns a corresponding Model. A one-off sequence of DataFrame transformations may be simpler than a custom stage unless persistence, tuning, and pipeline reuse are required.

Spark ML Versus Other ML Systems #

Spark ML is strongest when:

• training data is stored in a distributed data lake;
• feature preparation requires large joins, aggregations, or scans;
• classical algorithms meet the modeling requirement;
• batch training and inference fit the product; and
• the organization already operates Spark reliably.

It is less compelling when:

• the prepared dataset fits comfortably on one machine;
• GPU-based deep learning is required;
• the desired algorithm is not implemented;
• very low-latency online inference is the primary workload; or
• experiment iteration is dominated by cluster startup and scheduling.

For small and medium tabular datasets, scikit-learn, XGBoost, LightGBM, or a single-machine statistics package may be simpler and faster. Deep-learning frameworks such as PyTorch and TensorFlow provide richer neural-network and GPU support. Spark can still prepare their datasets, but using Spark for feature engineering does not require using Spark ML for the final model.

Common Mistakes #

1. Using spark.mllib for new code. Prefer the DataFrame-based spark.ml API unless maintaining an existing RDD application.
2. Fitting feature stages before splitting data. This leaks information into validation and test sets.
3. Collecting training or prediction data to the driver. Keep distributed data as DataFrames and aggregate it on executors.
4. Creating dense vectors from sparse data. Memory and shuffle size can grow dramatically.
5. Ignoring unseen categories and nulls. Define explicit production behavior with stage parameters and data-quality checks.
6. Choosing accuracy for an imbalanced problem. Evaluate metrics aligned with business costs.
7. Launching an enormous parameter grid. Model count is the Cartesian product of parameter values and folds.
8. Assuming every stage scales linearly. Model state, driver aggregation, iteration, skew, and shuffle can remain bottlenecks.
9. Logging or performing remote calls per row. Such side effects execute in parallel on executors and are difficult to retry safely.
10. Treating a saved model as complete lineage. Record data, code, parameters, schemas, and metrics separately.

Spark ML provides a consistent way to express classical machine-learning workflows over distributed DataFrames. Its central abstraction is not an individual algorithm but a fitted PipelineModel that carries feature preparation and prediction together. Use it when distributed data processing is genuinely part of the problem, and keep evaluation, model size, serving requirements, and data leakage in view throughout the workflow.