Come sottrarre una colonna di giorni da una colonna di date in Pyspark?

Question

dato i seguenti PySpark dataframe

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

Come può la colonna giorni sottratto dal colonna della data? In questo esempio, la colonna risultante dovrebbe essere ['2015-01-05', '2015-02-10'].

Ho esaminato pyspark.sql.functions.date_sub(), ma richiede una colonna di date e un solo giorno, ovvero date_sub(df['date_col'], 10). Idealmente, preferirei fare date_sub(df['date_col'], df['days_col']).

Ho provato anche la creazione di un'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']) 

funziona questo tecnicamente, ma ho letto che mettendosi tra Spark e Python può causare problemi di prestazioni per grandi insiemi di dati. Per ora posso restare con questa soluzione (non c'è bisogno di ottimizzare prematuramente), ma il mio istinto dice che ci deve essere un modo per fare questa cosa semplice senza usare una UDF Python.

Answer 1

Sono stato in grado di risolvere questo utilizzando selectExpr.

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

Se si desidera aggiungere la colonna all'originale DF, basta aggiungere * all'espressione

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

Answer 2

non la soluzione più elegante mai, ma se non si vuole incidere espressioni SQL a Scala (non che dovrebbe essere difficile, ma queste sono privato a sql) qualcosa di simile dovrebbe fare il trucco:

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")) 

Per Python 2.x basta rilasciare annotazioni di tipo.

Answer 3

formato leggermente diverso, ma funziona anche:

df.registerTempTable("dfTbl") 

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