2016-07-14 6 views
10

コンテキスト:私はDataFrameに単語とベクトルの2つの列があります。 「ベクトル」の列タイプがVectorUDTの場合ベクターを列に分割する方法 - PySparkを使用する

アン例:

word | vector 

assert | [435,323,324,212...] 

そして、私はこれを取得したい:

word | v1 | v2 | v3 | v4 | v5 | v6 ...... 

assert | 435 | 5435| 698| 356|.... 

質問:することができます私は、各ディメンションのためにいくつかの列にベクターで列を分割する方法

pysparkを使って?事前に

おかげ

答えて

16

1つの可能なアプローチは、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| 
## +-------+-----+-----+-----+ 
+0

パフォーマンスが賢明と、それは'.map/.toDF'関数を使うほうがはるかにスマートです。ほとんどの場合、UDFの実装よりも高速です。 [spark 2.2+の 'vectorized udf'定義を使用していない限り] – tmarthal

関連する問題