Scala 在预测值大于自定义阈值的情况下，向DataFrame添加值为1的列

我试图向

数据帧添加一列，当输出类概率较高时，该列的值应为1。大概是这样的：

val output = predictions
    .withColumn(
        "easy", 
        when( $"label" === $"prediction" && 
              $"probability" > 0.95, 1).otherwise(0)
    )

问题是，
概率
是一个
向量
，而
0.95
是一个
Double
，因此上述方法不起作用。我真正需要的是
max（$“probability”）>0.95
，当然这也不行

实现这一点的正确方法是什么？
定义自定义项

val findP = udf((label: <type>, prediction: <type>, probability: <type> ) => {
if (label == prediction && vector.toArray.max > 0.95) 1 else 0
})

使用udf，类似于：

val func = (label: String, prediction: String, vector: Vector) => {
  if(label == prediction && vector.toArray.max > 0.95) 1 else 0
}
val output = predictions
  .select($"label", func($"label", $"prediction", $"probability").as("easy"))

下面是一个简单的例子来实现您的问题。
创建一个udf和pass概率列，并为新添加的列返回0或1。在一行中，使用WrappedArray代替数组Vector

  val spark = SparkSession.builder().master("local").getOrCreate()

  import spark.implicits._

  val data = spark.sparkContext.parallelize(Seq(
    (Vector(0.78, 0.98, 0.97), 1), (Vector(0.78, 0.96), 2), (Vector(0.78, 0.50), 3)
  )).toDF("probability", "id")


  data.withColumn("label", label($"probability")).show()

  def label = udf((prob: mutable.WrappedArray[Double]) => {
    if (prob.max >= 0.95) 1 else 0
  })

输出：

+------------------+---+-----+
|       probability| id|label|
+------------------+---+-----+
|[0.78, 0.98, 0.97]|  1|    1|
|      [0.78, 0.96]|  2|    1|
|       [0.78, 0.5]|  3|    0|
+------------------+---+-----+

这起作用了。非常感谢。我要补充的一点是，为了让它工作，我必须找到正确的概率类型。那是一个灯塔。见这个问题：
+------------------+---+-----+
|       probability| id|label|
+------------------+---+-----+
|[0.78, 0.98, 0.97]|  1|    1|
|      [0.78, 0.96]|  2|    1|
|       [0.78, 0.5]|  3|    0|
+------------------+---+-----+