使用多个窗口过滤 PySpark 数据帧的行答案

【问题标题】：Filtering rows of a PySpark dataframe using multiple windows使用多个窗口过滤 PySpark 数据帧的行
【发布时间】：2019-02-07 17:55:08
【问题描述】：

我正在尝试根据时间戳元组列表过滤 Pyspark 数据帧[(start1, stop1), (start2, stop2), ...]。每个元组代表一个时间窗口。 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))] 组成，则过滤后的数据框应该是

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

我能够使用一个 udf 和一个 for 循环来设计一个工作代码 sn-p，该循环在所有时间窗口中对每一行进行迭代（参见下面的代码）。然而，在所有时间窗口上循环每一行是很慢的。

一些附加信息：

事先不知道时间窗口的大小和数量，即不可能进行硬编码
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()

【问题讨论】：

标签： python apache-spark pyspark apache-spark-sql user-defined-functions

【解决方案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])
            )
        )

【讨论】：