Как вычесть столбец дней из столбца дат в Pyspark?

учитывая следующий 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.

7
apache-spark pyspark python spark-dataframe udf

4 ответов

я смог решить эту проблему с помощью selectExpr.

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

если вы хотите добавить столбец к исходному DF, просто добавьте * выражение

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

не самое элегантное решение, но если вы не хотите взломать SQL-выражения в Scala (не то, что это должно быть сложно, но они являются частными для 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.х просто напишите аннотации типа.

1

С помощью withColumn функции, мы можем сделать date_sub функции

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

немного другой формат, но тоже работает:

df.registerTempTable("dfTbl")

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