From d744ddc17c436d12f64b2b141d54e6fad1ddd147 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Thu, 4 Jul 2024 16:36:39 +0300 Subject: [PATCH] Adding utils to help ensure that compile time schema ~ runtime schema --- .../jetbrains/kotlinx/dataframe/api/schema.kt | 5 ++++ .../kotlinx/dataframe/impl/schema/Utils.kt | 5 ++++ .../kotlinx/dataframe/api/toDataFrame.kt | 27 +++++++++++++++++++ .../testData/box/dataFrameOf.kt | 26 ++++++++++++++++++ .../testData/box/toDataFrame_superType.kt | 1 + .../kotlin-dataframe/testData/testUtils.kt | 19 +++++++++++++ ...DataFrameBlackBoxCodegenTestGenerated.java | 6 +++++ .../AbstractDataFrameBlackBoxCodegenTest.kt | 17 +++++++++--- 8 files changed, 103 insertions(+), 3 deletions(-) create mode 100644 plugins/kotlin-dataframe/testData/box/dataFrameOf.kt create mode 100644 plugins/kotlin-dataframe/testData/testUtils.kt 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 index 2c61d85b91..de2eef8419 100644 --- 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 DataFrame.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 index b083c5ab23..d5b83cbba0 100644 --- 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(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 index 21ec8c9645..cbd304ce8c 100644 --- 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>()), ).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) : 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 new file mode 100644 index 0000000000..6cee620c6f --- /dev/null +++ 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) + +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 index 892e20a8bc..aede9208ff 100644 --- 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) { 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 new file mode 100644 index 0000000000..b02c91d831 --- /dev/null +++ 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 DataFrame.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 index b6727312c6..dadc68e546 100644 --- 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 index 6df5441e9a..a9288aaf0a 100644 --- 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 { + return listOf(File(COMMON_SOURCE_PATH).toTestFile()) + } + } }