如何在 PySpark 中将向量列拆分为单独的列？

PySpark：将向量拆分为列

在 PySpark 中，您可能会遇到带有向量列的 DataFrame，并且需要将其拆分为多个列，每个维度一个。实现方法如下：

For Spark >= 3.0.0

从 Spark 3.0.0 开始，提取向量分量的便捷方法是使用 vector_to_array 函数:

<code class="python">df = df.withColumn("xs", vector_to_array("vector"))

# Pick the first three dimensions for illustration
result = df.select(["word"] + [col("xs")[i] for i in range(3)])</code>

对于 Spark 3.0.0

方法1：RDD转换

一种方法是将DataFrame转换为RDD并手动提取向量分量：

<code class="python">rdd = df.rdd.map(lambda row: (row.word, ) + tuple(row.vector.toArray().tolist()))
result = rdd.toDF(["word"])</code>

方法 2：创建 UDF

或者，您可以创建用户定义函数 (UDF) 并将其应用到向量列：

<code class="python">@udf(ArrayType(DoubleType()))
def to_array(vector):
    return vector.toArray().tolist()

result = df.withColumn("xs", to_array(col("vector"))).select(["word"] + [col("xs")[i] for i in range(3)])</code>

以上是如何在 PySpark 中将向量列拆分为单独的列？的详细内容。更多信息请关注PHP中文网其他相关文章！