如何将向量列拆分为两列？

Question

我使用PySpark.

Spark ML的随机森林输出DataFrame有一列"概率",它是一个具有两个值的向量.我只想两列添加到输出数据帧,"prob1"和"prob2",其对应于载体中的第一和第二值.

我尝试过以下方法:

output2 = output.withColumn('prob1', output.map(lambda r: r['probability'][0]))

但我得到'col应该是列'的错误.

关于如何将一列向量转换为其值列的任何建议？

Answer 1

您可能想要使用一个 UDF 来提取第一个值，并使用另一个 UDF 来提取第二个值。然后，您可以将 UDF 与随机 forrest 数据帧的输出上的 select 调用一起使用。例子：

from pyspark.sql.functions import udf, col

split1_udf = udf(lambda value: value[0], FloatType())
split2_udf = udf(lambda value: value[1], FloatType())
output2 = randomForrestOutput.select(split1_udf(col("probability")).alias("c1"),
                                     split2_udf(col("probability")).alias("c2"))

这应该给你一个数据帧输出2，其中有列c1和c2，对应于存储在列概率中的列表中的第一个和第二个值。