通过将日期列与两个参考列进行比较来计算从日期开始的期间

Question

我正在为我参与的一个项目做一些数据准备。我们在Databricks中完成了大部分工作，使用底层的Apache Spark在大型数据集上进行计算。一切都在PySpark中完成。

我的目标是将日期变量转换为变量yearperiod，它将一年划分为13个周期，每个周期为4周(有些例外(。该值是年份和周期的串联，例如yearperiod = 201513将是2015年，周期13。

我有两个表：yp_table，它包含年份的开始和结束日期(编辑：键入DateType()((从2012年到现在，编辑：约120行(：

+----------+----------+----------+
|     start|       end|yearperiod|
+----------+----------+----------+
|2012-01-16|2012-01-29|    201201|
|2012-01-30|2012-02-26|    201202|
|2012-02-27|2012-03-25|    201203|
|2012-03-26|2012-04-22|    201204|
|2012-04-23|2012-05-20|    201205|
|2012-05-21|2012-06-17|    201206|
....

我有一个实际的data表，它包含一个Date列(编辑：键入StringType()(：

+--------+--------+--------+-----+
|    Var1|    Var2|    Date| Var3|
+--------+--------+--------+-----+
|  xxxxxx|    xxxx|20191231| x,xx|
|  xxxxxx|    xxxx|20191231| x,xx|
|  xxxxxx|    xxxx|20191231| x,xx|
|  xxxxxx|    xxxx|20200101| x,xx|
|  xxxxxx|    xxxx|20200101| x,xx|
|  xxxxxx|    xxxx|20200101| x,xx|
|  xxxxxx|    xxxx|20200101| x,xx|
...

我的问题是：如何通过将data.Date与yp_table.start和yp_table.end进行比较来计算data表的列yearperiod？

到目前为止，我已经能够使它与常规Python(一种具有列表理解的解决方案(一起工作，但事实证明，对于大型数据集来说，它太慢了。非常感谢您的帮助！

编辑：出于隐私原因，我无法给出数据帧的实际模式。我在上面进行了编辑，以包含相关列的类型。

Answer 1

在datadf中添加一列，该列包含与yp_table匹配格式的日期，然后按日期间隔进行过滤。由于yp_table很小，您可以使用广播联接来加快速度。

import pyspark.sql.functions as fun
# Date lookup
start_dates = ["2012-01-16", "2012-01-30", "2012-02-27", "2012-03-26", "2012-04-23", "2012-05-21"]
end_dates = ["2012-01-29", "2012-02-26", "2012-03-25", "2012-04-22", "2012-05-20", "2012-06-17"]
yearperiod = ["201201", "201202", "201203", "201204", "201205", "201206"]
yp_table = spark.createDataFrame(pd.DataFrame({'start': start_dates, 'end': end_dates, 'yearperiod': yearperiod}))
# Data df
dates =  ["20120116", "20120130", "20120228", "20120301", "20200101", "20200101", "20200101"]
vals = range(0, len(dates))
data = spark.createDataFrame(pd.DataFrame({'Dates':dates, 'vals': vals}))
# Add formatted data_str column for joining
data = data.withColumn("date_str", fun.concat_ws("-", data.Dates.substr(0,4), data.Dates.substr(5,2), data.Dates.substr(7,2))) # + "-" + data.Dates.substr(6,8))
# Broadcase join small yp_table into the data table using conditional
joined = data.join(fun.broadcast(yp_table), (data.date_str >= yp_table.start) & (data.date_str < yp_table.end))

yp_table.show()
data.show()
joined.show()
+----------+----------+----------+
|     start|       end|yearperiod|
+----------+----------+----------+
|2012-01-16|2012-01-29|    201201|
|2012-01-30|2012-02-26|    201202|
|2012-02-27|2012-03-25|    201203|
|2012-03-26|2012-04-22|    201204|
|2012-04-23|2012-05-20|    201205|
|2012-05-21|2012-06-17|    201206|
+----------+----------+----------+
+--------+----+----------+
|   Dates|vals|  date_str|
+--------+----+----------+
|20120116|   0|2012-01-16|
|20120130|   1|2012-01-30|
|20120228|   2|2012-02-28|
|20120301|   3|2012-03-01|
|20200101|   4|2020-01-01|
|20200101|   5|2020-01-01|
|20200101|   6|2020-01-01|
+--------+----+----------+
+--------+----+----------+----------+----------+----------+
|   Dates|vals|  date_str|     start|       end|yearperiod|
+--------+----+----------+----------+----------+----------+
|20120116|   0|2012-01-16|2012-01-16|2012-01-29|    201201|
|20120130|   1|2012-01-30|2012-01-30|2012-02-26|    201202|
|20120228|   2|2012-02-28|2012-02-27|2012-03-25|    201203|
|20120301|   3|2012-03-01|2012-02-27|2012-03-25|    201203|
+--------+----+----------+----------+----------+----------+