python - 如何在 PySpark 中进行聚合和转置？

Question

我想对非数字字段进行一些聚合，例如collect_list，然后根据不同的列转置到多个列。

这是一个简单的数据集和我想要的结果。

df = sqlContext.createDataFrame([
    ("John", 'Male', 'GA'),
    ("Mary", "Female", 'GA'),
    ("Alex", "Male", "NY"),
    ("Ana", "Female", "NY"),
    ("Amy", "Female", "NY")   
], ["Name", "gender", "state"])

我可以按如下所示按 2 个字段进行分组。

(df
 .groupBy('state', 'gender')
 .agg(collect_list('Name'))
).show()

+-----+------+------------------+
|state|gender|collect_list(Name)|
+-----+------+------------------+
|   NY|  Male|            [Alex]|
|   GA|  Male|            [John]|
|   GA|Female|            [Mary]|
|   NY|Female|        [Ana, Amy]|
+-----+------+------------------+

如何根据性别字段将其转置为两个字段？这就是我要找的。

state |   male   | female
--------------------------
GA    | ['John'] | ['Mary']
NY    | ['Alex'] | ['Ana', 'Amy']

Answer 1

您可以使用 pivot() 转置在groupby之后。 Spark还有一个collect_list()聚合。

df.groupBy('state').pivot('gender').agg(F.collect_list('Name')).show()

# +-----+----------+------+
# |state|    Female|  Male|
# +-----+----------+------+
# |   NY|[Ana, Amy]|[Alex]|
# |   GA|    [Mary]|[John]|
# +-----+----------+------+