如何在PySpark中进行分组和查找列的唯一项

时间:2019-06-19 07:58:01

标签: python pandas pyspark

我有一个pySpark数据框,我想按一列分组,然后在每列的另一列中查找唯一项。

我可以在熊猫里做

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

我想对我的spark数据框执行相同的操作。我可以在组中找到distictCount项并进行计数,像这样

(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中的熊猫输出?

3 个答案:

答案 0 :(得分:1)

您可以使用以下代码,这些代码使用Window函数。

from pyspark.sql import functions as F
from pyspark.sql import Window

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

df.withColumn("distinct AB", 
  F.count("*").over(win)).distinct().show()

结果是:

+---+---+-----------+
|  A|  B|distinct AB|
+---+---+-----------+
|  2|  c|          1|
|  1|  a|          2|
|  1|  b|          1|
+---+---+-----------+

答案 1 :(得分:0)

我出于这种目的使用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]|
+---+--------------+

答案 2 :(得分:-1)

您可以为此使用列表理解 分组之后,假设数据框位于spark_df中,则可以使用:

[row.k for row in spark_df.select('k').distinct().collect()]