使用多个窗口过滤PySpark数据框的行

Question

我正在尝试根据时间戳[(start1, stop1), (start2, stop2), ...]的元组列表过滤Pyspark数据帧。 每个元组代表一个时间窗口。 Pyspark数据框具有以下结构：

+-------------------+------+
|                 ts|   var|
+-------------------+------+
|2018-09-01 20:10:00|     0|
|2018-09-01 20:12:00|     2|
|2018-09-01 20:13:00|     1|
|2018-09-01 20:17:00|     5|
+-------------------+------+

ts是时间戳列，而var是关注变量的列。 我正在寻找一种有效的方法来过滤掉不在一个时间窗口内的所有行。 例如，如果我的时间窗口列表包含单个窗口[(datetime(2018, 9, 1, 20, 11), datetime(2018, 9, 1, 20, 14))] 2018，9，1，1，20，11 [(datetime(2018, 9, 1, 20, 11), datetime(2018, 9, 1, 20, 14))] 2018，9，1，1，20，14 [(datetime(2018, 9, 1, 20, 11), datetime(2018, 9, 1, 20, 14))]则过滤后的数据帧应为

+-------------------+------+ 
|                 ts|   var| 
+-------------------+------+ 
|2018-09-01 20:12:00|     2| 
|2018-09-01 20:13:00|     1|
+-------------------+------+ 

我能够使用udf和for循环提供一个工作代码段，该代码段会在所有时间窗口内为每一行进行迭代（请参见下面的代码）。 但是，在所有时间范围内，每一行的循环都很慢。

一些其他信息：

• 提前不知道时间窗口的大小和数量，即不可能进行硬编码
• Pyspark数据框通常具有几百万行
• 时间窗口的数量通常在100-1000之间

如果有人可以指出一个更有效的解决方案，我将不胜感激。

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import BooleanType
import pandas as pd
from datetime import datetime

spark = SparkSession.builder.getOrCreate()

# create Pyspark dataframe
data = {'ts': [datetime(2018, 9, 1, 20, 10), datetime(2018, 9, 1, 20, 12),
               datetime(2018, 9, 1, 20, 13), datetime(2018, 9, 1, 20, 17)],
         'var': [0, 2, 1, 5]}
df = spark.createDataFrame(pd.DataFrame(data))

# list of windows [(start1, stop1), (start2, stop2), ...] for filtering
windows = [(datetime(2018, 9, 1, 20, 11), datetime(2018, 9, 1, 20, 14))]

# udf for filtering
def is_in_windows_udf(windows):
    def _is_in_windows(t, windows):
        for ts_l, ts_h in windows:
            if ts_l <= t <= ts_h:
                return True
            return False
    return udf(lambda t: _is_in_windows(t, windows), BooleanType())

# perform actual filtering operation
df.where(is_in_windows_udf(windows)(col("ts"))).show()

Answer 1

下面是一个更简单的解决方案，因为我们正在对同一数据集进行联合，因此它也将并行执行：

for count, item in enumerate(windows):
    if count == 0:
        result = df.filter(
            (F.col("ts")<= item[1]) &
            (F.col("ts")>= item[0])
        )
    else:
        result = result.union(
            df.filter(
            (F.col("ts")<= item[1]) &
            (F.col("ts")>= item[0])
            )
        )