python - 在 Spark 数据框中添加截止日期列

Question

我有一个场景，我想将月份添加到 spark DataFrame 中的日期列，它有两列数据类型为 (Date, Int)

例如

df.show()

data_date months_to_add
2015-06-23 5
2016-07-20 7

我想添加一个有新日期的新列(在现有日期上添加几个月后)，输出如下所示-

data_date month_to_add new_data_date
2015-06-23 5           2015-11-23
2016-07-20 1           2016-8-20

我试过下面的一段代码，但它似乎不起作用-

df = df.withColumn("new_data_date", a
  dd_months(col("data_date"), col("months_to_add")))

它给我错误-

'Column' object is not callable

如果有任何方法可以在不在数据框之上使用 SQL 查询的情况下实现这一点，请帮助我。

Answer 1

我会使用 expr:

from pyspark.sql.functions import expr

df = spark.createDataFrame(
    [("2015-06-23", 5), ("2016-07-20", 7)],
    ("data_date", "months_to_add")
).select(to_date("data_date").alias("data_date"), "months_to_add")

df.withColumn("new_data_date", expr("add_months(data_date, months_to_add)")).show()

+----------+-------------+-------------+
| data_date|months_to_add|new_data_date|
+----------+-------------+-------------+
|2015-06-23|            5|   2015-11-23|
|2016-07-20|            7|   2017-02-20|
+----------+-------------+-------------+