[SW-1281] Fix bad representation of predictionCol on H2OMOJOModel (#1199

) (cherry picked from commit d3f0662)
h2oai · May 15, 2019 · 90dc7c6 · 90dc7c6
1 parent a0b6c64
commit 90dc7c6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ml/src/main/scala/org/apache/spark/ml/h2o/models/H2OMOJOPipelineModel.scala b/ml/src/main/scala/org/apache/spark/ml/h2o/models/H2OMOJOPipelineModel.scala
@@ -190,10 +190,10 @@ class H2OMOJOPipelineModel(override val uid: String)
       val func = udf[Double, Double] {
         identity
       }
-      func(col(s"$getPredictionCol().`$column`")).alias(column)
+      func(col(s"${getPredictionCol()}.`$column`")).alias(column)
     } else {
       val func = selectFromArray(getOutputNames().indexOf(column))
-      func(col(s"$getPredictionCol().preds")).alias(column)
+      func(col(s"${getPredictionCol()}.preds")).alias(column)
     }
   }