具有数组的列，其中包含自开始日期起x年内的所有月份-Pyspark

Question

假设您有一个数据帧df，如下所示：

ID  Years     Date
A    5    2021-02-01
B    3    2021-02-01
C    6    2021-02-01

我希望能够创建一个额外的日期数组列，所有日期从最初的日期+1个月一直到年份列中的x年数。它看起来如下：

ID  Years     Date        Dates
A    5    2021-02-01     [2021-03-01,2021-04-01,...,2026-02-01]
B    3    2021-03-01     [2021-04-01,2021-04-01,...,2024-03-01]
C    6    2021-02-01     [2021-03-01,2021-04-01,...,2027-02-01]

Answer 1

对于火花>=2.4，您可以使用sequence和add_months函数来生成所需的日期序列。

df = df.withColumn('Dates',
F.expr('sequence(add_months(to_date(Date), 1), add_months(to_date(Date), int(Years) * 12), interval 1 month)')
)
df.show(truncate=False)

Answer 2

我不是PySpark的专家，我听说在某些情况下，将PySpark更改为pandas数据帧可能并不有趣。但如果可以的话，您可以将格式更改为panda，并尝试使用apply函数：

df = df_spark.toPandas()
def getRangeDate(row):
return list(map(lambda x: x.strftime("%Y-%m-%d"), list(pd.date_range(start = row["Date"], periods = 12*row["Years"]+1, freq = 'MS'))[1:]))
df['Dates'] = df.apply(getRangeDate, axis=1)
df

对于您的示例输入，它具有以下输出：

ID  Years   Date    Dates
0   A   5   2021-02-01  [2021-03-01, 2021-04-01, 2021-05-01, ..., 2026-02-01]
1   B   3   2021-03-01  [2021-04-01, 2021-05-01, 2021-06-01, ..., 2024-03-01]
2   C   6   2021-02-01  [2021-03-01, 2021-04-01, 2021-05-01, ..., 2027-02-01]