根据列值将月份添加到DateType

时间:2017-02-11 09:39:43

标签: apache-spark apache-spark-sql

假设数据框的日期列和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|
+----------+-----------+

有没有办法通过将月份添加到startDate而不将日期列转换回字符串来创建endDate列?

所以与add_months函数基本相同

def add_months(startDate: Column, numMonths: Int)

但是传递列而不是文字。

1 个答案:

答案 0 :(得分:1)

您可以使用 UDF (User Defined Functions) 来实现这一目标。下面我创建了myUDF函数,它添加了迄今为止的月份并以String格式返回结果日期,我将使用此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|
+----------+-----------+----------+