PySpark将Vector拆分为列

最新推荐文章于 2023-09-21 16:23:50 发布

转载最新推荐文章于 2023-09-21 16:23:50 发布 · 4.5k 阅读

6 ·

CC 4.0 BY-SA版权

原文链接：https://www.javaroad.cn/questions/44697

pyspark 专栏收录该内容

2 篇文章

订阅专栏

本文介绍了如何使用PySpark将DataFrame中的向量数据转换为RDD，以及通过定义User Defined Function (UDF)来提取向量值。方法包括RDD的map操作和UDF在DataFrame列转换的应用，适用于处理和操作大型数据集中的向量数据。

一种可能的方法是转换为RDD和从RDD转换：

from pyspark.ml.linalg import Vectors

df = sc.parallelize([
    ("assert", Vectors.dense([1, 2, 3])),
    ("require", Vectors.sparse(3, {1: 2}))
]).toDF(["word", "vector"])

def extract(row):
    return (row.word, ) + tuple(row.vector.toArray().tolist())

df.rdd.map(extract).toDF(["word"])  # Vector values will be named _2, _3, ...

## +-------+---+---+---+
## |   word| _2| _3| _4|
## +-------+---+---+---+
## | assert|1.0|2.0|3.0|
## |require|0.0|2.0|0.0|
## +-------+---+---+---+

另一种解决方案是创建UDF：

from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType, DoubleType

def to_array(col):
    def to_array_(v):
        return v.toArray().tolist()
    return udf(to_array_, ArrayType(DoubleType()))(col)

(df
    .withColumn("xs", to_array(col("vector")))
    .select(["word"] + [col("xs")[i] for i in range(3)]))

## +-------+-----+-----+-----+
## |   word|xs[0]|xs[1]|xs[2]|
## +-------+-----+-----+-----+
## | assert|  1.0|  2.0|  3.0|
## |require|  0.0|  2.0|  0.0|
## +-------+-----+-----+-----+

摘自:https://www.javaroad.cn/questions/44697