Apache Spark
SQL
AnalysisException
Error Handling
Data Processing

org.apache.spark.sql.AnalysisException Can't extract value from probability

System Design practice on Codemia

Work through 120+ system design problems with detailed solutions, from rate limiters to multi-region storage.

Practice system design

Introduction

This Spark error usually appears after a machine learning model adds a probability column and you try to query it like a normal struct or array column. The confusing part is that the column often prints like a list of numbers, so it looks extractable even when Spark does not treat it that way.

Why the Error Happens

In Spark ML pipelines, columns such as probability, rawPrediction, and features are often stored as ML vector types. A vector is not the same thing as a plain SQL array. Because of that, an expression that works for arrays or structs can fail on a vector column.

A common failing pattern is:

python
predictions.select("probability.0").show()

or:

python
predictions.selectExpr("probability[1] as positive_score").show()

If probability is a vector column, Spark may raise an AnalysisException saying it cannot extract a value from that field.

Reproducing the Problem

Here is a minimal PySpark example with logistic regression:

python
1from pyspark.ml.classification import LogisticRegression
2from pyspark.ml.linalg import Vectors
3from pyspark.sql import SparkSession
4
5spark = SparkSession.builder.getOrCreate()
6
7data = [
8    (0.0, Vectors.dense([0.0, 1.0])),
9    (1.0, Vectors.dense([1.0, 0.0])),
10    (1.0, Vectors.dense([0.8, 0.2])),
11]
12
13df = spark.createDataFrame(data, ["label", "features"])
14model = LogisticRegression().fit(df)
15predictions = model.transform(df)
16
17predictions.select("probability").show(truncate=False)

The probability column will display something like a dense vector, but it is still an ML vector under the hood.

The Correct Fix: Convert the Vector First

In modern Spark, the cleanest fix is vector_to_array:

python
1from pyspark.ml.functions import vector_to_array
2from pyspark.sql.functions import col
3
4result = predictions.withColumn(
5    "probability_array",
6    vector_to_array(col("probability"))
7)
8
9result.select(
10    col("probability_array")[0].alias("p0"),
11    col("probability_array")[1].alias("p1")
12).show()

Once the column becomes a real SQL array, normal array indexing works.

For binary classification, index 1 is commonly treated as the positive-class probability, but always confirm the label ordering used by your model.

SQL-Friendly Version

If you want to use Spark SQL, convert first and then query the derived array column:

python
1result.createOrReplaceTempView("predictions")
2
3spark.sql(
4    """
5    SELECT
6        probability_array[0] AS p0,
7        probability_array[1] AS p1,
8        prediction
9    FROM predictions
10    """
11).show()

The important step is still the same: convert the vector before trying to index into it.

Older Spark Versions

If vector_to_array is unavailable in your environment, a small UDF can work as a fallback:

python
1from pyspark.sql.functions import udf
2from pyspark.sql.types import ArrayType, DoubleType
3
4vector_as_array = udf(lambda v: v.toArray().tolist(), ArrayType(DoubleType()))
5
6result = predictions.withColumn("probability_array", vector_as_array("probability"))
7result.selectExpr("probability_array[1] as positive_score").show()

This is less elegant and may be slower than the built-in function, but it is a practical compatibility option.

Why This Often Appears After ML Pipelines

Spark SQL and Spark ML use related but not identical type systems. DataFrame operations often feel uniform, so it is easy to assume that a printed vector behaves like an array literal.

That assumption breaks when you start selecting nested values. The AnalysisException is Spark telling you that the extraction syntax does not match the underlying column type.

The error also appears when teams move from code that handled arrays to code that now handles VectorUDT, or when they upgrade examples from plain SQL data into ML output data.

Choosing the Right Probability Slot

Do not blindly hard-code index 1 without understanding the model output. For binary logistic regression, the vector is usually ordered by class index. In multiclass models, the vector can contain one probability per class, so you need to know which index corresponds to which label.

You can inspect labels and schema during debugging:

python
predictions.printSchema()
predictions.select("label", "prediction", "probability").show(truncate=False)

That quick check often prevents a second bug after the first one is fixed.

Common Pitfalls

Treating a vector as a SQL struct or array is the direct cause of the exception. Convert the vector explicitly before indexing.

Assuming the positive-class probability is always at the same index can produce silent logic errors. Verify class ordering.

Using a Python UDF when a built-in function exists can slow large jobs. Prefer vector_to_array when your Spark version supports it.

Ignoring the schema makes troubleshooting harder. printSchema() usually reveals whether the column is an array, struct, or vector-backed user-defined type.

Summary

  • Spark ML probability columns are often vectors, not plain SQL arrays.
  • Direct extraction syntax can fail with AnalysisException on vector columns.
  • Convert with vector_to_array before indexing into the values.
  • Use a UDF only as a fallback for older Spark versions.
  • Check label ordering before assuming which probability index you need.

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.

System Design practice on Codemia

Work through 120+ system design problems with detailed solutions, from rate limiters to multi-region storage.

Practice system design

All Rights Reserved.