Using Spark ML

SparkML is the distributed machine learning library that ships inside the pyspark.ml package. It exposes the same conceptual surface as scikit-learn (transformers, estimators, pipelines, evaluators) but every operation executes through Spark's DataFrame engine, which means training and inference run in parallel across executors. The decision to use SparkML is not driven by familiarity with the API; it is driven by where the data lives, how large it is, and what the inference target needs to look like at production scale.

The defining property of SparkML is that it operates on a Spark DataFrame, not on a pandas frame or a NumPy array. Every feature column the model reads must be present in that DataFrame, and the model assembles its feature vector through a VectorAssembler that compresses many columns into one features column of type VectorUDT. The estimator then fits on the features and label columns, and the resulting Model (a transformer) appends a prediction column when called via .transform(). This DataFrame-in, DataFrame-out pattern is what makes SparkML the right surface for batch inference jobs that already read terabytes of Delta data, and for Structured Streaming pipelines that score arriving records continuously.

Conceptually SparkML splits its components into three families. Transformers are stateless objects that map one DataFrame to another (a StringIndexer once fitted, a Tokenizer, a Bucketizer). Estimators have a .fit() method that consumes a DataFrame and returns a fitted Model that itself is a Transformer (a LogisticRegression, a GBTClassifier, a RandomForestRegressor). A Pipeline is itself an Estimator that wraps a sequence of stages and produces a PipelineModel when fitted. The PipelineModel is what gets logged through mlflow.spark.log_model and registered in Unity Catalog for downstream serving.

There are three sub-domains that sit underneath the SparkML umbrella and that every model developer at this tier is expected to operate fluently: scaling and tuning (how the pipeline grows beyond one node), advanced MLflow usage (how experiments are tracked and audited), and streaming feature generation (how features are produced in real time for online consumption).

Scaling and Tuning

SparkML estimators are themselves data-parallel; the data is partitioned across executors and gradient or tree-split computations happen in parallel. This handles the case where data is too large for one node, but it does not by itself address the orthogonal problem of searching across many hyperparameter combinations. For hyperparameter sweeps, Spark MLlib ships CrossValidator and TrainValidationSplit which evaluate a ParamGridBuilder grid in parallel via the parallelism parameter. For single-node estimators wrapped inside a tuning loop, Optuna with SparkTrials distributes the trials across executors. Either way, the scaling axis is "more trials in parallel," not "bigger model per trial."

from pyspark.ml.classification import LogisticRegression
from pyspark.ml.tuning import CrossValidator, ParamGridBuilder
from pyspark.ml.evaluation import BinaryClassificationEvaluator

lr = LogisticRegression(featuresCol="features", labelCol="churned")
grid = (
 ParamGridBuilder()
 .addGrid(lr.regParam, [0.0, 0.01, 0.1])
 .addGrid(lr.elasticNetParam, [0.0, 0.5, 1.0])
 .build()
)

cv = CrossValidator(
 estimator=lr,
 estimatorParamMaps=grid,
 evaluator=BinaryClassificationEvaluator(labelCol="churned"),
 numFolds=5,
 parallelism=4,
 seed=42,
)
cv_model = cv.fit(train_df)

💡 Exam Trap: parallelism on CrossValidator controls how many parameter combinations are evaluated concurrently, not how many partitions the input DataFrame is split into. A high parallelism value with a small cluster will starve each trial of executors and lengthen wall-clock time rather than shorten it.

Advanced MLflow Usage

Every SparkML training job belongs inside an mlflow.start_run() context. For SparkML specifically, mlflow.spark.autolog() captures the input DataFrame schema, the pipeline stages, the fitted parameters, and the model artifact in one call. When the experiment is more complex than a single training run (nested cross-validation, a child run per trial in a tuning sweep, a parent run that aggregates regional sub-models), the canonical pattern is the nested run: a parent start_run opens, and each inner training call opens another start_run(nested=True) that becomes a child in the MLflow UI tree.

import mlflow
import mlflow.spark

mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/Shared/acme_ml/churn_pipeline")

with mlflow.start_run(run_name="parent_sweep") as parent_run:
 mlflow.log_param("dataset_version", "v17")
 for region in ["NA", "EMEA", "APAC"]:
 with mlflow.start_run(run_name=f"child_{region}", nested=True):
 model = pipeline.fit(train_df.filter(f"region = '{region}'"))
 metric = evaluator.evaluate(model.transform(val_df))
 mlflow.log_metric("val_auc", metric)
 mlflow.spark.log_model(
 spark_model=model,
 artifact_path="model",
 registered_model_name=f"acme_ml.model_dev.churn_{region.lower()}",
 )

Streaming Feature Generation

Real-time features come from one of two sources: a Structured Streaming job that reads from a streaming source (Kafka, Event Hubs, Auto Loader on cloud storage) and writes engineered features into a Delta table that the online table mirrors; or an on-demand feature function defined in Unity Catalog that computes the feature inside the serving endpoint at request time from raw input arguments. The streaming pipeline is the right choice when the feature depends on historical aggregates (rolling 5-minute click counts, session-level conversion rates); the on-demand function is the right choice when the feature is a pure function of the request payload (haversine distance between two coordinates in the request, hour of day, payload size).

from pyspark.sql.functions import col, window, count

clicks_stream = (
 spark.readStream
 .format("delta")
 .table("acme_ml.bronze.clickstream")
)

(clicks_stream
 .withWatermark("event_time", "10 minutes")
 .groupBy(window("event_time", "5 minutes"), col("user_id"))
 .agg(count("*").alias("click_count_5m"))
 .writeStream
 .option("checkpointLocation", "/Volumes/acme_ml/state/click_features")
 .outputMode("append")
 .toTable("acme_ml.features.user_click_count_5m")
)

📊 Limit: A Structured Streaming feature pipeline requires a checkpointLocation on a UC Volume or DBFS path; without it, the stream cannot recover state across restarts and the offline feature table will skip or duplicate rows.

⚠️ Anti-Pattern: Computing a windowed aggregate inside a feature function called from the serving endpoint. The serving endpoint sees only the request payload and cannot look back across historical events. The aggregation must run upstream in Structured Streaming and write to the offline Delta table; the online table then mirrors the latest value for low-latency lookups.

🎯 Scenario: A team has a fraud model trained on a feature tx_count_last_hour. They configure an on-demand function that calls back to a Delta table at request time. Latency spikes from 30 ms to 800 ms under load because the request path is doing a full Delta scan. The fix is to precompute the feature in a Structured Streaming job, write to a feature Delta table, publish it as an online table, and let the endpoint resolve the value through a FeatureLookup. The on-demand function path is reserved for features that can be computed purely from the request payload.

Decision Anchor: Choose SparkML when the data does not fit one executor or when the inference path is itself a Spark job (batch or Structured Streaming). Choose single-node frameworks wrapped in Pandas Function APIs when each row independently fits memory but the volume requires parallelism across the cluster. Choose Mosaic AI Model Serving with a PyFunc model when the surface is request-response real time.