将数据集预测中的"probability"列转换为<row>数组的向量数据类型？

Question

我有一个数据集如下：

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

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

Answer 1

我遇到了类似的问题，我解决了编写自定义UDF的问题，该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数据帧中有一个新的列probability_1，其中包含概率向量位置 1 中的元素作为双精度。

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