如何从Pyspark中的一列日期中减去一列天数？

Question

给出以下PySpark DataFrame

df = sqlContext.createDataFrame([('2015-01-15', 10),
                                 ('2015-02-15', 5)],
                                 ('date_col', 'days_col'))

如何从日期列中减去日期列？在此示例中,结果列应为['2015-01-05', '2015-02-10'].

我调查了pyspark.sql.functions.date_sub(),但它需要一个日期列和一天,即date_sub(df['date_col'], 10).理想情况下,我更愿意这样做date_sub(df['date_col'], df['days_col']).

我也试过创建一个UDF:

from datetime import timedelta
def subtract_date(start_date, days_to_subtract):
    return start_date - timedelta(days_to_subtract)

subtract_date_udf = udf(subtract_date, DateType())
df.withColumn('subtracted_dates', subtract_date_udf(df['date_col'], df['days_col'])

这在技术上有效,但我已经读过Spark和Python之间的步骤可能会导致大型数据集的性能问题.我现在可以坚持使用这个解决方案(不需要过早优化),但我的直觉说只需要一种方法来做这个简单的事情而不使用Python UDF.

Answer 1

使用expr函数（如果你有dynamic values从列减去）：

>>> from pyspark.sql.functions import *
>>> df.withColumn('substracted_dates',expr("date_sub(date_col,days_col)"))

使用withColumn函数（如果必须literal values减去）：

>>> df.withColumn('substracted_dates',date_sub('date_col',<int_literal_value>))

Answer 2

我能够使用解决此问题selectExpr。

df.selectExpr('date_sub(date_col, day_col) as subtracted_dates')

如果要将列追加到原始DF，只需将其添加*到表达式中

df.selectExpr('*', 'date_sub(date_col, day_col) as subtracted_dates')