2015-09-14 55 views
3

我试图从使用ML管道和DataFrame API的逻辑模型中提取预测概率。预测概率的输出是一个列向量,存储每个类(0,1)的预测概率,如下所示。我想知道如何才能提取1级的概率。谢谢!如何拆分由ML Pileline逻辑回归生成的预测概率

概率
“[0.13293408418007766,0.8670659158199223]”
“[0.1335112097146626,0.8664887902853374]”

回答

5

UDF像这应该工作:

import org.apache.spark.sql.functions.udf 

val getPOne = udf((v: org.apache.spark.mllib.linalg.Vector) => v(1)) 
model.transform(testDf).select(getPOne($"probability")) 
0

这样的方式也可以达到目标。

import org.apache.spark.mllib.linalg.Vector 
model.transform(dataDF).select("prob").map(row => row.getAs[Vector]("prob").toArray) 

那么你可以得到你想要的东西。

0

您还可以首先将概率列列入org.apache.spark.mllib.linalg.DenseVector,然后提取元素。

model.transform(testDf).select("probability").map(e=> e.asInstanceOf[DenseVector]).map(e=> e(1))