java - 将数据集 预测中的 vector 数据类型 "probability"列转换为数组？

Question

我有一个数据集如下:

 Dataset<Row> result =   result.select("Probability","label","prediction");

概率的数据类型是Vector，我想将其转换为Array，以便将数据集保存到数据库中。谢谢!

Answer 1

我遇到了类似的问题，我解决了编写自定义 UDF 的问题，该 UDF 从 Vector 中提取元素并将其作为 double 值返回:

static UDF2<Vector, Integer, Double> probability_at = (UDF2<Vector, Integer, Double>) Vector::apply;

要使用 UDF，您首先需要在 SparkSession 中注册它:

spark.udf().register("probability_at", probability_at, DoubleType);

你可以使用它:

Dataset<Row> result = df.select(col("label"), col("prediction"), expr("probability_at(Probability, 1)").as("probability_1"));

此时，result DataFrame 中有一个新列 probability_1，其中包含概率 vector 位置 1 中的元素( double 值)。

如果您需要从 Vector 中提取多个项目，则可以多次调用该函数，或者您可以创建一个类似的 UDF 来循环整个 Vector 并返回一个 double 数组。