fix conversion of value classes to dataframe

Kotlin · Dec 18, 2023 · 9b853ed · 9b853ed
1 parent 63fdb46
commit 9b853ed
Show file tree

Hide file tree

Showing 4 changed files with 189 additions and 2 deletions.
diff --git a/...generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt b/...generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt
@@ -21,13 +21,16 @@ import org.jetbrains.kotlinx.dataframe.impl.getListType
 import org.jetbrains.kotlinx.dataframe.impl.projectUpTo
 import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertiesOrder
 import java.lang.reflect.InvocationTargetException
+import java.lang.reflect.Method
 import java.time.temporal.Temporal
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 import kotlin.reflect.KVisibility
 import kotlin.reflect.full.isSubclassOf
 import kotlin.reflect.full.memberProperties
+import kotlin.reflect.full.primaryConstructor
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.jvm.isAccessible
 import kotlin.reflect.jvm.javaField
 import kotlin.reflect.typeOf
 
@@ -141,7 +144,26 @@ internal fun convertToDataFrame(
         val property = it
         if (excludes.contains(property)) return@mapNotNull null
 
+        class ValueClassConverter(val unbox: Method, val box: Method)
+
+        val valueClassConverter = (it.returnType.classifier as? KClass<*>)?.let { kClass ->
+            if (!kClass.isValue) null else {
+                val constructor =
+                    requireNotNull(kClass.primaryConstructor) { "value class $kClass is expected to have primary constructor, but couldn't obtain it" }
+                val parameter = constructor.parameters.singleOrNull()
+                    ?: error("conversion of value class $kClass with multiple parameters in constructor is not yet supported")
+                // there's no need to unwrap if underlying field is nullable
+                if (parameter.type.isMarkedNullable) return@let null
+                // box and unbox impl methods are part of binary API of value classes
+                // https://youtrack.jetbrains.com/issue/KT-50518/Boxing-Unboxing-methods-for-JvmInline-value-classes-should-be-public-accessible
+                val unbox = kClass.java.getMethod("unbox-impl")
+                val box = kClass.java.methods.single { it.name == "box-impl" }
+                val valueClassConverter = ValueClassConverter(unbox, box)
+                valueClassConverter
+            }
+        }
         property.javaField?.isAccessible = true
+//        property.isAccessible = true
 
         var nullable = false
         var hasExceptions = false
@@ -151,7 +173,19 @@ internal fun convertToDataFrame(
                 null
             } else {
                 val value = try {
-                    it.call(obj)
+                    val value = it.call(obj)
+                    /**
+                     * here we do what compiler does
+                     * @see org.jetbrains.kotlinx.dataframe.api.CreateDataFrameTests.testKPropertyGetLibrary
+                     */
+                    if (valueClassConverter != null) {
+                        val var1 = value?.let {
+                            valueClassConverter.unbox.invoke(it)
+                        }
+                        var1?.let { valueClassConverter.box.invoke(null, var1) }
+                    } else {
+                        value
+                    }
                 } catch (e: InvocationTargetException) {
                     hasExceptions = true
                     e.targetException

diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.kind
 import org.jetbrains.kotlinx.dataframe.type
 import org.junit.Ignore
 import org.junit.Test
+import kotlin.reflect.KProperty
 import kotlin.reflect.typeOf
 
 class CreateDataFrameTests {
@@ -224,4 +225,63 @@ class CreateDataFrameTests {
             println()
         }
     }
+
+    // nullable field here - no generated unwrapping code
+    @JvmInline
+    internal value class Speed(val kmh: Number?)
+
+    internal class PathSegment(
+        val id: String,
+        val speedLimit: Speed? = null,
+    )
+
+    @Test
+    fun valueClassNullableField() {
+        val segments = listOf(PathSegment("foo", Speed(2.3)), PathSegment("bar"))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed(2.3), null)
+    }
+
+    @Test
+    fun valueClassNullableField1() {
+        val segments = listOf(PathSegment("foo", Speed(2.3)), PathSegment("bar", Speed(null)))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed(2.3), Speed(null))
+    }
+
+    @JvmInline
+    internal value class Speed1(val kmh: Number)
+
+    internal class PathSegment1(
+        val id: String,
+        val speedLimit: Speed1? = null,
+    )
+
+    @Test
+    fun valueClass() {
+        val segments = listOf(PathSegment1("foo", Speed1(2.3)), PathSegment1("bar"))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed1(2.3), null)
+    }
+
+    @Test
+    fun testKPropertyGet() {
+        val segment = PathSegment("bar")
+        val result = PathSegment::speedLimit.call(segment)
+        result shouldBe null
+    }
+
+    fun call(kProperty0: KProperty<*>, obj: Any) = kProperty0.call(obj)
+
+    @Test
+    fun testKPropertyCallLibrary() {
+        val segment = PathSegment1("bar")
+        val result = call(PathSegment1::speedLimit, segment)
+        // Sudden result! I cannot create this value, so toString.
+        // In the test above you can see decompiled code that "fixes" this strange wrapping
+        result.toString() shouldBe "Speed1(kmh=null)"
+    }
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt
@@ -21,13 +21,16 @@ import org.jetbrains.kotlinx.dataframe.impl.getListType
 import org.jetbrains.kotlinx.dataframe.impl.projectUpTo
 import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertiesOrder
 import java.lang.reflect.InvocationTargetException
+import java.lang.reflect.Method
 import java.time.temporal.Temporal
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 import kotlin.reflect.KVisibility
 import kotlin.reflect.full.isSubclassOf
 import kotlin.reflect.full.memberProperties
+import kotlin.reflect.full.primaryConstructor
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.jvm.isAccessible
 import kotlin.reflect.jvm.javaField
 import kotlin.reflect.typeOf
 
@@ -141,6 +144,24 @@ internal fun convertToDataFrame(
         val property = it
         if (excludes.contains(property)) return@mapNotNull null
 
+        class ValueClassConverter(val unbox: Method, val box: Method)
+
+        val valueClassConverter = (it.returnType.classifier as? KClass<*>)?.let { kClass ->
+            if (!kClass.isValue) null else {
+                val constructor =
+                    requireNotNull(kClass.primaryConstructor) { "value class $kClass is expected to have primary constructor, but couldn't obtain it" }
+                val parameter = constructor.parameters.singleOrNull()
+                    ?: error("conversion of value class $kClass with multiple parameters in constructor is not yet supported")
+                // there's no need to unwrap if underlying field is nullable
+                if (parameter.type.isMarkedNullable) return@let null
+                // box and unbox impl methods are part of binary API of value classes
+                // https://youtrack.jetbrains.com/issue/KT-50518/Boxing-Unboxing-methods-for-JvmInline-value-classes-should-be-public-accessible
+                val unbox = kClass.java.getMethod("unbox-impl")
+                val box = kClass.java.methods.single { it.name == "box-impl" }
+                val valueClassConverter = ValueClassConverter(unbox, box)
+                valueClassConverter
+            }
+        }
         property.javaField?.isAccessible = true
 
         var nullable = false
@@ -151,7 +172,19 @@ internal fun convertToDataFrame(
                 null
             } else {
                 val value = try {
-                    it.call(obj)
+                    val value = it.call(obj)
+                    /**
+                     * here we do what compiler does
+                     * @see org.jetbrains.kotlinx.dataframe.api.CreateDataFrameTests.testKPropertyGetLibrary
+                     */
+                    if (valueClassConverter != null) {
+                        val var1 = value?.let {
+                            valueClassConverter.unbox.invoke(it)
+                        }
+                        var1?.let { valueClassConverter.box.invoke(null, var1) }
+                    } else {
+                        value
+                    }
                 } catch (e: InvocationTargetException) {
                     hasExceptions = true
                     e.targetException

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.kind
 import org.jetbrains.kotlinx.dataframe.type
 import org.junit.Ignore
 import org.junit.Test
+import kotlin.reflect.KProperty
 import kotlin.reflect.typeOf
 
 class CreateDataFrameTests {
@@ -224,4 +225,63 @@ class CreateDataFrameTests {
             println()
         }
     }
+
+    // nullable field here - no generated unwrapping code
+    @JvmInline
+    internal value class Speed(val kmh: Number?)
+
+    internal class PathSegment(
+        val id: String,
+        val speedLimit: Speed? = null,
+    )
+
+    @Test
+    fun valueClassNullableField() {
+        val segments = listOf(PathSegment("foo", Speed(2.3)), PathSegment("bar"))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed(2.3), null)
+    }
+
+    @Test
+    fun valueClassNullableField1() {
+        val segments = listOf(PathSegment("foo", Speed(2.3)), PathSegment("bar", Speed(null)))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed(2.3), Speed(null))
+    }
+
+    @JvmInline
+    internal value class Speed1(val kmh: Number)
+
+    internal class PathSegment1(
+        val id: String,
+        val speedLimit: Speed1? = null,
+    )
+
+    @Test
+    fun valueClass() {
+        val segments = listOf(PathSegment1("foo", Speed1(2.3)), PathSegment1("bar"))
+
+        val df = segments.toDataFrame()
+        df["speedLimit"].values() shouldBe listOf(Speed1(2.3), null)
+    }
+
+    @Test
+    fun testKPropertyGet() {
+        val segment = PathSegment("bar")
+        val result = PathSegment::speedLimit.call(segment)
+        result shouldBe null
+    }
+
+    fun call(kProperty0: KProperty<*>, obj: Any) = kProperty0.call(obj)
+
+    @Test
+    fun testKPropertyCallLibrary() {
+        val segment = PathSegment1("bar")
+        val result = call(PathSegment1::speedLimit, segment)
+        // Sudden result! I cannot create this value, so toString.
+        // In the test above you can see decompiled code that "fixes" this strange wrapping
+        result.toString() shouldBe "Speed1(kmh=null)"
+    }
 }