PySpark中具有多列的日期算法

Question

我正在嘗試使用PySpark數據框中的多個列進行一些中等復雜的日期算術。 基本上，我有一個名為number的列，表示我需要過濾的created_at時間戳之后的周數。 在PostgreSQL中，你可以根據列中的值乘以一個間隔 ，但我似乎無法弄清楚如何使用SQL API或Python API在PySpark中執行此操作。 這里的任何幫助將不勝感激！

import datetime
from pyspark.sql import SQLContext
from pyspark.sql import Row
from pyspark import SparkContext

sc = SparkContext()
sqlContext = SQLContext(sc)
start_date = datetime.date(2020,1,1)

my_df = sc.parallelize([
        Row(id=1, created_at=datetime.datetime(2020, 1, 1), number=1,  metric=10),
        Row(id=1, created_at=datetime.datetime(2020, 1, 1), number=2,  metric=10),
        Row(id=1, created_at=datetime.datetime(2020, 1, 1), number=3,  metric=10),
        Row(id=2, created_at=datetime.datetime(2020, 1, 15), number=1,  metric=20),
        Row(id=2, created_at=datetime.datetime(2020, 1, 15), number=2,  metric=20),
        Row(id=3, created_at=datetime.datetime(2020, 7, 1), number=7,  metric=30),
        Row(id=3, created_at=datetime.datetime(2020, 7, 1), number=8,  metric=30),
        Row(id=3, created_at=datetime.datetime(2020, 7, 1), number=9,  metric=30),
        Row(id=3, created_at=datetime.datetime(2020, 7, 1), number=10, metric=30),
    ]).toDF()


# This doesn't work!
new_df = my_df.where("created_at + interval 7 days * number > '" + start_date.strftime("%Y-%m-%d") +"'")
# Neither does this!
new_df = my_df.filter(my_df.created_at + datetime.timedelta(days=my_df.number * 7)).date() > start_date.date()

這里有一個可能的解決方案 ，需要將日期轉換為字符串，使用python中的datetime庫將字符串轉換為datetime對象，然后執行操作，但這看起來很瘋狂。

Answer 1

好吧，我想出了使用expr和內置date_add函數的前進方法。

from pyspark.sql.functions import expr, date_add
new_df = my_df.withColumn('test', expr('date_add(created_at, number*7)'))
filtered = new_df.filter(new_df.test > start_date)
filtered.show()

但是，如果其他人想要添加，我會非常喜歡一般性地了解其工作方式/原因。