scala - mllib Vector 的最大值？

Question

我使用 mllib 创建了 Apache Spark 机器学习管道。评估器结果是一个带有“概率”列的 DataFrame，它是概率的 mllib 向量(类似于 scikit-learn 中的 predict_proba)。

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

我想要一个包含此概率最大值的新列。有什么想法吗？

Answer 1

Vector 没有 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"))