使用 groupby 计数的 Spark 过滤数据

Posted 2023-04-15

【中文标题】使用 groupby 计数的 Spark 过滤数据

【英文标题】：Spark Filter data with groupby count

【发布时间】：2017-08-15 06:41:20

【问题描述】：

一个数据框 A_df 像：

+------+----+-----+
|   uid|year|month|
+------+----+-----+
|     1|2017|   03|
      1|2017|   05|
|     2|2017|   01|
|     3|2017|   02|
|     3|2017|   04|
|     3|2017|   05|
+------+----+-----+

我想要过滤列 uid 出现时间超过 2 次， 预期结果：

+------+----+-----+
|   uid|year|month|
+------+----+-----+
|     3|2017|   02|
|     3|2017|   04|
|     3|2017|   05|
+------+----+-----+

如何通过 Scala 获得此结果？ 我的解决方案：

val condition_uid = A_df.groupBy("uid")
                  .agg(count("*").alias("cnt"))
                  .filter("cnt > 2").select("uid")
val results_df = A_df.join(condition_uid, Seq("uid"))

还有更好的答案吗？

【参考方案1】：

我认为使用窗口函数是完美的解决方案，因为您不必重新加入数据框。

val window = Window.partitionBy("uid").orderBy("year")

df.withColumn("count", count("uid").over(window))
  .filter($"count" > 2).drop("count").show

输出：

+---+----+-----+-----+
|uid|year|month|count|
+---+----+-----+-----+
|  1|2017|   03|    2|
|  1|2017|   05|    2|
|  2|2017|   01|    1|
+---+----+-----+-----+