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

时间:2016-03-17 03:45:41

标签: python apache-spark pyspark apache-spark-sql user-defined-functions

给出以下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就可以做到这一点。

4 个答案:

答案 0 :(得分:6)

我能够使用selectExpr来解决这个问题。

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

如果要将列附加到原始DF,只需将*添加到表达式

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

答案 1 :(得分:3)

通过使用 withColumn 功能,我们可以执行 date_sub 功能

>>> df.withColumn('substracted_dates',date_sub('date_col','day_col'))

答案 2 :(得分:1)

不是最优雅的解决方案,但是如果你不想在Scala中破解SQL表达式(不是它应该很难,但这些是私有的sql)这样的东西应该做的伎俩:

from pyspark.sql import Column

def date_sub_(c1: Column, c2: Column) -> Column:
    return ((c1.cast("timestamp").cast("long") - 60 * 60 * 24 * c2)
        .cast("timestamp").cast("date"))

对于Python 2.x,只需删除类型注释。

答案 3 :(得分:0)

格式略有不同,但也有效:

df.registerTempTable("dfTbl")

newdf = spark.sql("""
                     SELECT *, date_sub(d.date_col, d.day_col) AS DateSub 
                     FROM dfTbl d
                   """)