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.
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 , resulting in approximately 0.88.
• For class 1, this would instead be , resulting in approximately 0.12.
Technical Example: Logistic Regression
Consider the use of Spark MLlib's logistic regression for binary classification:
- Training the Model:• Here,
rawPredictionmight look like[-0.5, 0.5]indicating the log-odds for both classes. •probabilitymight 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,rawPredictionextends to vector outputs corresponding to each class, andprobabilityprovides probability vectors summing up to one. • Pipeline Component: Both columns are typically part of the DataFrame produced by thetransformmethod 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
- What do maskers really do in SHAP package and fit them to train or test?
- What do the functions tf.squeeze and tf.nn.rnn do?
- What do the options in ConfigProto like allow_soft_placement and log_device_placement mean?
- What do the TensorFlow Dataset's functions cache and prefetch do?
- What does -1 in numpy reshape mean?
- What does -1 mean in numpy reshape?
- What do foreachBatches contain in a streaming query from multiple Kafka topics?
- What do you use Apache Kafka for?

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseTrack 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.