Add support for enums (#52)

vincenzobaz · Oct 26, 2023 · 2c88a28 · 2c88a28
1 parent efed44b
commit 2c88a28
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 3 deletions.
diff --git a/encoders/src/main/scala/scala3encoders/EncoderDerivation.scala b/encoders/src/main/scala/scala3encoders/EncoderDerivation.scala
@@ -13,7 +13,7 @@ given encoder[T](using
     deserializer: Deserializer[T],
     classTag: ClassTag[T]
 ): ExpressionEncoder[T] =
-  val inputObject = BoundReference(0, serializer.inputType, true)
+  val inputObject = BoundReference(0, serializer.inputType, nullable = true)
   val path = GetColumnByOrdinal(0, deserializer.inputType)
 
   ExpressionEncoder(

diff --git a/encoders/src/main/scala/scala3encoders/derivation/Deserializer.scala b/encoders/src/main/scala/scala3encoders/derivation/Deserializer.scala
@@ -2,7 +2,7 @@ package scala3encoders.derivation
 
 import scala.compiletime.{constValue, summonInline, erasedValue}
 import scala.deriving.Mirror
-import scala.reflect.ClassTag
+import scala.reflect.{ClassTag, Enum}
 
 import org.apache.spark.sql.catalyst.expressions.{
   Expression,
@@ -144,6 +144,19 @@ object Deserializer:
     def deserialize(path: Expression): Expression =
       createDeserializerForScalaBigInt(path)
 
+  given[E <: Enum : ClassTag]: Deserializer[E] with
+    def inputType: DataType = StringType
+
+    def deserialize(path: Expression): Expression =
+      val string = summon[Deserializer[String]].deserialize(path)
+      StaticInvoke(
+        summon[ClassTag[E]].runtimeClass,
+        ObjectType(summon[ClassTag[E]].runtimeClass),
+        "valueOf",
+        string :: Nil,
+        returnNullable = false
+      )
+
   inline given deriveArray[T](using
       d: Deserializer[T],
       ct: ClassTag[T]

diff --git a/encoders/src/main/scala/scala3encoders/derivation/Serializer.scala b/encoders/src/main/scala/scala3encoders/derivation/Serializer.scala
@@ -2,7 +2,7 @@ package scala3encoders.derivation
 
 import scala.compiletime.{constValue, summonInline, erasedValue}
 import scala.deriving.Mirror
-import scala.reflect.ClassTag
+import scala.reflect.{ClassTag, Enum}
 
 import org.apache.spark.sql.catalyst.expressions.{Expression, KnownNotNull}
 import org.apache.spark.sql.catalyst.expressions.objects.Invoke
@@ -118,6 +118,17 @@ object Serializer:
     def inputType: DataType = ObjectType(classOf[String])
     def serialize(inputObject: Expression): Expression =
       createSerializerForString(inputObject)
+
+  given [E <: Enum: ClassTag]: Serializer[E] with
+    def inputType: DataType = ObjectType(summon[ClassTag[E]].runtimeClass)
+    def serialize(inputObject: Expression): Expression =
+      val string = Invoke(
+        inputObject,
+        "toString",
+        ObjectType(classOf[String]),
+        returnNullable = false
+      )
+      summon[Serializer[String]].serialize(string)
 
   given deriveSeq[F[_], T](using s: Serializer[T])(using
       F[T] <:< Seq[T]

diff --git a/encoders/src/test/scala/sql/EncoderDerivationSpec.scala b/encoders/src/test/scala/sql/EncoderDerivationSpec.scala
@@ -134,6 +134,10 @@ case class City(name: String, lat: Double, lon: Double)
 case class CityWithInts(name: String, lat: Int, lon: Int)
 case class Journey(id: Int, cities: Seq[City])
 
+enum Color:
+  case Red, Black
+case class ColorData(color: Color)
+
 val dSchema =
   StructType(
     Seq(
@@ -275,6 +279,16 @@ class EncoderDerivationSpec extends munit.FunSuite with SparkSqlTesting:
     assertEquals(input.toDS.collect.toSeq, input)
   }
 
+  test("derive encoder of FiniteDuration") {
+    val data = Seq(ColorData(Color.Black), ColorData(Color.Red)).toDS()
+      .map(_.copy(Color.Red))
+    assertEquals(
+      data.schema,
+      StructType(Seq(StructField("color", StringType, true)))
+    )
+    assertEquals(data.collect().toSeq, Seq(ColorData(Color.Red), ColorData(Color.Red)))
+  }
+
   test("List[Int]") {
     val ls = List(List(1, 2, 3), List(4, 5, 6))
     assertEquals(ls.toDS.collect().toList, ls)