PySpark：如何将行转换为向量？

Question

我在一个数据帧的工作有三列，可乐，COLB和COLC

+---+-----+-----+-----+ 
|id |colA |colB |colC | 
+---+-----+-----+-----+ 
| 1 | 5 | 8 | 3 | 
| 2 | 9 | 7 | 4 | 
| 3 | 3 | 0 | 6 | 
| 4 | 1 | 6 | 7 | 
+---+-----+-----+-----+ 

我需要合并的可乐，COLB和COLC列得到这样的下面新建一个数据帧：

+---+--------------+ 
|id |  colD  | 
+---+--------------+ 
| 1 | [5, 8, 3] | 
| 2 | [9, 7, 4] | 
| 3 | [3, 0, 6] | 
| 4 | [1, 6, 7] | 
+---+--------------+ 

这是获得第一个数据帧的pyspark代码：

l=[(1,5,8,3),(2,9,7,4), (3,3,0,6), (4,1,6,7)] 
names=["id","colA","colB","colC"] 
db=sqlContext.createDataFrame(l,names) 
db.show() 

如何将行转换为矢量？有谁能帮助我吗？ 感谢

Answer 1

它实际上取决于你想要的数据类型colD。如果你想要一个VectorUDT列，那么使用VectorAssembler是正确的转换。如果你只是想把这些字段组合成一个数组，那么UDF是不必要的。您可以使用内置的array功能列组合：

>>> from pyspark.sql.functions import array 
>>> db.select('id',array('colA','colB','colC').alias('colD')).show() 

+---+---------+ 
| id|  colD| 
+---+---------+ 
| 1|[5, 8, 3]| 
| 2|[9, 7, 4]| 
| 3|[3, 0, 6]| 
| 4|[1, 6, 7]| 
+---+---------+ 

这实际上给性能提升比其他转换，因为pyspark没有序列化你的UDF。

Answer 2

你可以从pyspark.ml使用vectorassembler，

from pyspark.ml.feature import VectorAssembler 
newdb = VectorAssembler(inputCols=["colA", "colB", "colC"], outputCol="colD").transform(db) 
newdb.show() 
+---+----+----+----+-------------+ 
| id|colA|colB|colC|   colD| 
+---+----+----+----+-------------+ 
| 1| 5| 8| 3|[5.0,8.0,3.0]| 
| 2| 9| 7| 4|[9.0,7.0,4.0]| 
| 3| 3| 0| 6|[3.0,0.0,6.0]| 
| 4| 1| 6| 7|[1.0,6.0,7.0]| 
+---+----+----+----+-------------+ 

，或者如果你愿意，可以使用UDF做行向组成，

from pyspark.sql import functions as F 
from pyspark.sql.types import * 
udf1 = F.udf(lambda x,y,z : [x,y,z],ArrayType(IntegerType())) 
df.select("id",udf1("colA","colB","colC").alias("colD")).show() 
+---+---------+ 
| id|  colD| 
+---+---------+ 
| 1|[5, 8, 3]| 
| 2|[9, 7, 4]| 
| 3|[3, 0, 6]| 
| 4|[1, 6, 7]| 
+---+---------+ 

希望这有助于！