与GROUPBY星火筛选数据统计

Question

一个数据帧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")) 

有没有更好的答案？

Answer 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| 
+---+----+-----+-----+