将数据集<row> 预测中的向量数据类型“概率”列转换为数组？

Question

我有一个数据集如下：

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

概率的DataType是Vector，我想把它转换成Array，这样数据集就可以保存到数据库了。谢谢！

Answer 1

我遇到了类似的问题，我解决了编写一个自定义 UDF，它从 Vector 中提取一个元素并将其作为双精度返回：

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，其中包含概率向量位置 1 中的元素，作为双精度数。

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