基于列值向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中而不将date列转换回string来创建startDate列？

所以与add_months函数基本相同

def add_months(startDate: Column, numMonths: Int)

而是传递一列而不是文字。

Answer 1

您可以使用UDF (User Defined Functions)来实现这一点。下面我创建了myUDF函数，它添加月份日期并以字符串格式返回结果日期，我将使用这个UDF创建一个新列，使用withColumn on 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|
+----------+-----------+----------+