2017-10-18 91 views
1

我使用mllib创建了一个使用Apache Spark的ML管道。 评估结果是一个DataFrame,其中有一列“probability”,它是概率的mllib向量(​​类似于scikit-learn中的predict_proba)。mllib矢量的最大值?

val rfPredictions = rfModels.bestModel.transform(testing) 
val precision = evaluator.evaluate(rfPredictions) 

我想这样的事情没有成功:

rfPredictions.select("probability").map{c => c.getAs[Vector](1).max} 
<console>:166: error: value max is not a member of 
org.apache.spark.mllib.linalg.Vector 

我想这个概率最大的新列。有任何想法吗?

回答

3

载体没有max方法。尝试toArray.max

rfPredictions.select("probability").map{ c => c.getAs[Vector](1).toArray.max } 

argmax

rfPredictions.select("probability").map{ c => { 
    val v = c.getAs[Vector](1) 
    v(v.argmax) 
}} 

要添加最大为新列定义一个UDF与withColumn功能使用:

val max_proba_udf = udf((v: Vector) => v.toArray.max) 

rfPredictions.withColumn("max_prob", max_proba_udf($"probability")) 
+0

是的,我认为作品。作为结果,我得到了一个RDD [Double],所以我无法将其添加为DataFrame中的列。我怎样才能做到这一点?谢谢! – marlanbar

1

星火> 2.0

随着毫升,不mllib这将在下一个工作方式:

import org.apache.spark.ml.linalg.DenseVector 

just_another_df.select("probability").map{ c => c.getAs[DenseVector](0).toArray.max } 

使用UDF

import org.apache.spark.ml.linalg.DenseVector 

val max_proba_udf = udf((v: DenseVector) => v.toArray.max) 

val rfPredictions = just_another_df.withColumn("MAX_PROB", max_proba_udf($"probability"))