Apache Spark - 筛选超前/滞后为特定值的行(带筛选器的窗口)

Question

我有一个这样的数据框：

  id x y
1  a 1 P
2  a 2 S
3  b 3 P
4  b 4 S

我想保留 y 的"领先"值为"S"的行，这样我的结果数据框将是：

      id     x      y
1      a     1      P
2      b     3      P

我能够使用 PySpark 执行以下操作：

getLeadPoint = udf(lambda x: 'S' if (y == 'S') else 'NOTS', StringType())
windowSpec = Window.partitionBy(df['id'])
df = df.withColumn('lead_point', getLeadPoint(lead(df.y).over(windowSpec)))
dfNew = df.filter(df.lead_point == 'S')

但是，在这里，我正在更改不必要的列，然后进行过滤。

我想做的是这样的事情，我使用 lead（） 过滤，但无法让它工作：

dfNew = df.filter(lead(df.y).over(windowSpec) == 'S')

关于如何使用窗口直接过滤器获得结果的任何想法？

R 等效值为：

library(dplyr)
df %>% group_by(id) %>% filter(lead(y) == 'S')

Answer 1

假设您的数据如下所示：

df = sc.parallelize([
    ("a", 1,  1, "P"), ("a", 2,  2, "S"),
    ("b", 4,  2, "S"), ("b", 3,  1, "P"), ("b", 2,  3, "P"), ("b", 3,  3, "S")
]).toDF(["id", "x", "timestamp", "y"])

和窗口规格等效于

from pyspark.sql.functions import lead, col
from pyspark.sql import Window
w = Window.partitionBy("id").orderBy("timestamp")

您可以简单地添加列并将其用于过滤：

(df
    .withColumn("lead_y", lead("y").over(w))
    .where(col("lead_y") == "S").drop("lead_y"))

它并不漂亮，但比UDF调用更有效。

Answer 2

效率

不高，但您可以使用index进行压缩，然后创建一个新的RDD，将1添加到索引中，然后连接索引，然后它变成一个简单的过滤器操作。