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

Posted 2023-04-17

【中文标题】Pyspark在第二个数据框中加入多行数据框

【英文标题】：Pyspark Joining dataframes with multiple rows in the second dataframe

【发布时间】：2018-07-14 22:15:55

【问题描述】：

我想在名为“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)

我希望连接的输出看起来像这样。有什么办法可以做到吗？

【问题讨论】：

您的意思是vector(7,8) 对应a_2 和vector(9) 对应a_3？

【参考方案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|
+-------+------+-------+

【讨论】：

感谢您的回答。如果有人能给我一个矢量输出的解决方案，我会等待。

【参考方案2】：

我只是添加一个udf 以将收集到的列表转换为@Psidom 的答案中的向量

#importing necessary libraries
from pyspark.sql.functions import udf, collect_list, col
from pyspark.ml.linalg import Vectors, VectorUDT

#udf for changing the collected list to vector
@udf(VectorUDT())
def vectorUdf(x):
    return Vectors.dense(x)

#grouping and aggregation for collecting values and calling the above udf function
vectorDf_2 = df_2.groupBy('TrackID').agg(vectorUdf(collect_list('Value')).alias('Value'))

#joining the two dataframes
Output = df_1.join(vectorDf_2, ['TrackID'])

这应该给你

+-------+-------+---------+
|TrackID|cluster|Value    |
+-------+-------+---------+
|a_1    |1      |[5.0,6.0]|
|a_1    |2      |[5.0,6.0]|
|a_2    |3      |[7.0,8.0]|
|a_3    |1      |[9.0]    |
+-------+-------+---------+

root
 |-- TrackID: string (nullable = true)
 |-- cluster: long (nullable = true)
 |-- Value: vector (nullable = true)

希望回答对你有帮助