Python 使用多个窗口筛选PySpark数据帧的行

我试图根据时间戳的元组列表过滤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循环生成一个工作代码段，该循环在所有时间窗口内对每一行进行迭代（请参见下面的代码）。但是，在所有时间窗口中每行的循环都很慢

一些补充资料：

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

从pyspark.sql导入SparkSession
从pyspark.sql.functions导入udf，col
从pyspark.sql.types导入BooleanType
作为pd进口熊猫
从日期时间导入日期时间
spark=SparkSession.builder.getOrCreate（）
#创建Pyspark数据帧
数据={'ts'：[datetime（2018,9,1,20,10），datetime（2018,9,1,20,12），
日期时间（2018,9,1,20,13），日期时间（2018,9,1,20,17）]，
'var'：[0,2,1,5]}
df=spark.createDataFrame（pd.DataFrame（数据））
#用于筛选的窗口列表[（开始1，停止1），（开始2，停止2），…]
windows=[（日期时间（2018,9,1,20,11），日期时间（2018,9,1,20,14））]
#用于过滤的自定义项
def是windows中的自定义项（windows）：
def_在_窗口中（t，窗口）：
对于windows中的ts_l和ts_h：
如果ts_l一个更简单的解决方案可以是下面的解决方案，因为我们在同一数据集上进行联合，所以它也将并行执行：

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])
            )
        )

对于计数，枚举中的项（窗口）：
如果计数=0：
结果=df.filter(
（F.col（“ts”）=第[0]项）
)
其他：
result=result.union(
测向滤波器(
（F.col（“ts”）=第[0]项）
)
)