我有一个pyspark数据框,其中已对数据进行分组以用collect_list列出。
from pyspark.sql.functions import udf, collect_list
from itertools import combinations, chain
#Create Dataframe
df = spark.createDataFrame( [(1,'a'), (1,'b'), (2,'c')] , ["id", "colA"])
df.show()
>>>
+---+----+
| id|colA|
+---+----+
| 1| a|
| 1| b|
| 2| c|
+---+----+
#Group by and collect to list
df = df.groupBy(df.id).agg(collect_list("colA").alias("colAlist"))
df.show()
>>>
+---+--------+
| id|colAList|
+---+--------+
| 1| [a, b]|
| 2| [c]|
+---+--------+
然后我使用一个函数来查找列表元素到新列表的所有组合
allsubsets = lambda l: list(chain(*[combinations(l , n) for n in range(1,len(l)+1)]))
df = df.withColumn('colAsubsets',udf(allsubsets)(df['colAList']))
所以我会期待
+---+--------------------+
| id| colAsubsets |
+---+--------------------+
| 1| [[a], [b], [a,b]] |
| 2| [[b]] |
+---+--------------------+
但是我得到了
df.show()
>>>
+---+--------+-----------------------------------------------------------------------------------------+
|id |colAList|colAsubsets |
+---+--------+-----------------------------------------------------------------------------------------+
|1 |[a, b] |[[Ljava.lang.Object;@75e2d657, [Ljava.lang.Object;@7f662637, [Ljava.lang.Object;@b572639]|
|2 |[c] |[[Ljava.lang.Object;@26f67148] |
+---+--------+-----------------------------------------------------------------------------------------+
任何想法怎么办?然后也许如何将列表展平到不同的行?
答案 0 :(得分:3)
您需要做的就是以扁平的方式从chain
和combinations
创建的对象中提取元素
正在更改
allsubsets = lambda l: list(chain(*[combinations(l , n) for n in range(1,len(l)+1)]))
到以下
allsubsets = lambda l: [[z for z in y] for y in chain(*[combinations(l , n) for n in range(1,len(l)+1)])]
应该给您
+---+---------+------------------+
|id |colA_list|colAsubsets |
+---+---------+------------------+
|1 |[a, b] |[[a], [b], [a, b]]|
|2 |[c] |[[c]] |
+---+---------+------------------+
我希望答案会有所帮助
答案 1 :(得分:2)
改进@RameshMaharjan答案,以将列表展平到不同的行:
您必须在数组上使用explode。您必须先指定udf的类型,以便它不返回StringType。
from pyspark.sql.functions import explode
from pyspark.sql.types import ArrayType, StringType
allsubsets = lambda l: [[z for z in y] for y in chain(*[combinations(l , n) for n in range(1,len(l)+1)])]
df = df.withColumn('colAsubsets', udf(allsubsets, ArrayType(ArrayType(StringType())))(df['colAList']))
df = df.select('id', explode('colAsubsets'))
结果:
+---+------+
| id| col|
+---+------+
| 1| [a]|
| 1| [b]|
| 1|[a, b]|
| 2| [c]|
+---+------+