title: Writing Spark functions that take Column arguments and return a Column
body: <p>I would like to write a <code>yeardiff</code> function that works similarly to <code>datediff</code>.  <code>yeardiff</code> should take two <code>Column</code> arguments and return a Column with the number of years between those parameter Columns.</p>
  
  <p>Let's use the following sample data:</p>
  
  <pre><code>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"))
  </code></pre>
  
  <p>We can run this to get the date difference:</p>
  
  <pre><code>testDf.withColumn("num_days", datediff(col("first_datetime"), col("second_datetime")))
  </code></pre>
  
  <p>I want to be able to run this:</p>
  
  <pre><code>testDf.withColumn("num_years", yeardiff(col("first_datetime"), col("second_datetime")))
  </code></pre>
  
  <p>I tried to define a <code>yeardiff</code> function with the necessary method signature and didn't get anywhere:</p>
  
  <pre><code>def yeardiff(end: Column, start: Column): Column = {
    // what do I do here
  }    
  </code></pre>
  
  <p>Here is a hacked transformation solution that I came up with and don't like:</p>
  
  <pre><code>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)
  }
  </code></pre>
  
  <p>EDIT</p>
  
  <p>I started digging into the Spark source code to see how <code>datediff</code> works.  Here is <a href="https://github.com/apache/spark/blob/7a7ce272fe9a703f58b0180a9d2001ecb5c4b8db/sql/core/src/main/scala/org/apache/spark/sql/functions.scala#L2516" rel="nofollow noreferrer">the <code>datediff</code> function definition</a>:</p>
  
  <pre><code>def datediff(end: Column, start: Column): Column = withExpr { DateDiff(end.expr, start.expr) }
  </code></pre>
  
  <p>Here is <a href="https://github.com/apache/spark/blob/7a7ce272fe9a703f58b0180a9d2001ecb5c4b8db/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala#L1329" rel="nofollow noreferrer">the <code>DateDiff</code> case class</a>:</p>
  
  <pre><code>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) =&gt; s"$end - $start")
    }
  }
  </code></pre>