包含从开始日期起 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

我不是 PySpark 方面的专家，而且我听说在某些情况下将 PySpark 更改为 pandas 数据框并不有趣。但如果没问题，您可以将格式更改为 pandas 并尝试使用 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]

Answer 2

对于 spark >= 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)

包含从开始日期起 x 年中所有月份的数组的列 - Pyspark

Column with array with all months in x amount of years from starting date - Pyspark

python

date

dataframe

pyspark