Map support

Kotlin · Mar 17, 2024 · 0c8f4b1 · 0c8f4b1
1 parent 0d061b6
commit 0c8f4b1
Showing 1 changed file with 18 additions and 2 deletions.
diff --git a/kotlin-spark-api/src/main/kotlin/org/jetbrains/kotlinx/spark/api/Encoding.kt b/kotlin-spark-api/src/main/kotlin/org/jetbrains/kotlinx/spark/api/Encoding.kt
@@ -463,8 +463,24 @@ object KotlinTypeInference {
                 )
             }
 
-            currentType.isSubtypeOf<Map<*, *>?>() -> TODO()
-            currentType.isSubtypeOf<scala.collection.Map<*, *>?>() -> TODO()
+            currentType.isSubtypeOf<Map<*, *>?>() || currentType.isSubtypeOf<scala.collection.Map<*, *>?>() -> {
+                val keyEncoder = encoderFor(
+                    currentType = tArguments[0].type!!,
+                    seenTypeSet = seenTypeSet,
+                    typeVariables = typeVariables,
+                )
+                val valueEncoder = encoderFor(
+                    currentType = tArguments[1].type!!,
+                    seenTypeSet = seenTypeSet,
+                    typeVariables = typeVariables,
+                )
+                AgnosticEncoders.MapEncoder(
+                    /* clsTag = */ ClassTag.apply<Map<*, *>>(jClass),
+                    /* keyEncoder = */ keyEncoder,
+                    /* valueEncoder = */ valueEncoder,
+                    /* valueContainsNull = */ tArguments[1].type!!.isMarkedNullable,
+                )
+            }
 
             kClass.isData -> {
                 if (currentType in seenTypeSet) throw IllegalStateException("Circular reference detected for type $currentType")