Adding utils to help ensure that compile time schema ~ runtime schema

koperagen · koperagen · commit d744ddc17c43 · 2024-07-05T18:04:37.000+03:00
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/schema.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/schema.kt
@@ -2,8 +2,10 @@ package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
+import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.impl.owner
 import org.jetbrains.kotlinx.dataframe.impl.schema.extractSchema
+import org.jetbrains.kotlinx.dataframe.impl.schema.getSchema
 import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
 
 // region DataRow
@@ -23,3 +25,6 @@ public fun AnyFrame.schema(): DataFrameSchema = extractSchema()
 public fun GroupBy<*, *>.schema(): DataFrameSchema = toDataFrame().schema()
 
 // endregion
+
+@Suppress("UnusedReceiverParameter")
+public inline fun <reified T> DataFrame<T>.compileTimeSchema(): DataFrameSchema = getSchema(T::class)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/schema/Utils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/schema/Utils.kt
@@ -96,6 +96,11 @@ internal fun AnyCol.extractSchema(): ColumnSchema = when (this) {
     else -> throw RuntimeException("Unknown column type: $this")
 }
 
+@PublishedApi
+internal fun getSchema(kClass: KClass<*>): DataFrameSchema {
+    return MarkersExtractor.get(kClass).schema
+}
+
 internal fun ColumnSchema.createEmptyColumn(name: String): AnyCol = when (this) {
     is ColumnSchema.Value -> DataColumn.createValueColumn<Any?>(name, emptyList(), type)
     is ColumnSchema.Group -> DataColumn.createColumnGroup(name, schema.createEmptyDataFrame()) as AnyCol
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
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import io.kotest.assertions.throwables.shouldNotThrowAny
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
@@ -397,4 +398,30 @@ class CreateDataFrameTests {
             DataColumn.createValueColumn("c", listOf(arrayOf(5, null)), typeOf<Array<Int?>>()),
         ).schema()
     }
+
+    @DataSchema
+    data class Person(val firstName: String, val lastName: String, val age: Int, val city: String?) : DataRowSchema
+
+    @DataSchema
+    data class Group(val id: String, val participants: List<Person>) : DataRowSchema
+
+    @Test
+    fun `deeply convert data schema and list of data schema`() {
+        val participants1 = listOf(
+            Person("Alice", "Cooper", 15, "London"),
+            Person("Bob", "Dylan", 45, "Dubai"),
+        )
+        val participants2 = listOf(
+            Person("Charlie", "Daniels", 20, "Moscow"),
+            Person("Charlie", "Chaplin", 40, "Milan"),
+        )
+        val df = dataFrameOf(
+            Group("1", participants1),
+            Group("2", participants2),
+        )
+        shouldNotThrowAny {
+            df.participants[0].firstName
+            df.participants[0].city
+        }
+    }
 }
diff --git a/plugins/kotlin-dataframe/testData/box/dataFrameOf.kt b/plugins/kotlin-dataframe/testData/box/dataFrameOf.kt
@@ -0,0 +1,26 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Person(val firstName: String, val lastName: String, val age: Int, val city: String?)
+
+@DataSchema
+data class Group(val id: String, val participants: List<Person>)
+
+fun box(): String {
+    val df = dataFrameOf(
+        Group("1", listOf(
+            Person("Alice", "Cooper", 15, "London"),
+            Person("Bob", "Dylan", 45, "Dubai")
+        )),
+        Group("2", listOf(
+            Person("Charlie", "Daniels", 20, "Moscow"),
+            Person("Charlie", "Chaplin", 40, "Milan"),
+        )),
+    )
+
+    df.compareSchemas()
+    return "OK"
+}
diff --git a/plugins/kotlin-dataframe/testData/box/toDataFrame_superType.kt b/plugins/kotlin-dataframe/testData/box/toDataFrame_superType.kt
@@ -16,6 +16,7 @@ class MyImpl(override val i: Int, override val a: Int) : MyInterface
 fun test(list: List<MyInterface>) {
     val df = list.toDataFrame()
     df.schema().print()
+    df.compareSchemas(strict = true)
     df.i
     df.a
 }
diff --git a/plugins/kotlin-dataframe/testData/testUtils.kt b/plugins/kotlin-dataframe/testData/testUtils.kt
@@ -0,0 +1,19 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+inline fun <reified T> DataFrame<T>.compareSchemas(strict: Boolean = false) {
+    val schema = schema()
+    val compileTimeSchema = compileTimeSchema()
+    val compare = compileTimeSchema.compare(schema)
+    require(if (strict) compare.isEqual() else compare.isSuperOrEqual()) {
+        buildString {
+            appendLine("Comparison result: $compare")
+            appendLine("Runtime:")
+            appendLine(schema.toString())
+            appendLine("Compile:")
+            appendLine(compileTimeSchema.toString())
+        }
+    }
+}
diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
@@ -52,6 +52,12 @@ public void testConvertToDataFrame() {
     runTest("testData/box/convertToDataFrame.kt");
   }
 
+  @Test
+  @TestMetadata("dataFrameOf.kt")
+  public void testDataFrameOf() {
+    runTest("testData/box/dataFrameOf.kt");
+  }
+
   @Test
   @TestMetadata("dataRowSchemaApi.kt")
   public void testDataRowSchemaApi() {
diff --git a/plugins/kotlin-dataframe/tests/org/jetbrains/kotlin/fir/dataframe/AbstractDataFrameBlackBoxCodegenTest.kt b/plugins/kotlin-dataframe/tests/org/jetbrains/kotlin/fir/dataframe/AbstractDataFrameBlackBoxCodegenTest.kt
@@ -5,20 +5,20 @@
 
 package org.jetbrains.kotlin.fir.dataframe
 
-import org.jetbrains.kotlin.fir.dataframe.AbstractDataFrameBlackBoxCodegenTest.MyClasspathProvider
 import org.jetbrains.kotlin.fir.dataframe.services.DataFramePluginAnnotationsProvider
 import org.jetbrains.kotlin.fir.dataframe.services.ExperimentalExtensionRegistrarConfigurator
 import org.jetbrains.kotlin.fir.dataframe.services.TemporaryDirectoryManagerImplFixed
 import org.jetbrains.kotlin.fir.dataframe.services.classpath.classpathFromClassloader
 import org.jetbrains.kotlin.test.TestJdkKind
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
-import org.jetbrains.kotlin.test.directives.CodegenTestDirectives
 import org.jetbrains.kotlin.test.directives.CodegenTestDirectives.IGNORE_DEXING
-import org.jetbrains.kotlin.test.directives.FirDiagnosticsDirectives
 import org.jetbrains.kotlin.test.directives.JvmEnvironmentConfigurationDirectives
+import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
 import org.jetbrains.kotlin.test.initIdeaConfiguration
+import org.jetbrains.kotlin.test.model.TestFile
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.runners.codegen.AbstractFirLightTreeBlackBoxCodegenTest
+import org.jetbrains.kotlin.test.services.AdditionalSourceProvider
 import org.jetbrains.kotlin.test.services.EnvironmentBasedStandardLibrariesPathProvider
 import org.jetbrains.kotlin.test.services.KotlinStandardLibrariesPathProvider
 import org.jetbrains.kotlin.test.services.RuntimeClasspathProvider
@@ -48,6 +48,7 @@ open class AbstractDataFrameBlackBoxCodegenTest : AbstractFirLightTreeBlackBoxCo
         builder.useConfigurators(::DataFramePluginAnnotationsProvider)
         builder.useConfigurators(::ExperimentalExtensionRegistrarConfigurator)
         builder.useCustomRuntimeClasspathProviders(::MyClasspathProvider)
+        builder.useAdditionalSourceProviders(::TestUtilsSourceProvider)
     }
 
     override fun runTest(filePath: String) {
@@ -65,5 +66,15 @@ open class AbstractDataFrameBlackBoxCodegenTest : AbstractFirLightTreeBlackBoxCo
     override fun createKotlinStandardLibrariesPathProvider(): KotlinStandardLibrariesPathProvider {
         return EnvironmentBasedStandardLibrariesPathProvider
     }
+
+    class TestUtilsSourceProvider(testServices: TestServices) : AdditionalSourceProvider(testServices) {
+        companion object {
+            const val COMMON_SOURCE_PATH = "testData/testUtils.kt"
+        }
+
+        override fun produceAdditionalFiles(globalDirectives: RegisteredDirectives, module: TestModule): List<TestFile> {
+            return listOf(File(COMMON_SOURCE_PATH).toTestFile())
+        }
+    }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ class MyImpl(override val i: Int, override val a: Int) : MyInterface`
`16`	`16`	`fun test(list: List<MyInterface>) {`
`17`	`17`	`val df = list.toDataFrame()`
`18`	`18`	`df.schema().print()`
	`19`	`+ df.compareSchemas(strict = true)`
`19`	`20`	`df.i`
`20`	`21`	`df.a`
`21`	`22`	`}`