apache-spark - 如何在PySpark中按多列分组并收集列表？

这是我的问题：
我有这个RDD：

a = [[u'PNR1',u'TKT1',u'TEST',u'a2',u'a3'],[u'PNR1',u'TKT1',u'TEST',u'a5',u'a6'],[u'PNR1',u'TKT1',u'TEST',u'a8',u'a9']]

rdd= sc.parallelize (a)

然后我尝试：

rdd.map(lambda x: (x[0],x[1],x[2], list(x[3:])))

.toDF(["col1","col2","col3","col4"])

.groupBy("col1","col2","col3")

.agg(collect_list("col4")).show

最后我应该找到这个：

[col1,col2,col3,col4]=[u'PNR1',u'TKT1',u'TEST',[[u'a2',u'a3'][u'a5',u'a6'][u'a8',u'a9']]]

但是问题是我无法收集列表。

如果有人可以帮助我，我将不胜感激

最佳答案

我终于找到了解决方案，这不是最好的方法，但是我可以继续工作...

from pyspark.sql.functions import udf
from pyspark.sql.functions import *

def example(lista):
    d = [[] for x in range(len(lista))]
    for index, elem in enumerate(lista):
      d[index] = elem.split("@")
    return d
example_udf = udf(example, LongType())

a = [[u'PNR1',u'TKT1',u'TEST',u'a2',u'a3'],[u'PNR1',u'TKT1',u'TEST',u'a5',u'a6'],[u'PNR1',u'TKT1',u'TEST',u'a8',u'a9']]

rdd= sc.parallelize (a)

df = rdd.toDF(["col1","col2","col3","col4","col5"])

df2=df.withColumn('col6', concat(col('col4'),lit('@'),col('col5'))).drop(col("col4")).drop(col("col5")).groupBy([col("col1"),col("col2"),col("col3")]).agg(collect_set(col("col6")).alias("col6"))

df2.map(lambda x: (x[0],x[1],x[2],example(x[3]))).collect()

它给出：

[(u'PNR1', u'TKT1', u'TEST', [[u'a2', u'a3'], [u'a5', u'a6'], [u'a8', u'a9']])]

希望此解决方案可以对其他人有所帮助。

感谢您的所有答复。