mllib

Spark ships with two ML libraries:

• spark.ml (the ml package) - the recommended DataFrame-based API with Estimator/Transformer/Pipeline abstractions.
• spark.mllib (the mllib package) - the legacy RDD-based API, in maintenance mode.

Both live in the mllib/ source tree. A small mllib-local/ module holds the linear-algebra primitives that have no Spark dependency, so they can be reused by both mllib and ml.

Purpose

• Provide distributed implementations of common machine-learning algorithms (classification, regression, clustering, recommendation, dimensionality reduction).
• Express end-to-end ML workflows as Pipelines of Transformers and Estimators.
• Persist trained models with versioned, schema-aware writers.

Directory layout

mllib-local/                        - vectors, matrices, BLAS bindings (no SparkContext)
  src/main/scala/org/apache/spark/ml/
    linalg/                         - Vector, Matrix, BLAS
    util/                           - utility code with no Spark dep

mllib/
  src/main/scala/org/apache/spark/
    ml/                             - DataFrame-based ML
      Pipeline.scala
      Estimator.scala / Transformer.scala
      Model.scala
      classification/               - LogisticRegression, RandomForest, ...
      regression/                   - LinearRegression, GBT, ...
      clustering/                   - KMeans, BisectingKMeans, GMM, LDA, ...
      recommendation/               - ALS
      feature/                      - StringIndexer, OneHotEncoder, VectorAssembler, ...
      tuning/                       - CrossValidator, TrainValidationSplit
      evaluation/                   - BinaryClassificationEvaluator, ...
      tree/                         - decision trees and ensembles
      r/                            - R-side wrappers for SparkR users
      util/                         - Identifiable, MLReadable, MLWritable
    mllib/                          - RDD-based legacy API (org.apache.spark.mllib)
      classification/, regression/, clustering/, recommendation/,
      feature/, fpm/, optimization/, evaluation/, ...

Key abstractions

Type	What it is
Vector / Matrix (mllib-local/.../linalg/)	Local vector/matrix primitives with sparse and dense forms.
Estimator[E] / Transformer (mllib/.../ml/)	The pipeline contract. Estimator.fit returns a Model.
Pipeline	A linear sequence of stages that fits to a PipelineModel.
MLReadable / MLWritable	Pluggable model persistence to versioned Parquet directories.
Param[T]	Typed, named parameter for an Estimator or Transformer.
Predictor / PredictionModel	Base classes for supervised algorithms.
BarrierTaskContext	Used by distributed gradient-descent style algorithms.
RDD[(label, features)] (legacy spark.mllib)	The original RDD-based input to legacy algorithms.

Pipeline mechanics

graph LR
    A["Raw DataFrame"] --> B[StringIndexer]
    B --> C[OneHotEncoder]
    C --> D[VectorAssembler]
    D --> E[LogisticRegression]
    E --> F[LogisticRegressionModel]
    F -.fit and transform.- A

Pipeline.fit(df) walks the stages in order. For each Estimator, it calls fit on the current DataFrame and pipes the resulting Model.transform(df) into the next stage. The final PipelineModel contains the fitted models from each estimator stage and pure transformers from the rest.

Linear algebra

mllib-local/ wraps a netlib-java BLAS backend by default (org.apache.spark.ml.linalg.BLAS) and can switch to a native implementation when one is on the classpath. Vectors come in two shapes:

• DenseVector - backed by a double[].
• SparseVector - backed by parallel int[]/double[] for indices and values.

Matrices have analogous dense/sparse forms. The wire format used for model serialization is defined in org.apache.spark.ml.linalg.SQLDataTypes.

Model persistence

Models persist as a directory of Parquet files plus a JSON metadata file. The contract is implemented by the DefaultParamsReader/DefaultParamsWriter helpers in mllib/.../ml/util/. Persistence is forward-compatible: each stage records the Spark version it was written under so the loader can apply the right upgrade rules.

Spark Connect for ML

sql/connect/common/.../protobuf/spark/connect/ml.proto and ml_common.proto carry ML operations. The Python client lives in python/pyspark/ml/connect/. Connect ML uses the existing ml.Pipeline API on the server and serializes models to Arrow batches for prediction calls.

Integration points

• Pipelines run on top of the DataFrame API (sql/core).
• Distributed iterative algorithms (gradient descent, ALS) use Spark's broadcast variables, accumulators, and barrier-execution mode.
• ML model persistence uses Parquet (sql/core/.../execution/datasources/parquet/).
• The R wrappers (mllib/.../ml/r/) expose every ML algorithm to SparkR and follow the same versioning rules.

Entry points for modification

• Adding a new ML algorithm: extend Predictor or Estimator, define Params, implement fit/transform, and provide an MLReadable companion. Add a Python wrapper in python/pyspark/ml/.
• Adding a feature transformer: extend Transformer (or UnaryTransformer) and add a Python wrapper.
• Modifying linear-algebra primitives: edit mllib-local/ (no Spark dependency) and run mllib-local/test.
• Adding Connect ML support for an algorithm: extend ml.proto and the Python and JVM client code in python/pyspark/ml/connect/ and sql/connect/server/.../ml/.

Note on spark.mllib

The RDD-based spark.mllib API receives bug fixes only. New algorithms go to spark.ml. If you are starting a new project, target spark.ml.