如何使用 groupby 和聚合将 pyspark 数据框中的行与多列连接起来

Posted 2023-04-15

【中文标题】如何使用 groupby 和聚合将 pyspark 数据框中的行与多列连接起来

【英文标题】：How can I concatenate the rows in a pyspark dataframe with multiple columns using groupby and aggregate

【发布时间】：2020-07-07 21:24:35

【问题描述】：

我有一个包含多列的 pyspark 数据框。比如下面这个。

from pyspark.sql import Row
l = [('Jack',"a","p"),('Jack',"b","q"),('Bell',"c","r"),('Bell',"d","s")]
rdd = sc.parallelize(l)
score_rdd = rdd.map(lambda x: Row(name=x[0], letters1=x[1], letters2=x[2]))
score_card = sqlContext.createDataFrame(score_rdd)

+----+--------+--------+
|name|letters1|letters2|
+----+--------+--------+
|Jack|       a|       p|
|Jack|       b|       q|
|Bell|       c|       r|
|Bell|       d|       s|
+----+--------+--------+

现在我想按“名称”进行分组，并将每行中两列的值连接起来。 我知道该怎么做，但是假设有数千行，那么我的代码就会变得非常难看。 这是我的解决方案。

import pyspark.sql.functions as f
t = score_card.groupby("name").agg(
    f.concat_ws("",collect_list("letters1").alias("letters1")),
    f.concat_ws("",collect_list("letters2").alias("letters2"))
)

这是我将其保存在 CSV 文件中时得到的输出。

+----+--------+--------+
|name|letters1|letters2|
+----+--------+--------+
|Jack|      ab|      pq|
|Bell|      cd|      rs|
+----+--------+--------+

但我主要关心的是这两行代码

f.concat_ws("",collect_list("letters1").alias("letters1")),
f.concat_ws("",collect_list("letters2").alias("letters2"))

如果有数千列，那么我将不得不重复上述代码数千次。有没有更简单的解决方案，这样我就不必为每一列重复 f.concat_ws() 了？

我到处搜索，但找不到解决方案。

【参考方案1】：

是的，您可以在 agg 函数中使用 for 循环并遍历 df.columns。如果有帮助，请告诉我。

    from pyspark.sql import functions as F
    df.show()

    # +--------+--------+----+
    # |letters1|letters2|name|
    # +--------+--------+----+
    # |       a|       p|Jack|
    # |       b|       q|Jack|
    # |       c|       r|Bell|
    # |       d|       s|Bell|
    # +--------+--------+----+

    df.groupBy("name").agg( *[F.array_join(F.collect_list(column), "").alias(column) for column in df.columns if column !='name' ]).show()

    # +----+--------+--------+
    # |name|letters1|letters2|
    # +----+--------+--------+
    # |Bell|      cd|      rs|
    # |Jack|      ab|      pq|
    # +----+--------+--------+