Substitute type parameters with type arguments in toDataFrame conversion

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlinx.dataframe.impl.isGetterLike
 import org.jetbrains.kotlinx.dataframe.impl.isJavaRecord
 import org.jetbrains.kotlinx.dataframe.impl.projectUpTo
 import org.jetbrains.kotlinx.dataframe.impl.recordComponentNames
+import org.jetbrains.kotlinx.dataframe.impl.replace
 import org.jetbrains.kotlinx.dataframe.impl.schema.sortWithConstructor
 import java.lang.reflect.InvocationTargetException
 import java.lang.reflect.Method
@@ -31,13 +32,13 @@ import kotlin.reflect.KCallable
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
+import kotlin.reflect.KTypeParameter
 import kotlin.reflect.KVisibility
 import kotlin.reflect.full.isSubclassOf
 import kotlin.reflect.full.memberFunctions
 import kotlin.reflect.full.memberProperties
 import kotlin.reflect.full.primaryConstructor
 import kotlin.reflect.full.starProjectedType
-import kotlin.reflect.full.valueParameters
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.jvm.isAccessible
 import kotlin.reflect.jvm.javaField
@@ -237,6 +238,8 @@ internal fun convertToDataFrame(
         return dataFrameOf(column)
     }
 
+    val substitution: Map<KTypeParameter, KType> = type.typeParametersSubstitution()
+
     val properties: List<KCallable<*>> = roots
         .ifEmpty {
             clazz.properties()
@@ -307,7 +310,7 @@ internal fun convertToDataFrame(
             }
         }
 
-        val returnType = property.returnType.let { type ->
+        val returnType = property.returnType.replace(substitution).let { type ->
             if (type.classifier is KClass<*>) {
                 type
             } else {
@@ -429,6 +432,13 @@ internal fun convertToDataFrame(
 
 private fun KType.classifierOrAny(): KClass<*> = classifier as? KClass<*> ?: Any::class
 
+internal fun KType.typeParametersSubstitution(): Map<KTypeParameter, KType> {
+    val klass = classifier as? KClass<*> ?: return emptyMap()
+    return klass.typeParameters.zip(arguments).mapNotNull { (param, projection) ->
+        projection.type?.let { param to it }
+    }.toMap()
+}
+
 private fun KClass<*>.properties(): List<KCallable<*>> {
     // fall back to getter functions for pojo-like classes if no member properties were found
     return memberProperties

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt

@@ -1,8 +1,14 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import io.kotest.assertions.asClue
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.types.shouldBeInstanceOf
 import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
+import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.junit.Test
 import kotlin.reflect.typeOf
 
@@ -65,4 +71,66 @@ class UnfoldTests {
     )
 
     data class Group(val id: String, val participants: List<Person>)
+
+    @Test
+    fun `unfold pair of dataframe structures`() {
+        val schema = dataFrameOf("b" to columnOf(42)).cast<SimpleDataSchema>()
+
+        val df = dataFrameOf("pairs" to columnOf(schema to schema.first()))
+            .unfold("pairs")
+
+        df.schema().asClue {
+            val pairsGroup = df.shouldHaveColumnGroup("pairs")
+            pairsGroup.shouldHaveFrameColumn("first") {
+                it[0].shouldHaveColumn<Int>("b")
+            }
+            pairsGroup.shouldHaveColumnGroup("second") {
+                it.shouldHaveColumn<Int>("b")
+            }
+        }
+    }
+
+    @DataSchema
+    data class SimpleDataSchema(val b: Int)
+
+    @Test
+    fun `unfold pair of dataschema object structures`() {
+        val element = SimpleDataSchema(42)
+        val df = dataFrameOf("pairs" to columnOf(listOf(element) to element))
+            .unfold("pairs")
+
+        df.schema().asClue {
+            val pairsGroup = df.shouldHaveColumnGroup("pairs")
+            pairsGroup.shouldHaveFrameColumn("first") {
+                it[0].shouldHaveColumn<Int>("b")
+            }
+            pairsGroup.shouldHaveColumnGroup("second") {
+                it.shouldHaveColumn<Int>("b")
+            }
+        }
+    }
+
+    fun DataFrame<*>.shouldHaveColumnGroup(
+        name: String,
+        block: (ColumnGroup<*>) -> Unit = {
+        },
+    ): ColumnGroup<*> = getColumnOrNull(name).shouldBeInstanceOf<ColumnGroup<*>>(block)
+
+    fun DataFrame<*>.shouldHaveFrameColumn(
+        name: String,
+        block: (FrameColumn<*>) -> Unit = {
+        },
+    ): FrameColumn<*> = getColumnOrNull(name).shouldBeInstanceOf<FrameColumn<*>>(block)
+
+    inline fun <reified T> DataFrame<*>.shouldHaveColumn(
+        name: String,
+        block: (DataColumn<T>) -> Unit = {
+        },
+    ): DataColumn<T> {
+        val shouldBeInstanceOf = getColumnOrNull(name).shouldBeInstanceOf<DataColumn<*>>()
+        shouldBeInstanceOf.type() shouldBe typeOf<T>()
+        val cast = shouldBeInstanceOf.cast<T>()
+        block(cast)
+        return cast
+    }
 }