如何在 PySpark 中进行 groupby 并查找列的唯一项

Question

我有一个 pySpark 数据框，我想按列进行分组，然后在每个组的另一列中找到唯一的项目。

在熊猫中我可以做，

data.groupby(by=['A'])['B'].unique()

我想对我的 Spark 数据框做同样的事情。我可以找到组中项目的 distectCount 并进行计数，如下所示

(spark_df.groupby('A')
    .agg(
        fn.countDistinct(col('B'))
            .alias('unique_count_B'),
        fn.count(col('B'))
            .alias('count_B')
        )
    .show())

但我找不到一些功能来查找组中的独特项目。

为了更清楚地说明，请考虑示例数据框，

df = spark.createDataFrame(
  [(1, "a"), (1, "b"), (1, "a"), (2, "c")],
  ["A", "B"])

我期待得到这样的输出，

+---+----------+
|  A|  unique_B|
+---+----------+
|  1|  [a, b]  |
|  2|  [c]     |
+---+----------+

如何在 pySpark 中获得 pandas 的输出？

Answer 1

我collect_set这样用于我的目的，

(df.groupby('A')
    .agg(
        fn.collect_set(col('B')).alias('unique_count_B')
    )
    .show())

我根据需要得到以下输出，

+---+--------------+
|  A|unique_count_B|
+---+--------------+
|  1|        [b, a]|
|  2|           [c]|
+---+--------------+