Spark MLlib
DataFrame
rawPrediction
probability
machine learning

What do columns ‘rawPrediction’ and ‘probability’ of DataFrame mean in Spark MLlib?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Spark MLlib is a powerful library within the Apache Spark ecosystem, which facilitates scalable machine learning. When using Spark MLlib's models, you often encounter columns such as rawPrediction and probability in the resulting DataFrame from transformations made using models such as logistic regression or decision trees. Understanding these columns is crucial for interpreting model outputs accurately.

Understanding rawPrediction

and probability Columns in Spark MLlib

When you train a machine learning model using Spark MLlib, the output typically includes predictions in the form of a DataFrame. For classification problems, this DataFrame can have up to three primary columns: prediction , rawPrediction , and probability . Here, we focus on rawPrediction and probability .

rawPrediction

The rawPrediction column contains the raw output scores from the classification model. These scores often represent the decision function—for instance, in SVM, these would be the margins. In the context of logistic regression, these are the log-odds or the values before the logistic function is applied. As such, they don’t constitute probabilities directly but are instead unbounded values which need further processing to become probabilities.

For example, in a binary logistic regression model:

• The model might output a rawPrediction of 2.0 for class 0 and -2.0 for class 1. • These results reflect the distance of the sample point from the decision boundary in feature space.

probability

The probability column, on the other hand, displays the predicted probabilities of each class. This is generated by applying a link function—most commonly the softmax function—to the rawPrediction . For binary classification, this is equivalent to transforming the log-odds from logistic regression to a probability using the logistic function.

For example:

• Given a rawPrediction of 2.0, the probability for class 0 might be calculated as 1/(1+e2.0)1 / (1 + e^{-2.0}), resulting in approximately 0.88. • For class 1, this would instead be e2.0/(1+e2.0)e^{-2.0} / (1 + e^{-2.0}), resulting in approximately 0.12.

Technical Example: Logistic Regression

Consider the use of Spark MLlib's logistic regression for binary classification:

  1. Training the Model:
    • Here, rawPrediction might look like [-0.5, 0.5] indicating the log-odds for both classes. • probability might convert this into [0.38, 0.62] using the logistic function, indicating a 38% chance of class 0 and 62% chance of class 1. • Multi-Class Classification: For multi-class problems, rawPrediction extends to vector outputs corresponding to each class, and probability provides probability vectors summing up to one. • Pipeline Component: Both columns are typically part of the DataFrame produced by the transform method in an MLlib-based pipeline, ideal for further evaluation or debugging. • Threshold Adjustments: While probabilities are often used directly to assign predictions, custom thresholds can be applied to adjust sensitivity/specificity tradeoffs.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.