Implement map case for SchemaEncoder

grouzen · Dec 10, 2023 · 7480a82 · 7480a82
1 parent cec2966
commit 7480a82
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/Schemas.scala b/modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/Schemas.scala
@@ -81,6 +81,23 @@ object Schemas {
 
   }
 
+  case class MapDef(key: Type, value: Type, isOptional: Boolean = false) extends Def[MapDef] {
+
+    override def named(name: String): Type =
+      Types
+        .map(repetition(isOptional))
+        .key(key)
+        .value(value)
+        .named(name)
+
+    override def required: MapDef =
+      this.copy(isOptional = false)
+
+    override def optional: MapDef =
+      this.copy(isOptional = true)
+
+  }
+
   def repetition(optional: Boolean): Repetition =
     if (optional) Repetition.OPTIONAL else Repetition.REQUIRED
 
@@ -97,5 +114,6 @@ object Schemas {
 
   def record(fields: Chunk[Type]): RecordDef = RecordDef(fields)
   def list(element: Type): ListDef           = ListDef(element)
+  def map(key: Type, value: Type)            = MapDef(key, value)
 
 }
diff --git a/...re/src/main/scala/me/mnedokushev/zio/apache/parquet/core/codec/SchemaEncoderDeriver.scala b/...re/src/main/scala/me/mnedokushev/zio/apache/parquet/core/codec/SchemaEncoderDeriver.scala
@@ -85,7 +85,16 @@ object SchemaEncoderDeriver {
       key: => SchemaEncoder[K],
       value: => SchemaEncoder[V],
       summoned: => Option[SchemaEncoder[Map[K, V]]]
-    ): SchemaEncoder[Map[K, V]] = ???
+    ): SchemaEncoder[Map[K, V]] = new SchemaEncoder[Map[K, V]] {
+      override def encode(schema: Schema[Map[K, V]], name: String, optional: Boolean): Type =
+        Schemas
+          .map(
+            key.encode(map.keySchema, "key", optional = false),
+            value.encode(map.valueSchema, "value", optional = isSchemaOptional(map.valueSchema))
+          )
+          .optionality(optional)
+          .named(name)
+    }
 
     override def deriveTransformedRecord[A, B](
       record: Schema.Record[A],

diff --git a/...rc/test/scala/me/mnedokushev/zio/apache/parquet/core/codec/SchemaEncoderDeriverSpec.scala b/...rc/test/scala/me/mnedokushev/zio/apache/parquet/core/codec/SchemaEncoderDeriverSpec.scala
@@ -164,6 +164,19 @@ object SchemaEncoderDeriverSpec extends ZIOSpecDefault {
             )
           }
           .reduce(_ && _)
+      },
+      test("map") {
+        val name = "mymap"
+
+        val encoder = Derive.derive[SchemaEncoder, Map[String, Int]](SchemaEncoderDeriver.default)
+        val tpe     = encoder.encode(Schema.map[String, Int], name, optional = true)
+
+        assertTrue(
+          tpe == Schemas
+            .map(Schemas.string.required.named("key"), Schemas.int.required.named("value"))
+            .optional
+            .named(name)
+        )
       }
 //      test("summoned") {
       //        // @nowarn annotation is needed to avoid having 'variable is not used' compiler error