Question

我有一个这样的数据框：

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

我想保留y的'lead'值为'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

效率不高，但你可以用索引压缩，然后创建一个新的RDD，你在索引中加1，然后连接索引然后变成一个简单的过滤操作。

过滤前导/滞后是特定值的行（带过滤器的窗口）

2 个答案: