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

Question

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

我的目标是将 date 变量转换为 variable yearperiod，它将一年分为 13 个 4 周的时期（有一些例外）。该值是年份和期间的串联，例如yearperiod = 2015132015 年，期间 13。

我有两个表：yp_table其中包含年份的开始日期和结束日期（编辑： type 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 列（编辑： type 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|
                 ...

我的问题：如何通过与和进行比较来计算表yearperiod的列？datadata.Dateyp_table.startyp_table.end

到目前为止，我已经能够使其与常规 Python（具有列表理解的解决方案）一起使用，但事实证明它对于大型数据集来说太慢了。任何帮助是极大的赞赏！

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

标签： pythonapache-sparkdatepyspark

Answer 1

向您的 df 添加一列data，其中包含匹配格式的日期，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|
+--------+----+----------+----------+----------+----------+