PySpark从sql Dataframe获取嵌套数组

Question

我有一个来自sql的数据框，如：

+----+----------+
|  id|  category|
+----+----------+
|   1|         1|
|   1|         2|
|   2|         4|
|   3|         1|
|   3|         4|
+----+----------+

我想得到一个嵌套数组，如：

[[Row(category=1), Row(category=2)],
 [Row(category=4)],
 [Row(category=1), Row(category=4)]]

是id的顺序，返回相同id的类别数组的嵌套数组。

我写的是

orders = df.select(df.id).distinct().collect()
map(lambda order: df.where(df.id == order.id).select(df.category).collect(), orders)

非常慢，所以我需要一个更好的方法。 最好的是我可以通过一个查询从数据帧中获取嵌套数组。