Skip to content

Add enum support for ApiV1Kt.generateEncoder #99

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jul 14, 2021
Merged

Add enum support for ApiV1Kt.generateEncoder #99

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5c6caa4
add enum support
nonpool Jul 14, 2021
c50b1f5
revert format
nonpool Jul 14, 2021
dcc715f
support enum for encoder for List<Enum>
nonpool Jul 14, 2021
ceb4d22
add array enum to test
nonpool Jul 14, 2021
90c2375
add enum map for test
nonpool Jul 14, 2021
25b380d
clear test
nonpool Jul 14, 2021
6f1f966
support enum for 2.4
nonpool Jul 14, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion core/3.0/src/main/scala/org/apache/spark/sql/KotlinReflection.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,10 @@ object KotlinReflection extends KotlinReflection {
mirror.runtimeClass(t.typeSymbol.asClass)
)

case t if isSubtype(t, localTypeOf[java.lang.Enum[_]]) =>
createDeserializerForTypesSupportValueOf(
createDeserializerForString(path, returnNullable = false), Class.forName(t.toString))

case t if t.typeSymbol.annotations.exists(_.tree.tpe =:= typeOf[SQLUserDefinedType]) =>
val udt = getClassFromType(t).getAnnotation(classOf[SQLUserDefinedType]).udt().
getConstructor().newInstance()
Expand Down Expand Up @@ -612,7 +616,7 @@ object KotlinReflection extends KotlinReflection {
case _: StringType =>
val clsName = getClassNameFromType(typeOf[String])
val newPath = walkedTypePath.recordArray(clsName)
createSerializerForMapObjects(input, ObjectType(classOf[String]),
createSerializerForMapObjects(input, ObjectType(Class.forName(getClassNameFromType(elementType))),
serializerFor(_, elementType, newPath, seenTypeSet))


Expand Down Expand Up @@ -718,6 +722,10 @@ object KotlinReflection extends KotlinReflection {
case t if isSubtype(t, localTypeOf[java.lang.Boolean]) => createSerializerForBoolean(inputObject)
case t if isSubtype(t, localTypeOf[Boolean]) => createSerializerForBoolean(inputObject)

case t if isSubtype(t, localTypeOf[java.lang.Enum[_]]) =>
createSerializerForString(
Invoke(inputObject, "name", ObjectType(classOf[String]), returnNullable = false))

case t if t.typeSymbol.annotations.exists(_.tree.tpe =:= typeOf[SQLUserDefinedType]) =>
val udt = getClassFromType(t)
.getAnnotation(classOf[SQLUserDefinedType]).udt().getConstructor().newInstance()
Expand Down
3 changes: 3 additions & 0 deletions kotlin-spark-api/2.4/src/main/kotlin/org/jetbrains/kotlinx/spark/api/ApiV1.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -858,6 +858,9 @@ fun schema(type: KType, map: Map<String, KType> = mapOf()): DataType {
it.first.name to it.second.type!!
}.toMap())
return when {
klass.isSubclassOf(Enum::class) -> {
KSimpleTypeWrapper(DataTypes.StringType, klass.java, type.isMarkedNullable)
}
klass.isSubclassOf(Iterable::class) || klass.java.isArray -> {
val listParam = if (klass.java.isArray) {
when (klass) {
Expand Down
48 changes: 48 additions & 0 deletions kotlin-spark-api/2.4/src/test/kotlin/org/jetbrains/kotlinx/spark/api/ApiTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,37 @@ class ApiTest : ShouldSpec({
dataset.sort(SomeClass::a, SomeClass::b)
dataset.takeAsList(1).first().b shouldBe 2
}
should("Generate encoder correctly with complex enum data class") {
val dataset: Dataset<ComplexEnumDataClass> =
dsOf(
ComplexEnumDataClass(
1,
"string",
listOf("1", "2"),
SomeEnum.A,
SomeOtherEnum.C,
listOf(SomeEnum.A, SomeEnum.B),
listOf(SomeOtherEnum.C, SomeOtherEnum.D),
arrayOf(SomeEnum.A, SomeEnum.B),
arrayOf(SomeOtherEnum.C, SomeOtherEnum.D),
mapOf(SomeEnum.A to SomeOtherEnum.C)
)
)

dataset.show(false)
val first = dataset.takeAsList(1).first()

first.int shouldBe 1
first.string shouldBe "string"
first.strings shouldBe listOf("1","2")
first.someEnum shouldBe SomeEnum.A
first.someOtherEnum shouldBe SomeOtherEnum.C
first.someEnums shouldBe listOf(SomeEnum.A, SomeEnum.B)
first.someOtherEnums shouldBe listOf(SomeOtherEnum.C, SomeOtherEnum.D)
first.someEnumArray shouldBe arrayOf(SomeEnum.A, SomeEnum.B)
first.someOtherArray shouldBe arrayOf(SomeOtherEnum.C, SomeOtherEnum.D)
first.enumMap shouldBe mapOf(SomeEnum.A to SomeOtherEnum.C)
}
}
}
})
Expand Down Expand Up @@ -505,3 +536,20 @@ data class Test<Z>(val id: Long, val data: Array<Pair<Z, Int>>) {
data class SomeClass(val a: IntArray, val b: Int) : Serializable

data class SomeOtherClass(val a: IntArray, val b: Int, val c: Boolean) : Serializable

enum class SomeEnum { A, B }

enum class SomeOtherEnum(val value: Int) { C(1), D(2) }

data class ComplexEnumDataClass(
val int: Int,
val string: String,
val strings: List<String>,
val someEnum: SomeEnum,
val someOtherEnum: SomeOtherEnum,
val someEnums: List<SomeEnum>,
val someOtherEnums: List<SomeOtherEnum>,
val someEnumArray: Array<SomeEnum>,
val someOtherArray: Array<SomeOtherEnum>,
val enumMap: Map<SomeEnum, SomeOtherEnum>
)
3 changes: 3 additions & 0 deletions kotlin-spark-api/3.0/src/main/kotlin/org/jetbrains/kotlinx/spark/api/ApiV1.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -851,6 +851,9 @@ fun schema(type: KType, map: Map<String, KType> = mapOf()): DataType {
it.first.name to it.second.type!!
}.toMap())
return when {
klass.isSubclassOf(Enum::class) -> {
KSimpleTypeWrapper(DataTypes.StringType, klass.java, type.isMarkedNullable)
}
klass.isSubclassOf(Iterable::class) || klass.java.isArray -> {
val listParam = if (klass.java.isArray) {
when (klass) {
Expand Down
48 changes: 48 additions & 0 deletions kotlin-spark-api/3.0/src/test/kotlin/org/jetbrains/kotlinx/spark/api/ApiTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -515,6 +515,37 @@ class ApiTest : ShouldSpec({
dataset.sort(SomeClass::a, SomeClass::b)
dataset.takeAsList(1).first().b shouldBe 2
}
should("Generate encoder correctly with complex enum data class") {
val dataset: Dataset<ComplexEnumDataClass> =
dsOf(
ComplexEnumDataClass(
1,
"string",
listOf("1", "2"),
SomeEnum.A,
SomeOtherEnum.C,
listOf(SomeEnum.A, SomeEnum.B),
listOf(SomeOtherEnum.C, SomeOtherEnum.D),
arrayOf(SomeEnum.A, SomeEnum.B),
arrayOf(SomeOtherEnum.C, SomeOtherEnum.D),
mapOf(SomeEnum.A to SomeOtherEnum.C)
)
)

dataset.show(false)
val first = dataset.takeAsList(1).first()

first.int shouldBe 1
first.string shouldBe "string"
first.strings shouldBe listOf("1","2")
first.someEnum shouldBe SomeEnum.A
first.someOtherEnum shouldBe SomeOtherEnum.C
first.someEnums shouldBe listOf(SomeEnum.A, SomeEnum.B)
first.someOtherEnums shouldBe listOf(SomeOtherEnum.C, SomeOtherEnum.D)
first.someEnumArray shouldBe arrayOf(SomeEnum.A, SomeEnum.B)
first.someOtherArray shouldBe arrayOf(SomeOtherEnum.C, SomeOtherEnum.D)
first.enumMap shouldBe mapOf(SomeEnum.A to SomeOtherEnum.C)
}
}
}
})
Expand All @@ -527,3 +558,20 @@ data class LonLat(val lon: Double, val lat: Double)
data class SomeClass(val a: IntArray, val b: Int) : Serializable

data class SomeOtherClass(val a: IntArray, val b: Int, val c: Boolean) : Serializable

enum class SomeEnum { A, B }

enum class SomeOtherEnum(val value: Int) { C(1), D(2) }

data class ComplexEnumDataClass(
val int: Int,
val string: String,
val strings: List<String>,
val someEnum: SomeEnum,
val someOtherEnum: SomeOtherEnum,
val someEnums: List<SomeEnum>,
val someOtherEnums: List<SomeOtherEnum>,
val someEnumArray: Array<SomeEnum>,
val someOtherArray: Array<SomeOtherEnum>,
val enumMap: Map<SomeEnum, SomeOtherEnum>
)