scala - 如何在Spark中使用窗口函数过滤数据

我有以下数据:

rowid uid time code
   1  1      5    a
   2  1      6    b
   3  1      7    c
   4  2      8    a
   5  2      9    c
   6  2      9    c
   7  2     10    c
   8  2     11    a
   9  2     12    c

现在，我想以一种可以删除第6行和第7行的方式对数据进行过滤，因为对于特定的uid，我只想在代码中保留一行带有值'c'的行。

因此，预期数据应为:

rowid uid time code
   1  1      5    a
   2  1      6    b
   3  1      7    c
   4  2      8    a
   5  2      9    c
   8  2     11    a
   9  2     12    c

我正在使用像这样的窗口函数:

val window = Window.partitionBy("uid").orderBy("time")
val change = ((lag("code", 1).over(window) <=> "c")).cast("int")

这将帮助我们用代码“c”标识每一行。我可以扩展它以过滤出行以获取预期数据

最佳答案

如果只想删除code =“c”的行(每个uid的第一个除外)，则可以尝试以下操作:

val window = Window.partitionBy("uid", "code").orderBy("time")
val result = df
  .withColumn("rank", row_number().over(window))
  .where(
    (col("code") !== "c") ||
    col("rank") === 1
  )
  .drop("rank")

根据新信息进行编辑:

val window = Window.partitionBy("uid").orderBy("time")
val result = df
  .withColumn("lagValue", coalesce(lag(col("code"), 1).over(window), lit("")))
  .where(
    (col("code") !== "c") ||
    (col("lagValue") !== "c")
  )
  .drop("lagValue")