Как найти, если какой-либо элемент массива находится в пределах диапазона в pyspark

Question

У меня есть столбец массивов в кадре данных, и я хочу знать, находится ли какой-либо из элементов массива в определенном диапазоне. Пример:
input:

+------------------------------------------------------------------------------------------+
|dateTimeValue                                                                             |
+------------------------------------------------------------------------------------------+
|[2019-11-11T20:08:47.453+0000, 2020-10-15T20:08:47.453+0000, 2021-09-19T20:08:47.453+0000]|
|[2017-11-05T20:08:47.453+0000, 2020-05-05T20:08:47.453+0000, 2021-11-11T20:08:47.453+0000]|
+------------------------------------------------------------------------------------------+

Интересующий диапазон дат - 8 августа 2018 г. и 8 декабря 2019 г.
output:

+------------------------------------------------------------------------------------------+------------+
|dateTimeValue                                                                             |includedFlag|
+------------------------------------------------------------------------------------------+------------+
|[2019-11-11T20:08:47.453+0000, 2020-10-15T20:08:47.453+0000, 2021-09-19T20:08:47.453+0000]|True        |
|[2017-11-05T20:08:47.453+0000, 2020-05-05T20:08:47.453+0000, 2021-11-11T20:08:47.453+0000]|False       |
+------------------------------------------------------------------------------------------+------------+

Схема моегофрейм данных:

root 
|-- dateTimeValue: array (nullable = true) | 
    |-- element: timestamp (containsNull = true)

Ввод может быть сгенерирован с помощью:

import datetime 
df = spark.createDataFrame([([datetime.datetime(2019,11,11,20,8,47), datetime.datetime(2020,10,15,20,8,47), datetime.datetime(2021,9,19,20,8,47)],), ([datetime.datetime(2017,11,5,20,8,47), datetime.datetime(2020,5,5,20,8,47), datetime.datetime(2021,11,11,20,8,47)],)], ['dateTimeValue'])

Спасибо.

Answer 1

С помощью взорваться вы можете создать новую строку для каждого элемента вашего массива и сравнить каждый отдельный элемент с границами вашего диапазона.

import datetime 
import pyspark.sql.functions as F
df = spark.createDataFrame([([datetime.datetime(2019,11,11,20,8,47), datetime.datetime(2020,10,15,20,8,47), datetime.datetime(2021,9,19,20,8,47)],), ([datetime.datetime(2017,11,5,20,8,47), datetime.datetime(2020,5,5,20,8,47), datetime.datetime(2021,11,11,20,8,47)],)], ['dateTimeValue'])
df.show(truncate=False)

df= df.withColumn('ex', F.explode('dateTimeValue'))
df= df.withColumn('includedFlag', F.when((F.col('ex') < datetime.datetime(2019,12,8,00,00,00)) & (F.col("ex") > datetime.datetime(2018,8,8,00,00,00)) , 1).otherwise(0))
df.groupby('dateTimeValue').agg(F.max('includedFlag').alias('includedFlag')).show(truncate=False)

Вывод:

+---------------------------------------------------------------+
|dateTimeValue                                                  |
+---------------------------------------------------------------+
|[2019-11-11 20:08:47, 2020-10-15 20:08:47, 2021-09-19 20:08:47]|
|[2017-11-05 20:08:47, 2020-05-05 20:08:47, 2021-11-11 20:08:47]|
+---------------------------------------------------------------+

+---------------------------------------------------------------+------------+
|dateTimeValue                                                  |includedFlag|
+---------------------------------------------------------------+------------+
|[2017-11-05 20:08:47, 2020-05-05 20:08:47, 2021-11-11 20:08:47]|0           |
|[2019-11-11 20:08:47, 2020-10-15 20:08:47, 2021-09-19 20:08:47]|1           |
+---------------------------------------------------------------+------------+