Index back probability in a classification prediction in Spark

Question

I'm trying to index back the prediction probability in a classification prediction in Spark. I have an input data for multiclass classifier with labels red, green, blue.

Input dataframe:

+-----+---+---+---+---+---+---+---+---+---+----+----+----+----+
|  _c0|_c1|_c2|_c3|_c4|_c5|_c6|_c7|_c8|_c9|_c10|_c11|_c12|_c13|
+-----+---+---+---+---+---+---+---+---+---+----+----+----+----+
|  red|  0|  0|  0|  1|  0|  0|  0|  2|  3|   2|   2|   0|   5|
|green|  5|  6|  0| 14|  0|  5|  0| 95|  2| 120|   0|   0|   9|
|green|  6|  1|  0|  3|  0|  4|  0| 21| 22|  11|   0|   0|  23|
|  red|  0|  1|  0|  1|  0|  4|  0|  1|  4|   2|   0|   0|   5|
|green| 37|  9|  0| 19|  0| 31|  0| 87|  9| 108|   0|   0| 170|
+-----+---+---+---+---+---+---+---+---+---+----+----+----+----+
only showing top 5 rows

I use StringIndexer to index the label column and VectorAssembler to create feature vector from feature columns.

Parsed Dataframe:

+-----+--------------------+
|label|            features|
+-----+--------------------+
|  1.0|(13,[3,7,8,9,10,1...|
|  0.0|[5.0,6.0,0.0,14.0...|
|  0.0|[6.0,1.0,0.0,3.0,...|
|  1.0|(13,[1,3,5,7,8,9,...|
|  0.0|[37.0,9.0,0.0,19....|
+-----+--------------------+
only showing top 5 rows

A Random Forest Classification model is generated with this data. While querying I will be providing the feature columns to predict the label and its probability.

Query dataframe:

+---+---+---+---+---+---+---+---+---+---+----+----+----+
|_c0|_c1|_c2|_c3|_c4|_c5|_c6|_c7|_c8|_c9|_c10|_c11|_c12|
+---+---+---+---+---+---+---+---+---+---+----+----+----+
| 11| 11|  0| 23|  0|  7|  2| 70| 81| 76|   7|   0|  23|
|  4|  0|  0|  0|  0|  0|  2|  2|  3|  2|   7|   0|   2|
+---+---+---+---+---+---+---+---+---+---+----+----+----+

Parsed query dataframe:

+--------------------+--------------------+
|          queryValue|            features|
+--------------------+--------------------+
|11,11,0,23,0,7,2,...|[11.0,11.0,0.0,23...|
|4,0,0,0,0,0,2,2,3...|(13,[0,6,7,8,9,10...|
+--------------------+--------------------+

Raw prediction from the RFCModel:

+--------------------+--------------------+--------------------+----------+
|          queryValue|            features|         probability|prediction|
+--------------------+--------------------+--------------------+----------+
|11,11,0,23,0,7,2,...|[11.0,11.0,0.0,23...|        [0.67, 0.32]|       0.0|
|4,0,0,0,0,0,2,2,3...|(13,[0,6,7,8,9,10...|        [0.05, 0.94]|       1.0|
+--------------------+--------------------+--------------------+----------+

In the raw prediction, probability column is an array of double with probabilities in the coresponding class index. Say if a row in probability column is [0.67,0.32], it means class 0.0 has probability of 0.67 and class 1.0 has probability of 0.32. The probability column makes sense only when the labels are 0,1,2... In this case, when I use IndexToString to index back the predictions to original labels, probability column will make no sense.

Indexed dataframe:

+--------------------+--------------------+--------------------+----------+
|          queryValue|            features|         probability|prediction|
+--------------------+--------------------+--------------------+----------+
|11,11,0,23,0,7,2,...|[11.0,11.0,0.0,23...|        [0.67, 0.32]|     green|
|4,0,0,0,0,0,2,2,3...|(13,[0,6,7,8,9,10...|        [0.05, 0.94]|       red|
+--------------------+--------------------+--------------------+----------+

I want to index back probability column like below,

+--------------------+--------------------+--------------------------+----------+
|          queryValue|            features|              probability |prediction|
+--------------------+--------------------+--------------------------+----------+
|11,11,0,23,0,7,2,...|[11.0,11.0,0.0,23...|{"red":0.32,"green":0.67} |     green|
|4,0,0,0,0,0,2,2,3...|(13,[0,6,7,8,9,10...|{"red":0.94,"green":0.05} |       red|
+--------------------+--------------------+--------------------------+----------+

For now I'm indexing the probability column by converting the dataframe to List. Is there any feature transformer available in spark to do this?

Answer 1

Tried to solve this using below approach-

I used Iris data to solve this.

1. Sample Input (Top 5 rows)

+------------+-----------+------------+-----------+-----------+
|sepal_length|sepal_width|petal_length|petal_width|      label|
+------------+-----------+------------+-----------+-----------+
|         5.1|        3.5|         1.4|        0.2|Iris-setosa|
|         4.9|        3.0|         1.4|        0.2|Iris-setosa|
|         4.7|        3.2|         1.3|        0.2|Iris-setosa|
|         4.6|        3.1|         1.5|        0.2|Iris-setosa|
|         5.0|        3.6|         1.4|        0.2|Iris-setosa|
+------------+-----------+------------+-----------+-----------+

2. Capture the labels with there indices from StringIndexerModel

   You mentioned -

   I use StringIndexer to index the label column and VectorAssembler to create feature vector from feature columns.

We will use the stringIndexerModel here to get the Map[index, Label]

// in my case, StringIndexerModel is referenced as labelIndexer
val labelToIndex = labelIndexer.labels.zipWithIndex.map(_.swap).toMap
println(labelToIndex)

Result-

Map(0 -> Iris-setosa, 1 -> Iris-versicolor, 2 -> Iris-virginica)

3. Use this map to generate probability json

  import org.apache.spark.ml.linalg.Vector
  val mapToLabel = udf((vector: Vector) => vector.toArray.zipWithIndex.toMap.map{
      case(prob, index) => labelToIndex(index) -> prob
    })
    predictions.select(
      col("features"),
      col("probability"),
      to_json(mapToLabel(col("probability"))).as("probability_json"),
      col("prediction"),
      col("predictedLabel"))
      .show(5,false)

Result-

+-------------------------------------+------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+----------+--------------+
|features                             |probability                                                 |probability_json                                                                                             |prediction|predictedLabel|
+-------------------------------------+------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+----------+--------------+
|(123,[0,37,82,101],[1.0,1.0,1.0,1.0])|[0.7094347002635046,0.174338768115942,0.11622653162055337]  |{"Iris-setosa":0.7094347002635046,"Iris-versicolor":0.174338768115942,"Iris-virginica":0.11622653162055337}  |0.0       |Iris-setosa   |
|(123,[0,39,58,101],[1.0,1.0,1.0,1.0])|[0.7867074275362319,0.12433876811594202,0.0889538043478261] |{"Iris-setosa":0.7867074275362319,"Iris-versicolor":0.12433876811594202,"Iris-virginica":0.0889538043478261} |0.0       |Iris-setosa   |
|(123,[0,39,62,107],[1.0,1.0,1.0,1.0])|[0.5159492704509036,0.2794443583750028,0.2046063711740936]  |{"Iris-setosa":0.5159492704509036,"Iris-versicolor":0.2794443583750028,"Iris-virginica":0.2046063711740936}  |0.0       |Iris-setosa   |
|(123,[2,39,58,101],[1.0,1.0,1.0,1.0])|[0.7822379507920459,0.12164981462756994,0.09611223458038423]|{"Iris-setosa":0.7822379507920459,"Iris-versicolor":0.12164981462756994,"Iris-virginica":0.09611223458038423}|0.0       |Iris-setosa   |
|(123,[2,43,62,101],[1.0,1.0,1.0,1.0])|[0.7049652235193186,0.17164981462756992,0.1233849618531115] |{"Iris-setosa":0.7049652235193186,"Iris-versicolor":0.17164981462756992,"Iris-virginica":0.1233849618531115} |0.0       |Iris-setosa   |
+-------------------------------------+------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+----------+--------------+
only showing top 5 rows