TypedColumn#year and LocalDateTime generator

dataset/src/main/scala/frameless/functions/NonAggregateFunctions.scala

@@ -1,8 +1,7 @@
 package frameless
 package functions
 
-import org.apache.spark.sql.Column
-import org.apache.spark.sql.{functions => untyped}
+import org.apache.spark.sql.{Column, functions => untyped}
 
 import scala.util.matching.Regex
 
@@ -279,4 +278,13 @@ trait NonAggregateFunctions {
   def upper[T](str: TypedColumn[T, String]): TypedColumn[T, String] = {
     new TypedColumn[T, String](untyped.upper(str.untyped))
   }
-}
+
+  /**
+    * Non-Aggregate function: Extracts the year as an integer from a given date/timestamp/string.
+    *
+    * apache/spark
+    */
+  def year[T](col: TypedColumn[T, String]): TypedColumn[T, Int] = {
+    new TypedColumn[T, Int](untyped.year(col.untyped))
+  }
+}

dataset/src/test/scala/frameless/functions/NonAggregateFunctionsTests.scala

@@ -1,10 +1,11 @@
 package frameless
 package functions
+import java.time.{LocalDateTime => JavaLocalDateTime}
+
 import frameless.functions.nonAggregate._
-import org.apache.spark.sql.{ Column, Encoder }
-import org.scalacheck.Gen
+import org.apache.spark.sql.{Column, Encoder, functions => untyped}
 import org.scalacheck.Prop._
-import org.apache.spark.sql.{ functions => untyped }
+import org.scalacheck.{Gen, Prop}
 
 class NonAggregateFunctionsTests extends TypedDatasetSuite {
 
@@ -611,23 +612,46 @@ class NonAggregateFunctionsTests extends TypedDatasetSuite {
     check(stringFuncProp(upper, untyped.upper))
   }
 
-  def stringFuncProp[A : Encoder](strFunc: TypedColumn[X1[String], String] => TypedColumn[X1[String], A], sparkFunc: Column => Column) = {
+  test("year") {
+    val spark = session
+    import spark.implicits._
+
+    check(dateTimeStringFuncProp(year, untyped.year))
+  }
+
+  def stringFuncProp[A: Encoder](strFunc: TypedColumn[X1[String], String] => TypedColumn[X1[String], A],
+                                 sparkFunc: Column => Column): Prop = {
     forAll { values: List[X1[String]] =>
       val ds = TypedDataset.create(values)
 
-      val sparkResult: List[A] = ds.toDF()
-        .select(sparkFunc(untyped.col("a")))
-        .map(_.getAs[A](0))
-        .collect()
-        .toList
+      funcProp(ds)(strFunc, sparkFunc)
+    }
+  }
 
-      val typed: List[A] = ds
-        .select(strFunc(ds[String]('a)))
-        .collect()
-        .run()
-        .toList
+  def dateTimeStringFuncProp[A: Encoder](strFunc: TypedColumn[X1[String], String] => TypedColumn[X1[String], A],
+                                         sparkFunc: Column => Column): Prop =
+    forAll { values: List[JavaLocalDateTime] =>
+      val ds = TypedDataset.create(values.map(v => X1[String](v.format(dateTimeFormatter))))
 
-      typed ?= sparkResult
+      funcProp(ds)(strFunc, sparkFunc)
     }
+
+  def funcProp[A: Encoder](ds: TypedDataset[X1[String]])
+                          (strFunc: TypedColumn[X1[String], String] => TypedColumn[X1[String], A],
+                           sparkFunc: Column => Column): Prop = {
+    val sparkResult = ds.toDF()
+      .select(sparkFunc(untyped.col("a")))
+      .map(_.getAs[A](0))
+      .collect()
+      .toList
+
+    val typed = ds
+      .select(strFunc(ds[String]('a)))
+      .collect()
+      .run()
+      .toList
+
+    typed ?= sparkResult
   }
-}
+
+}

dataset/src/test/scala/frameless/package.scala

@@ -1,3 +1,6 @@
+import java.time.format.DateTimeFormatter
+import java.time.{LocalDateTime => JavaLocalDateTime}
+
 import org.scalacheck.{Arbitrary, Gen}
 
 package object frameless {
@@ -37,4 +40,15 @@ package object frameless {
     } yield new UdtEncodedClass(int, doubles.toArray)
   }
 
+  val dateTimeFormatter: DateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm")
+
+  implicit val localDateArb: Arbitrary[JavaLocalDateTime] = Arbitrary {
+    for {
+      year <- Gen.chooseNum(1900, 2027)
+      month <- Gen.chooseNum(1, 12)
+      dayOfMonth <- Gen.chooseNum(1, 28)
+      hour <- Gen.chooseNum(1, 23)
+      minute <- Gen.chooseNum(1, 59)
+    } yield JavaLocalDateTime.of(year, month, dayOfMonth, hour, minute)
+  }
 }