Scala Spark - 将矢量列拆分为Spark DataFrame中的单独列

时间:2018-04-19 02:07:19

标签: scala apache-spark dataframe countvectorizer

我有一个Spark DataFrame,其中有一个包含Vector值的列。矢量值都是n维的,也就是具有相同的长度。我还有一个列名Array("f1", "f2", "f3", ..., "fn")的列表,每个列对应于向量中的一个元素。

some_columns... | Features
      ...       | [0,1,0,..., 0]

to

some_columns... | f1 | f2 | f3 | ... | fn

      ...       | 0  | 1  | 0  | ... | 0

实现这一目标的最佳方法是什么?我想到了一种方法,即使用createDataFrame(Row(Features), featureNameList)创建一个新的DataFrame然后与旧的一起加入,但它需要使用spark context来使用createDataFrame。我只想转换现有的数据框。我也知道.withColumn("fi", value)但如果n很大,该怎么办?

我是Scala和Spark的新手,但却找不到任何好的例子。我认为这可能是一项常见任务。我的具体情况是我使用了CountVectorizer并希望单独恢复每一列以获得更好的可读性,而不是只有矢量结果。

1 个答案:

答案 0 :(得分:10)

一种方法是将vector列转换为array<double>,然后使用getItem提取单个元素。

import org.apache.spark.sql.functions._
import org.apache.spark.ml._

val df = Seq( (1 , linalg.Vectors.dense(1,0,1,1,0) ) ).toDF("id", "features")
//df: org.apache.spark.sql.DataFrame = [id: int, features: vector]

df.show
//+---+---------------------+
//|id |features             |
//+---+---------------------+
//|1  |[1.0,0.0,1.0,1.0,0.0]|
//+---+---------------------+

// A UDF to convert VectorUDT to ArrayType
val vecToArray = udf( (xs: linalg.Vector) => xs.toArray )

// Add a ArrayType Column   
val dfArr = df.withColumn("featuresArr" , vecToArray($"features") )

// Array of element names that need to be fetched
// ArrayIndexOutOfBounds is not checked.
// sizeof `elements` should be equal to the number of entries in column `features`
val elements = Array("f1", "f2", "f3", "f4", "f5")

// Create a SQL-like expression using the array 
val sqlExpr = elements.zipWithIndex.map{ case (alias, idx) => col("featuresArr").getItem(idx).as(alias) }

// Extract Elements from dfArr    
dfArr.select(sqlExpr : _*).show
//+---+---+---+---+---+
//| f1| f2| f3| f4| f5|
//+---+---+---+---+---+
//|1.0|0.0|1.0|1.0|0.0|
//+---+---+---+---+---+