apache-spark - 编写采用 Column 参数并返回 Column 的 Spark 函数

Question

我想写一个yeardiff功能类似于 datediff . yeardiff应该带两个Column参数并返回一个列，其中包含这些参数列之间的年数。

让我们使用以下示例数据:

val testDf = Seq(
  ("2016-09-10", "2001-08-10"),
  ("2016-04-18", "2010-05-18"),
  ("2016-01-10", "2013-08-10")
)
  .toDF("first_datetime", "second_datetime")
  .withColumn("first_datetime", $"first_datetime".cast("timestamp"))
  .withColumn("second_datetime", $"second_datetime".cast("timestamp"))

我们可以运行它来获取日期差异:

testDf.withColumn("num_days", datediff(col("first_datetime"), col("second_datetime")))

我希望能够运行这个:

testDf.withColumn("num_years", yeardiff(col("first_datetime"), col("second_datetime")))

我试图定义一个 yeardiff具有必要的方法签名的函数并且没有得到任何地方:

def yeardiff(end: Column, start: Column): Column = {
  // what do I do here
}    

这是我想出但不喜欢的黑客转换解决方案:

def yearDiff(end: String, start: String)(df: DataFrame): DataFrame = {
  val c = s"${end}_${start}_datediff"
  df
    .withColumn(c, datediff(col(end), col(start)))
    .withColumn("yeardiff", col(c) / 365)
}

编辑

我开始深入研究 Spark 源代码以了解如何 datediff作品。这是 the datediff function definition :

def datediff(end: Column, start: Column): Column = withExpr { DateDiff(end.expr, start.expr) }

这是 the DateDiff case class :

case class DateDiff(endDate: Expression, startDate: Expression)
  extends BinaryExpression with ImplicitCastInputTypes {

  override def left: Expression = endDate
  override def right: Expression = startDate
  override def inputTypes: Seq[AbstractDataType] = Seq(DateType, DateType)
  override def dataType: DataType = IntegerType

  override def nullSafeEval(end: Any, start: Any): Any = {
    end.asInstanceOf[Int] - start.asInstanceOf[Int]
  }

  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    defineCodeGen(ctx, ev, (end, start) => s"$end - $start")
  }
}

Answer 1

这可能会解决您的问题:

def yearDiff(end: Column, start: Column): Column = {
  datediff(end, start)/365
}