使用 pyspark 使用 interval_start 和 interval_end 日期分割日期

Question

我正在尝试在 pyspark 中构建一个用于拆分日期的脚本：以下是我的输入数据框

+--------------+------------+----+---+
|interval_start|interval_end|name|val|
+--------------+------------+----+---+
|2018-10-31    |2020-09-05  | abc|1  |
|2020-09-05    |2020-10-05  | abc|1  |
|2020-01-31    |2020-04-30  | def|2  |
+--------------+------------+----+---+

预期的输出数据帧如下：

+--------------+------------+----+---+
|interval_start|interval_end|name|val|
+--------------+------------+----+---+
|2018-10-31    |2018-11-30  | abc|1  |
|2018-11-30    |2018-12-31  | abc|1  |
|2018-12-31    |2019-01-31  | abc|1  |
|2019-01-31    |2019-02-28  | abc|1  |
|2019-02-28    |2019-03-31  | abc|1  |
|2019-03-31    |2019-04-30  | abc|1  |
|2019-04-30    |2019-05-31  | abc|1  |
|2019-05-31    |2019-06-30  | abc|1  |
|2019-06-30    |2019-07-31  | abc|1  |
|2019-07-31    |2019-08-31  | abc|1  |
|2019-08-31    |2019-09-30  | abc|1  |
|2019-09-30    |2019-10-31  | abc|1  |
|2019-10-31    |2019-11-30  | abc|1  |
|2019-11-30    |2019-12-31  | abc|1  |
|2019-12-31    |2020-01-31  | abc|1  |
|2020-01-31    |2020-02-29  | abc|1  |
|2020-02-29    |2020-03-31  | abc|1  |
|2020-03-31    |2020-04-30  | abc|1  |
|2020-04-30    |2020-05-31  | abc|1  |
|2020-05-31    |2020-06-30  | abc|1  |
|2020-06-30    |2020-07-31  | abc|1  |
|2020-07-31    |2020-08-31  | abc|1  |
|2020-08-31    |2020-09-05  | abc|1  |
|2020-09-05    |2020-10-05  | abc|1  |
|2020-01-31    |2020-02-29  | def|2  |
|2020-02-29    |2020-03-31  | def|2  |
|2020-03-31    |2020-04-30  | def|2  |
+--------------+------------+----+---+

有什么方法可以实现吗？

我在 for 循环中使用了以下代码，但这不是拆分数据帧的最佳解决方案。代码也有问题：

df = spark.sql(f"""
    select
        explode(
            arrays_zip(
                sequence(date('{interval_start}'), date('{interval_end}') - interval 1 month, interval 1 month),
                sequence(date('{interval_start}') + interval 1 month, date('{interval_end}'), interval 1 month)
           )
       )
""").selectExpr('col.*').toDF('interval_start', 'interval_end')

如果我通过以下日期

interval_start = '2018-10-31'
interval_end = '2020-09-05'

它将按如下方式拆分，但我们缺少从 2020-08-31 到 2020-09-05

的日期

+--------------+------------+
|interval_start|interval_end|
+--------------+------------+
|2018-10-31    |2018-11-30  |
|2018-11-30    |2018-12-31  |
|2018-12-31    |2019-01-31  |
|2019-01-31    |2019-02-28  |
|2019-02-28    |2019-03-31  |
|2019-03-31    |2019-04-30  |
|2019-04-30    |2019-05-31  |
|2019-05-31    |2019-06-30  |
|2019-06-30    |2019-07-31  |
|2019-07-31    |2019-08-31  |
|2019-08-31    |2019-09-30  |
|2019-09-30    |2019-10-31  |
|2019-10-31    |2019-11-30  |
|2019-11-30    |2019-12-31  |
|2019-12-31    |2020-01-31  |
|2020-01-31    |2020-02-29  |
|2020-02-29    |2020-03-31  |
|2020-03-31    |2020-04-30  |
|2020-04-30    |2020-05-31  |
|2020-05-31    |2020-06-30  |
|2020-06-30    |2020-07-31  |
|2020-07-31    |2020-08-31  |
+--------------+------------+

有人可以帮我们解决这个问题吗？

Answer 1

import pyspark.sql.functions as f
from pyspark.sql.functions import desc, row_number, monotonically_increasing_id
from pyspark.sql.window import Windowenter code here

interval_start = '2018-10-31'
interval_end = '2020-09-05'

data = [(interval_start, interval_end)]
df = spark.createDataFrame(data, ["interval_start", "interval_end"])

df.show()

+--------------+------------+
|interval_start|interval_end|
+--------------+------------+
|    2018-10-31|  2020-09-05|
+--------------+------------+

df1 = df.withColumn("monthsDiff", f.months_between("interval_end", "interval_start"))\
    .withColumn("repeat", f.expr("split(repeat(',', monthsDiff), ',')"))\
    .select("*", f.posexplode("repeat").alias("date", "val"))\
    .withColumn("start_date", f.expr("add_months(interval_start, date)"))\
    .select('start_date')\
    .withColumn('index', row_number().over(Window.orderBy(monotonically_increasing_id())))

df1.show()
+----------+-----+
|start_date|index|
+----------+-----+
|2018-10-31|    1|
|2018-11-30|    2|
|2018-12-31|    3|
|2019-01-31|    4|
|2019-02-28|    5|
|2019-03-31|    6|
|2019-04-30|    7|
|2019-05-31|    8|
|2019-06-30|    9|
|2019-07-31|   10|
|2019-08-31|   11|
|2019-09-30|   12|
|2019-10-31|   13|
|2019-11-30|   14|
|2019-12-31|   15|
|2020-01-31|   16|
|2020-02-29|   17|
|2020-03-31|   18|
|2020-04-30|   19|
|2020-05-31|   20|
|2020-06-30|   21|
|2020-07-31|   22|
|2020-08-31|   23|
+----------+-----+

df2 = df1.where(col('index') != 1).select(col("start_date").alias("end_date"), col("index"))
columns = ['end_date','index']
vals = [(interval_end, 0)]

newRow = spark.createDataFrame(vals, columns)
df2 = df2.union(newRow).drop("index").withColumn('index', row_number().over(Window.orderBy(monotonically_increasing_id())))

df2.show()
+----------+-----+
|  end_date|index|
+----------+-----+
|2018-11-30|    1|
|2018-12-31|    2|
|2019-01-31|    3|
|2019-02-28|    4|
|2019-03-31|    5|
|2019-04-30|    6|
|2019-05-31|    7|
|2019-06-30|    8|
|2019-07-31|    9|
|2019-08-31|   10|
|2019-09-30|   11|
|2019-10-31|   12|
|2019-11-30|   13|
|2019-12-31|   14|
|2020-01-31|   15|
|2020-02-29|   16|
|2020-03-31|   17|
|2020-04-30|   18|
|2020-05-31|   19|
|2020-06-30|   20|
|2020-07-31|   21|
|2020-08-31|   22|
|2020-09-05|   23|
+----------+-----+

result_df = df1.join(df2, ("index")).drop("index")
result_df.show()

+----------+----------+
|start_date|  end_date|
+----------+----------+
|2018-10-31|2018-11-30|
|2018-11-30|2018-12-31|
|2018-12-31|2019-01-31|
|2019-01-31|2019-02-28|
|2019-02-28|2019-03-31|
|2019-03-31|2019-04-30|
|2019-04-30|2019-05-31|
|2019-05-31|2019-06-30|
|2019-06-30|2019-07-31|
|2019-07-31|2019-08-31|
|2019-08-31|2019-09-30|
|2019-09-30|2019-10-31|
|2019-10-31|2019-11-30|
|2019-11-30|2019-12-31|
|2019-12-31|2020-01-31|
|2020-01-31|2020-02-29|
|2020-02-29|2020-03-31|
|2020-03-31|2020-04-30|
|2020-04-30|2020-05-31|
|2020-05-31|2020-06-30|
|2020-06-30|2020-07-31|
|2020-07-31|2020-08-31|
|2020-08-31|2020-09-05|
+----------+----------+