Introduction

PySpark’s MultilayerPerceptronClassifier can return class probabilities directly, but many users look only at the prediction column and miss the richer output. The classifier exposes both rawPrediction and probability, and those columns are what you want for confidence-based decisions, ranking, or threshold tuning. A clean workflow is to train the model, transform a DataFrame, and inspect the probability vector per row.

Understand the Output Columns

After calling transform, Spark ML classifiers usually produce these columns:

• 'prediction'
• 'rawPrediction'
• 'probability'

For most application use cases, probability is the one you want. It is a vector where each element corresponds to the estimated probability of one class.

For example, in a three-class problem a probability vector might look like:

[0.05, 0.90, 0.05]

That means the model is assigning the highest confidence to class index 1.

Train a Simple Multilayer Perceptron Model

Here is a small runnable example using a local Spark session.

1from pyspark.sql import SparkSession
2from pyspark.ml.classification import MultilayerPerceptronClassifier
3from pyspark.ml.linalg import Vectors
4
5spark = SparkSession.builder.master("local[*]").appName("mlp-prob-demo").getOrCreate()
6
7train = spark.createDataFrame([
8    (0.0, Vectors.dense([0.0, 0.1])),
9    (0.0, Vectors.dense([0.1, 0.2])),
10    (1.0, Vectors.dense([0.9, 1.0])),
11    (1.0, Vectors.dense([1.0, 0.9])),
12], ["label", "features"])
13
14layers = [2, 4, 2]
15
16clf = MultilayerPerceptronClassifier(
17    layers=layers,
18    blockSize=2,
19    seed=42,
20    maxIter=100,
21    featuresCol="features",
22    labelCol="label",
23)
24
25model = clf.fit(train)

Once the model is trained, run inference on new rows.

Read the probability Column

The probability column is included automatically by default.

1test = spark.createDataFrame([
2    (Vectors.dense([0.05, 0.1]),),
3    (Vectors.dense([0.95, 0.9]),),
4], ["features"])
5
6pred = model.transform(test)
7pred.select("features", "prediction", "probability").show(truncate=False)

This will display a dense vector of per-class probabilities for each input row.

If you only need prediction plus confidence, you can extract the maximum probability from the vector.

Convert Probability Vectors to Usable Columns

Spark stores probabilities as vector objects. For downstream SQL-like work, convert them into arrays and individual columns.

1from pyspark.ml.functions import vector_to_array
2from pyspark.sql import functions as F
3
4result = pred.withColumn("prob_array", vector_to_array("probability")) \
5    .withColumn("p_class_0", F.col("prob_array")[0]) \
6    .withColumn("p_class_1", F.col("prob_array")[1]) \
7    .withColumn("confidence", F.greatest(F.col("prob_array")[0], F.col("prob_array")[1]))
8
9result.select("prediction", "p_class_0", "p_class_1", "confidence").show()

This is especially useful when writing predictions to tables or applying threshold rules in Spark SQL pipelines.

Map Probabilities Back to Class Labels

If your labels were encoded numerically, keep a mapping so the probability vector remains interpretable.

1label_map = {
2    0: "negative",
3    1: "positive",
4}
5
6rows = result.select("prediction", "prob_array").collect()
7for row in rows:
8    predicted_name = label_map[int(row["prediction"])]
9    print(predicted_name, row["prob_array"])

In multiclass problems, documenting that index-to-label mapping is essential. Otherwise, the probability vector is easy to misread.

Use Probabilities for Thresholding, Not Just Ranking

The default prediction picks the class with highest probability, but sometimes you want a custom rule. For binary classification, you might require a higher threshold before predicting the positive class.

1thresholded = result.withColumn(
2    "custom_prediction",
3    F.when(F.col("p_class_1") >= 0.80, F.lit(1.0)).otherwise(F.lit(0.0))
4)
5
6thresholded.select("p_class_1", "prediction", "custom_prediction").show()

This is common in fraud, moderation, and alerting systems where false positives are costly.

rawPrediction Is Not the Same as Probability

Users often confuse rawPrediction with normalized class probabilities. For MultilayerPerceptronClassifier, rawPrediction contains intermediate scores before probability normalization.

You should usually:

• use probability for reporting and thresholds
• use prediction for final default class
• inspect rawPrediction only if you are debugging model behavior

Do not build business thresholds on rawPrediction unless you have a specific reason and understand the model output semantics.

Common Pitfalls

One common mistake is selecting only prediction and ignoring probability, then later trying to reconstruct model confidence.

Another issue is treating the probability vector as if it were already labeled with class names. Spark keeps only numeric positions, so you need to maintain your own class mapping.

A third mistake is using rawPrediction for thresholding when the probability column is the correctly normalized output for most downstream tasks.

Summary

• 'MultilayerPerceptronClassifier exposes probabilities in the probability column after transform.'
• Use select("probability") or convert the vector into array columns for downstream processing.
• Keep an explicit mapping from class index to business label.
• Use probabilities for thresholds and confidence-based decisions.
• Treat rawPrediction and probability as different outputs with different purposes.