[英]Time filtering a struct column in Pyspark dataframe
我有一个 dataframe 的列,其中的结构具有日期和值,因此架构看起来像
root
|-- col1: struct (nullable = true)
| |-- dates: array (nullable = true)
| | |-- element: timestamp (containsNull = true)
| |-- values: array (nullable = true)
| | |-- element: double (containsNull = true)
|-- col2: struct (nullable = true)
| |-- dates: array (nullable = true)
| | |-- element: timestamp (containsNull = true)
| |-- values: array (nullable = true)
| | |-- element: double (containsNull = true)
|-- id: string (nullable = true)
给定一些时间索引:
time_index = datetime.datetime(2015, 12, 12, 4, 45)
以及前后天数:
min_diff = -1 and max_diff = 2
我想要新的列col1_filt
和col2_filt
,它们具有相同的结构,返回落入由time_index
和min_diff
和max_diff
定义的 window 以及相应值的日期。 如果没有日期或值属于该 window 我希望它返回None
。
下面是一个可以使用的示例 DataFrame。
示例 DataFrame:
example_input = [
Row(
id = "A",
col1 = Row(
dates = [datetime.datetime(2015, 12, 11, 5, 28), datetime.datetime(2015, 12, 12, 4, 45), datetime.datetime(2015, 12, 13, 5, 9)],
values = [17.7, 19.1, 19.1]
),
col2 = Row(
dates = [datetime.datetime(2015, 12, 13, 4, 48), datetime.datetime(2015, 12, 15, 5, 8)],
values = [19.1, 19.1]
)
),
Row(
id = "B",
col1 = Row(
dates = [datetime.datetime(2017, 1, 13, 5, 9)],
values = [19.1]
),
col2 = Row(
dates = [datetime.datetime(2017, 1, 12, 2, 48), datetime.datetime(2017, 1, 15, 5, 8)],
values = [19.5, 29.1]
)
),
]
df = spark.createDataFrame(example_input)
显示df:
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+
|col1 |col2 |id |
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+
|[[2015-12-11 05:28:00, 2015-12-12 04:45:00, 2015-12-13 05:09:00], [17.7, 19.1, 19.1]]|[[2015-12-13 04:48:00, 2015-12-15 05:08:00], [19.1, 19.1]]|A |
|[[2017-01-13 05:09:00], [19.1]] |[[2017-01-12 02:48:00, 2017-01-15 05:08:00], [19.5, 29.1]]|B |
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+
我有一些代码将采用 Pyspark 行 object 并返回过滤后的 Pyspark 行 ZA8CFDE6331BD59EB62AC96F891 怎么做
下面是使用 UDF 进行过滤的示例:
import datetime
import pyspark.sql.functions as F
from pyspark.sql import Row
time_index = datetime.datetime(2015, 12, 12, 4, 45)
min_diff = -1
max_diff = 2
def time_filter(r):
ret = list(zip(*[
x for x in list(zip(r['dates'], r['values']))
if x[0] > time_index + datetime.timedelta(days=min_diff)
and x[0] < time_index + datetime.timedelta(days=max_diff)
]))
return Row(dates=ret[0], values=ret[1]) if len(ret) != 0 else None
time_filter_udf = F.udf(time_filter, 'struct<dates:array<timestamp>,values:array<double>>')
df2 = df.withColumn('col1_filt', time_filter_udf('col1')).withColumn('col2_filt', time_filter_udf('col2'))
df2.show(truncate=False)
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+-------------------------------------------------------------------------------------+-------------------------------+
|col1 |col2 |id |col1_filt |col2_filt |
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+-------------------------------------------------------------------------------------+-------------------------------+
|[[2015-12-11 05:28:00, 2015-12-12 04:45:00, 2015-12-13 05:09:00], [17.7, 19.1, 19.1]]|[[2015-12-13 04:48:00, 2015-12-15 05:08:00], [19.1, 19.1]]|A |[[2015-12-11 05:28:00, 2015-12-12 04:45:00, 2015-12-13 05:09:00], [17.7, 19.1, 19.1]]|[[2015-12-13 04:48:00], [19.1]]|
|[[2017-01-13 05:09:00], [19.1]] |[[2017-01-12 02:48:00, 2017-01-15 05:08:00], [19.5, 29.1]]|B |null |null |
+-------------------------------------------------------------------------------------+----------------------------------------------------------+---+-------------------------------------------------------------------------------------+-------------------------------+
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.