Implement TypeTag enum deriver

grouzen · Jan 21, 2024 · 363a334 · 363a334
1 parent 4d5fab6
commit 363a334
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 10 deletions.
diff --git a/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/FilterError.scala b/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/FilterError.scala
@@ -0,0 +1,6 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+final case class FilterError(
+  message: String,
+  cause: Option[Throwable] = None
+) extends IllegalArgumentException(message, cause.getOrElse(new Throwable()))
diff --git a/...s/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/OperatorSupport.scala b/...s/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/OperatorSupport.scala
@@ -62,6 +62,8 @@ object OperatorSupport {
   }
 
   object EqNotEq {
+    implicit def enum0[A: TypeTag]: EqNotEq[A] = new EqNotEq[A] {}
+
     implicit case object string         extends EqNotEq[String]
     implicit case object boolean        extends EqNotEq[Boolean]
     implicit case object byte           extends EqNotEq[Byte]

diff --git a/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/TypeTag.scala b/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/TypeTag.scala
@@ -114,6 +114,12 @@ object TypeTag {
 
     }
 
+  implicit def enum0[A](casesMap: Map[A, String]): TypeTag.EqNotEq[A] =
+    eqnoteq[A, Binary, BinaryColumn](
+      FilterApi.binaryColumn,
+      v => Value.string(casesMap.getOrElse(v, throw FilterError(s"Failed to encode enum for value $v"))).value
+    )
+
   implicit val string: TypeTag.EqNotEq[String]                =
     eqnoteq[String, Binary, BinaryColumn](
       FilterApi.binaryColumn,

diff --git a/...es/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/TypeTagDeriver.scala b/...es/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/TypeTagDeriver.scala
@@ -23,8 +23,13 @@ object TypeTagDeriver {
       `enum`: Schema.Enum[A],
       cases: => Chunk[Deriver.WrappedF[TypeTag, _]],
       summoned: => Option[TypeTag[A]]
-    ): TypeTag[A] =
-      TypeTag.dummy[A]
+    ): TypeTag[A] = {
+      val casesMap = `enum`.cases.map { case0 =>
+        case0.schema.asInstanceOf[Schema.CaseClass0[A]].defaultConstruct() -> case0.id
+      }.toMap
+
+      TypeTag.enum0(casesMap)
+    }
 
     override def derivePrimitive[A](
       st: StandardType[A],

diff --git a/...les/core/src/test/scala-2.12/me/mnedokushev/zio/apache/parquet/core/filter/Fixtures.scala b/...les/core/src/test/scala-2.12/me/mnedokushev/zio/apache/parquet/core/filter/Fixtures.scala
@@ -12,7 +12,7 @@ import java.util.UUID
 
 object Fixtures {
 
-  case class MyRecord(a: String, b: Int, child: MyRecord.Child)
+  case class MyRecord(a: String, b: Int, child: MyRecord.Child, enm: MyRecord.Enum)
 
   object MyRecord {
     implicit val schema                     =
@@ -28,6 +28,18 @@ object Fixtures {
         Derive.derive[TypeTag, Child](TypeTagDeriver.default)
     }
 
+    sealed trait Enum
+    object Enum {
+      case object Started    extends Enum
+      case object InProgress extends Enum
+      case object Done       extends Enum
+
+      implicit val schema: Schema[Enum]   =
+        DeriveSchema.gen[Enum]
+      implicit val typeTag: TypeTag[Enum] =
+        Derive.derive[TypeTag, Enum](TypeTagDeriver.default)
+    }
+
   }
 
   case class MyRecordSummoned(a: Int, b: String)

diff --git a/...es/core/src/test/scala-2.13+/me/mnedokushev/zio/apache/parquet/core/filter/Fixtures.scala b/...es/core/src/test/scala-2.13+/me/mnedokushev/zio/apache/parquet/core/filter/Fixtures.scala
@@ -12,12 +12,13 @@ import java.util.UUID
 
 object Fixtures {
 
-  case class MyRecord(a: String, b: Int, child: MyRecord.Child)
+  case class MyRecord(a: String, b: Int, child: MyRecord.Child, enm: MyRecord.Enum)
 
   object MyRecord {
-    implicit val schema: Schema.CaseClass3.WithFields["a", "b", "child", String, Int, MyRecord.Child, MyRecord] =
+    implicit val schema
+      : Schema.CaseClass4.WithFields["a", "b", "child", "enm", String, Int, MyRecord.Child, MyRecord.Enum, MyRecord] =
       DeriveSchema.gen[MyRecord]
-    implicit val typeTag: TypeTag[MyRecord]                                                                     =
+    implicit val typeTag: TypeTag[MyRecord] =
       Derive.derive[TypeTag, MyRecord](TypeTagDeriver.default)
 
     case class Child(c: Int, d: Option[Long])
@@ -26,7 +27,18 @@ object Fixtures {
         DeriveSchema.gen[Child]
       implicit val typeTag: TypeTag[Child]                                                           =
         Derive.derive[TypeTag, Child](TypeTagDeriver.default)
+    }
+
+    sealed trait Enum
+    object Enum {
+      case object Started    extends Enum
+      case object InProgress extends Enum
+      case object Done       extends Enum
 
+      implicit val schema: Schema[Enum]   =
+        DeriveSchema.gen[Enum]
+      implicit val typeTag: TypeTag[Enum] =
+        Derive.derive[TypeTag, Enum](TypeTagDeriver.default)
     }
   }
 

diff --git a/modules/core/src/test/scala/me/mnedokushev/zio/apache/parquet/core/filter/ExprSpec.scala b/modules/core/src/test/scala/me/mnedokushev/zio/apache/parquet/core/filter/ExprSpec.scala
@@ -18,7 +18,7 @@ object ExprSpec extends ZIOSpecDefault {
   override def spec: Spec[TestEnvironment with Scope, Any] =
     suite("ExprSpec")(
       test("compile all operators") {
-        val (a, b, _) = Filter.columns[MyRecord]
+        val (a, b, _, _) = Filter.columns[MyRecord]
 
         val result = Expr.compile(
           Filter.not(
@@ -77,7 +77,7 @@ object ExprSpec extends ZIOSpecDefault {
 
         assert(result)(isRight(equalTo(expected)))
       },
-      test("type tag compilation") {
+      test("compile all primitive types") {
         val (
           string,
           boolean,
@@ -307,9 +307,17 @@ object ExprSpec extends ZIOSpecDefault {
         assert(result1)(isRight(equalTo(expected1))) &&
         assert(result2)(isRight(equalTo(expected2)))
       },
+      test("compile enum") {
+        val (_, _, _, enm) = Filter.columns[MyRecord]
+
+        val result   = Expr.compile(enm === MyRecord.Enum.Done)
+        val expected = FilterApi.eq(FilterApi.binaryColumn("enm"), Value.string("Done").value)
+
+        assert(result)(isRight(equalTo(expected)))
+      },
       test("column path concatenation") {
-        val (a, b, child) = Filter.columns[MyRecord]
-        val (c, d)        = Filter.columns[MyRecord.Child]
+        val (a, b, child, _) = Filter.columns[MyRecord]
+        val (c, d)           = Filter.columns[MyRecord.Child]
 
         assert(a.path)(equalTo("a")) &&
         assert(b.path)(equalTo("b")) &&