Introduction

In PySpark ML, you usually do not call a standalone predict function directly like some local libraries. Instead, you train an estimator to get a model, then call model.transform on a DataFrame to generate prediction columns. This DataFrame-first pattern is central to scalable prediction in Spark.

Train, Fit, and Transform Workflow

Spark ML separates workflow into three stages:

• Estimator, for example LogisticRegression
• Model produced by fit
• Predictions produced by transform

Minimal binary classification example:

1from pyspark.sql import SparkSession
2from pyspark.ml.feature import VectorAssembler
3from pyspark.ml.classification import LogisticRegression
4
5spark = SparkSession.builder.appName("predict-demo").getOrCreate()
6
7rows = [
8    (0, 0.2, 1.1, 0),
9    (1, 0.5, 0.9, 0),
10    (2, 1.3, 2.1, 1),
11    (3, 1.0, 1.8, 1),
12]
13
14df = spark.createDataFrame(rows, ["id", "x1", "x2", "label"])
15
16assembler = VectorAssembler(inputCols=["x1", "x2"], outputCol="features")
17train_df = assembler.transform(df)
18
19lr = LogisticRegression(featuresCol="features", labelCol="label")
20model = lr.fit(train_df)
21
22pred_df = model.transform(train_df)
23pred_df.select("id", "label", "prediction", "probability").show(truncate=False)

The prediction output is a new DataFrame with additional columns, usually prediction, probability, and rawPrediction.

Predict on New Data

For inference, prepare input schema exactly as training expected and run transform.

1new_rows = [
2    (10, 0.3, 1.0),
3    (11, 1.2, 2.0),
4]
5new_df = spark.createDataFrame(new_rows, ["id", "x1", "x2"])
6new_feat_df = assembler.transform(new_df)
7
8new_pred = model.transform(new_feat_df)
9new_pred.select("id", "prediction", "probability").show(truncate=False)

If feature columns or types differ, Spark may fail at runtime, so keep schema contracts explicit.

Use Pipeline for Cleaner Production Inference

Pipelines bundle preprocessing and model into one reusable object.

1from pyspark.ml import Pipeline
2from pyspark.ml.feature import StringIndexer
3from pyspark.ml.classification import RandomForestClassifier
4
5train = spark.createDataFrame([
6    ("A", 0.2, 1.1, 0),
7    ("B", 1.4, 2.2, 1),
8    ("A", 0.3, 0.8, 0),
9], ["category", "x1", "x2", "label"])
10
11indexer = StringIndexer(inputCol="category", outputCol="category_idx")
12assembler = VectorAssembler(inputCols=["category_idx", "x1", "x2"], outputCol="features")
13clf = RandomForestClassifier(featuresCol="features", labelCol="label", numTrees=20)
14
15pipeline = Pipeline(stages=[indexer, assembler, clf])
16pipe_model = pipeline.fit(train)
17
18scored = pipe_model.transform(train)
19scored.select("category", "prediction").show()

With a pipeline model, you call transform once and avoid manual step orchestration.

Save and Load Models for Reuse

For batch or streaming inference jobs, persist models and reload in separate processes.

1model_path = "/tmp/lr_model"
2model.write().overwrite().save(model_path)
3
4from pyspark.ml.classification import LogisticRegressionModel
5loaded_model = LogisticRegressionModel.load(model_path)
6
7loaded_pred = loaded_model.transform(train_df)
8loaded_pred.select("id", "prediction").show()

Always version model artifacts together with feature engineering code and schema expectations.

Evaluate Prediction Quality in Spark

After scoring, validate output quality before shipping results to downstream systems. Spark evaluators let you measure accuracy without collecting full datasets to driver memory.

1from pyspark.ml.evaluation import BinaryClassificationEvaluator
2
3evaluator = BinaryClassificationEvaluator(labelCol=\"label\", rawPredictionCol=\"rawPrediction\")
4auc = evaluator.evaluate(pred_df)
5print(\"AUC:\", auc)

Including this check in pipelines catches model regressions early.

UDF-Based Prediction Is Usually a Last Resort

Some teams try to use Python UDFs for prediction because it resembles local inference APIs. This is usually slower and harder to manage than native Spark ML transforms. Prefer Spark ML transformers unless you must call an external model format unsupported by Spark.

If you must use UDF-based scoring, benchmark carefully and consider Pandas UDFs for better throughput.

Common Pitfalls

• Looking for predict() method instead of using model.transform().
• Scoring DataFrames without the same feature preparation used in training.
• Schema drift between training and inference datasets.
• Saving only model weights and forgetting preprocessing pipeline.
• Overusing Python UDF prediction when native Spark ML can handle the workload.

Summary

• In PySpark ML, prediction is performed through model.transform.
• Keep feature engineering consistent between training and inference.
• Pipelines simplify production scoring and reduce preprocessing mismatch.
• Persist and reload models with explicit versioning.
• Prefer native Spark ML transforms over custom UDF prediction paths.