如何将Vector拆分为列 - 使用PySpark

1 回答

• 36

  一种可能的方法是转换为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|
  ## +-------+-----+-----+-----+
  

  对于Scala等效，请参阅Spark Scala: How to convert Dataframe[vector] to DataFrame[f1:Double, ..., fn: Double)] .

  回复于 2024-04-28T12:21:52+08:00