Scala中如何提取XGBoost评分结果里概率向量的第二个元素？

搞定这个需求很简单，根据你的Spark版本，有两种实用的方法可以选择：

方法1：用Spark内置函数（Spark 3.0+首选） #

Spark 3.0及以后版本自带了vector_to_array函数，能直接把ML库的Vector类型转换成普通数组，然后我们就可以用getItem取出第二个元素（注意：数组和向量的索引都是从0开始的，所以第二个元素对应索引1）：

import org.apache.spark.sql.functions._

// 转换向量为数组，提取目标元素，最后保留需要的列
val resultDF = predictions
  .withColumn("prob_array", vector_to_array(col("probabilities")))
  .withColumn("probability", col("prob_array").getItem(1))
  .select("some_key", "probability")

// 查看最终结果
resultDF.show()

这种方法不用自己写UDF，代码简洁还自带性能优化，推荐优先用它。

方法2：自定义UDF（兼容Spark 2.x） #

如果你的Spark版本低于3.0，那就自定义一个UDF来处理Vector类型：

import org.apache.spark.sql.functions._
import org.apache.spark.ml.linalg.Vector

// 写个UDF，输入向量就返回它的第二个元素
val getSecondProb = udf((probVec: Vector) => probVec(1))

// 生成目标DataFrame
val resultDF = predictions
  .withColumn("probability", getSecondProb(col("probabilities")))
  .select("some_key", "probability")

resultDF.show()

小提示 #

两种方法最终都会生成你想要的结构：some_key搭配单个概率值的probability列。记得别搞混索引——向量里的第一个元素是索引0，第二个才是索引1哦。

内容的提问来源于stack exchange，提问作者BlindMowingIsMindBlowing