根据条件修改spark日期列中的日期（月）

Question

我想修改 spark df 中的日期列，以便仅在出现某些月份时减去 1 个月。 IE。仅当日期为 yyyy-07-31 或日期为 yyyy-04-30 时，才分别将其更改为 yyyy-06-31 和 yyyy-03-30。任何想法如何使用 pyspark 函数来做到这一点？

|DateColumn|    
|2016-04-30|
|2015-04-30|
|2017-09-03|
|2017-07-31|
|2016-09-01|
|2018-07-31|
     ...

预期结果：

|DateColumn|    
|2016-03-30| <- changed
|2015-03-30| <- changed
|2017-09-03|
|2017-06-31| <- changed
|2016-09-01|
|2018-06-31| <- changed
     ...

Answer 1

我建议使用 functions 模块，然后结合几个功能：

.when() 然后 otherwise()
.month()
.date_format()
.add_months(date, -1)

例如，它可以归结为：

import pyspark.sql.functions as F

df = spark.createDataFrame([{'date': '2022-04-15'}, {'date': '2022-05-17'}])

df \
    .withColumn('new_date', 
                F.when(F.month(F.col('date')).isin([4, 7]), 
                       F.add_months(F.date_format('date', 'yyyy-MM-dd'), -1))
                .otherwise(F.col('date'))) \
    .show()

然后你会得到：

+----------+----------+
|      date|  new_date|
+----------+----------+
|2022-04-15|2022-03-15|
|2022-05-17|2022-05-17|
+----------+----------+

Update （事实证明这是一个子字符串问题，而不是一般的月份减法，详见评论）

实际上，您可以堆叠 .when 功能以在 Python 中执行某种 if-elif-else。下面的代码显示了如何在您分享的场景中执行它：

from pyspark.sql.functions import when

df \
    .withColumn('new_date', 
                 when(F.substring(F.col("date"), -5, 5) == '01-31', F.concat(F.year(F.col("date")), F.lit('-12-31')))
                .when(F.substring(F.col("date"), -5, 5) == '04-30', F.concat(F.year(F.col("date")), F.lit('-03-30')))
                .when(F.substring(F.col("date"), -5, 5) == '07-31', F.concat(F.year(F.col("date")), F.lit('-06-30')))
                .when(F.substring(F.col("date"), -5, 5) == '10-31', F.concat(F.year(F.col("date")), F.lit('-09-30')))
                .otherwise(F.col('date'))) \
    .show()

根据条件修改spark日期列中的日期（月）

Modify date (month) in spark date column based on condition

apache-spark

apache-spark-sql

pyspark