從 pyspark 中的 dataframe 創建歷史數據

Question

我有一個 dataframe 如下：

日期	一些數量
...	...
2021-01-01	4
2021-01-02	1
2021-01-03	6
2021-01-04	2
2021-01-05	2
2021-01-06	8
2021-01-07	9
2021-01-08	1
...	...

我想為每個日歷日創建歷史數據，並在最后一步進行一些聚合。 中間 dataframe 應如下所示：

日歷日期	日期	一些數量
...	...	...
2021-01-03	2021-01-01	4
2021-01-03	2021-01-02	1
2021-01-04	...	...
2021-01-04	2021-01-01	4
2021-01-04	2021-01-02	1
2021-01-04	2021-01-03	6
2021-01-05	...	...
2021-01-05	2021-01-01	4
2021-01-05	2021-01-02	1
2021-01-05	2021-01-03	6
2021-01-05	2021-01-04	2
2021-01-06	...	...
2021-01-06	2021-01-01	4
2021-01-06	2021-01-02	1
2021-01-06	2021-01-03	6
2021-01-06	2021-01-04	2
2021-01-06	2021-01-05	2
2021-01-06	...	...

有了這個 dataframe，日歷日期上的任何聚合都很容易（例如，指出當天之前售出的數量、平均 7 天、平均 30 天等）。

我嘗試運行日歷日期的 for 循環：

for i, date in enumerate(pd.data_range("2021-01-01","2021-03-01"):

   df_output = []

   df_transformed = df.where(F.col("date") < date)
   df_transformed = df_transformed.withColumn("calendar_date", date)

   if i == 0:
      df_output = df_transformed
   else:
      df_output = df_output.union(df_transformed)

但是，這是非常低效的，並且一旦我開始包含更多列，它就會崩潰。

是否可以創建帶有日歷日期的 dataframe 並進行重新創建 dataframe 我期望的連接？

謝謝你的幫助。

Answer 1

您可以簡單地將日歷 dataframe 與您的主要 dataframe 連接條件“小於”：

# Let's call your main dataframe as `df`

# Extracting first and last date
_, min_date, max_date = (df
    .groupBy(F.lit(1))
    .agg(
        F.min('date').alias('min_date'),
        F.max('date').alias('max_date'),
    )
    .first()
)

# Then create a temporary dataframe to hold all calendar dates
dates = [{'calendar_date': str(d.date())} for d in pd.date_range(min_date, max_date)]
calendar_df = spark.createDataFrame(dates)
calendar_df.show(10, False)
# +-------------+
# |calendar_date|
# +-------------+
# |2021-01-01   |
# |2021-01-02   |
# |2021-01-03   |
# |2021-01-04   |
# |2021-01-05   |
# |2021-01-06   |
# |2021-01-07   |
# |2021-01-08   |
# +-------------+

# Now you can join to build your expected dataframe, note the join condition
(calendar_df
    .join(df, on=[calendar_df.calendar_date > df.date])
    .show()
)
# +-------------+----------+---+
# |calendar_date|      date|qty|
# +-------------+----------+---+
# |   2021-01-02|2021-01-01|  4|
# |   2021-01-03|2021-01-01|  4|
# |   2021-01-03|2021-01-02|  1|
# |   2021-01-04|2021-01-01|  4|
# |   2021-01-04|2021-01-02|  1|
# |   2021-01-04|2021-01-03|  6|
# |   2021-01-05|2021-01-01|  4|
# |   2021-01-05|2021-01-02|  1|
# |   2021-01-05|2021-01-03|  6|
# |   2021-01-05|2021-01-04|  2|
# |   2021-01-06|2021-01-01|  4|
# |   2021-01-06|2021-01-02|  1|
# |   2021-01-06|2021-01-03|  6|
# |   2021-01-06|2021-01-04|  2|
# |   2021-01-06|2021-01-05|  2|
# |   2021-01-07|2021-01-01|  4|
# |   2021-01-07|2021-01-02|  1|
# |   2021-01-07|2021-01-03|  6|
# |   2021-01-07|2021-01-04|  2|
# |   2021-01-07|2021-01-05|  2|
# +-------------+----------+---+
# only showing top 20 rows