Добавление месяц DateType на основе значения столбца

Question

Предполагая dataframe с столбца даты и столбца Int, представляющий число месяцев:

val df = Seq(("2011-11-11",1),("2010-11-11",3),("2012-11-11",5)) 
       .toDF("startDate","monthsToAdd") 
       .withColumn("startDate",'startDate.cast(DateType)) 

+----------+-----------+ 
| startDate|monthsToAdd| 
+----------+-----------+ 
|2011-11-11|   1| 
|2010-11-11|   3| 
|2012-11-11|   5| 
+----------+-----------+ 

есть способ создания столбца EndDate путем добавления месяцев StartDate без преобразования столбца даты в строку?

Так в основном такие же как add_months функционировать

def add_months(startDate: Column, numMonths: Int) 

но проходя колонку вместо буквальным.

Answer 1

Вы можете использовать UDF (User Defined Functions) для этого. Ниже я есть создать myUDF функцию, добавить месяцы на сегодняшний день и возвращает дату результата в формате Строка и я буду использовать эту UDF, чтобы создать новый столбец с помощью withColumn на DataFrame

import java.text.SimpleDateFormat 
import java.util.Calendar 
import javax.xml.bind.DatatypeConverter 

import org.apache.spark.sql.functions._ 
import sparkSession.sqlContext.implicits._ 

val df = Seq(("2011-11-11",1),("2010-11-11",3),("2012-11-11",5)).toDF("startDate","monthsToAdd") 

val myUDF = udf { 
    val simpleDateFormat = new SimpleDateFormat("yyyy-MM-dd") 
    (startDate: String, monthValue: Int) => { 
    val calendar = DatatypeConverter.parseDateTime(startDate) 
    calendar.add(Calendar.MONTH, monthValue) 
    simpleDateFormat.format(calendar.getTime) 
    } 
} 

val newDf = df.withColumn("endDate", myUDF(df("startDate"), df("monthsToAdd"))) 
newDf.show() 

Выход:

+----------+-----------+----------+ 
| startDate|monthsToAdd| endDate| 
+----------+-----------+----------+ 
|2011-11-11|   1|2011-12-11| 
|2010-11-11|   3|2011-02-11| 
|2012-11-11|   5|2013-04-11| 
+----------+-----------+----------+