如何使用 groupby 和聚合將 pyspark dataframe 中的行與多列連接起來

Question

我有一個多列的 pyspark dataframe。 比如下面這張。

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 文件中時得到的 output。

+----+--------+--------+
|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() 了？

我到處搜索，但無法找到解決方案。

Answer 1

是的，您可以在 agg function 中使用 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|
    # +----+--------+--------+