python - Pyspark在第二个数据框中加入多行数据框

Question

我想在名为“TrackID”的列上加入数据框“df_1”和“df_2”。

 df_1:   cluster    TrackID
           1           a_1
           2           a_1
           3           a_2
           1           a_3

 df_2:   TrackID     Value
           a_1         5
           a_1         6
           a_2         7
           a_2         8
           a_3         9
Output:   
         cluster    TrackID   Value
          1           a_1    Vector(5,6)
          2           a_1    Vector(5,6)
          3           a_2    Vector(7,8)
          1           a_3    Vetor(9)

我希望连接的输出看起来像这样。有没有办法我可以做到这一点？

标签： pythonapache-sparkpysparkapache-spark-sql

Answer 1

如果你对ArrayType 没问题，你可以先通过TrackID聚合第二个数据帧，然后加入第一个数据帧：

import pyspark.sql.functions as F

df_2.groupBy('TrackID').agg(
    F.collect_list('Value').alias('Value')
).join(df_1, ['TrackID']).show()

+-------+------+-------+
|TrackID| Value|cluster|
+-------+------+-------+
|    a_1|[5, 6]|      1|
|    a_1|[5, 6]|      2|
|    a_2|[7, 8]|      3|
|    a_3|   [9]|      1|
+-------+------+-------+