From 95b9ade3eef92315799abddaf641dbf07e84959b Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Fri, 10 Jul 2026 00:57:39 +0300 Subject: [PATCH] Clarify that schema of property with `DataFrame?` type is not a FrameColumn This makes extracted schema represent runtime more accurately --- .../dataframe/codeGen/MarkersExtractor.kt | 2 +- .../kotlinx/dataframe/impl/api/convertTo.kt | 9 +- .../kotlinx/dataframe/impl/api/toDataFrame.kt | 31 ++++--- .../kotlinx/dataframe/api/convertTo.kt | 86 +++++++++++-------- .../dataframe/codeGen/CodeGenerationTests.kt | 17 ++++ 5 files changed, 98 insertions(+), 47 deletions(-) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt index 98cbb59acc..4854d14b39 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt @@ -22,7 +22,7 @@ import kotlin.reflect.typeOf internal fun KType.getFieldKind(): FieldKind = FieldKind.of( this, - isDataFrame = { jvmErasure == DataFrame::class }, + isDataFrame = { jvmErasure == DataFrame::class && !isMarkedNullable }, isListToFrame = { jvmErasure == List::class && (arguments[0].type?.jvmErasure?.hasAnnotation() == true) }, diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/convertTo.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/convertTo.kt index eca83dfcd7..4fc1745e8c 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/convertTo.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/convertTo.kt @@ -193,7 +193,14 @@ internal fun AnyFrame.convertToImpl( when (targetSchema.kind) { ColumnKind.Value -> - convertedColumn ?: originalColumn.convertTo(to) + when { + convertedColumn != null -> convertedColumn + + originalColumn.kind == ColumnKind.Frame && to.jvmErasure == DataFrame::class -> + originalColumn + + else -> originalColumn.convertTo(to) + } ColumnKind.Group -> { val column = when { 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 index ceca88942c..ce77bfa619 100644 --- 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 @@ -320,18 +320,27 @@ internal fun convertToDataFrame( val kClass = returnType.classifier as KClass<*> val fieldKind = returnType.getFieldKind() + // property type might not be of FrameColumn kind if `AnyFrame?`, but we still have to narrow it to FrameCol + // when no actual nulls are met + val shouldCreateFrameCol = kClass == DataFrame::class && !nullable + val keepSubtree = - maxDepth <= 0 && !fieldKind.shouldBeConvertedToFrameColumn && !fieldKind.shouldBeConvertedToColumnGroup - val shouldCreateValueCol = keepSubtree || - kClass in preserveClasses || - property in preserveProperties || - ( - !kClass.canBeUnfolded && - !fieldKind.shouldBeConvertedToFrameColumn && - !fieldKind.shouldBeConvertedToColumnGroup - ) + maxDepth <= 0 && + !fieldKind.shouldBeConvertedToFrameColumn && + !fieldKind.shouldBeConvertedToColumnGroup && + !shouldCreateFrameCol + + val shouldCreateValueCol = + keepSubtree || + kClass in preserveClasses || + property in preserveProperties || + ( + !kClass.canBeUnfolded && + !fieldKind.shouldBeConvertedToFrameColumn && + !fieldKind.shouldBeConvertedToColumnGroup && + !shouldCreateFrameCol + ) - val shouldCreateFrameCol = kClass == DataFrame::class && !nullable val shouldCreateColumnGroup = kClass == DataRow::class if (shouldCreateFrameCol && shouldCreateValueCol) { @@ -358,7 +367,7 @@ internal fun convertToDataFrame( shouldCreateColumnGroup -> DataColumn.createColumnGroup( name = it.columnName, - df = (values as List).concat(), + df = (values as List).concat(), ) kClass.isSubclassOf(Iterable::class) -> diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/convertTo.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/convertTo.kt index aa16906ac3..5cb278efcf 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/convertTo.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/convertTo.kt @@ -184,6 +184,21 @@ class ConvertToTests { @DataSchema data class DataSchemaWithAnyFrame(val dfs: AnyFrame?) + private fun locationsFrame(): DataFrame = + listOf( + Location("Home", Gps(0.0, 0.0)), + Location("Away", null), + null, + ).toDataFrame() + .alsoDebug("locations:") + + private fun gpsFrame(): DataFrame = + listOf( + Gps(0.0, 0.0), + null, + ).toDataFrame() + .alsoDebug("gps:") + @Test fun test() { val df1 = dataFrameOf("a")(1, 2, 3) @@ -202,67 +217,70 @@ class ConvertToTests { } @Test - fun `convert df with AnyFrame to itself`() { - val locationsList = listOf( - Location("Home", Gps(0.0, 0.0)), - Location("Away", null), - null, - ) - val locations = locationsList - .toDataFrame() - .alsoDebug("locations:") + fun `convert df with AnyFrame containing locations to itself`() { + val locations = locationsFrame() - val gpsList = listOf( - Gps(0.0, 0.0), - null, - ) - val gps = gpsList + listOf(DataSchemaWithAnyFrame(locations)) .toDataFrame() - .alsoDebug("gps:") - - val df1 = listOf( - DataSchemaWithAnyFrame(locations), - ).toDataFrame() .alsoDebug("df1:") + .convertTo() + } - df1.convertTo() + @Test + fun `convert df with AnyFrame containing gps to itself`() { + val gps = gpsFrame() - val df2 = listOf( - DataSchemaWithAnyFrame(gps), - ).toDataFrame() + listOf(DataSchemaWithAnyFrame(gps)) + .toDataFrame() .alsoDebug("df2:") + .convertTo() + } - df2.convertTo() + @Test + fun `convert df with preserved AnyFrame containing null and gps to itself`() { + val gps = gpsFrame() - val df3 = listOf( + listOf( DataSchemaWithAnyFrame(null), DataSchemaWithAnyFrame(gps), ).toDataFrame { properties { preserve(DataFrame::class) } } .alsoDebug("df3 before convert:") + .convertTo() + } - df3.convertTo() - - val df4 = listOf( + @Test + fun `convert df with preserved null AnyFrame to itself`() { + listOf( DataSchemaWithAnyFrame(null), ).toDataFrame { properties { preserve(DataFrame::class) } } .alsoDebug("df4 before convert:") + .convertTo() + } - df4.convertTo() + @Test + fun `convert raw df with AnyFrame column to itself`() { + val locations = locationsFrame() + val gps = gpsFrame() - val df5a: DataFrame<*> = dataFrameOf( + val df: DataFrame<*> = dataFrameOf( columnOf(locations, gps, null).named("dfs"), ).alsoDebug("df5a:") - df5a.convertTo() + df.convertTo() + } + + @Test + fun `convert df with preserved mixed AnyFrame values to itself repeatedly`() { + val locations = locationsFrame() + val gps = gpsFrame() - val df5 = listOf( + listOf( DataSchemaWithAnyFrame(null), DataSchemaWithAnyFrame(locations), DataSchemaWithAnyFrame(gps), ).toDataFrame { properties { preserve(DataFrame::class) } } .alsoDebug("df5 before convert:") - - df5.convertTo() + .convertTo() .alsoDebug("df5 after convert:") .convertTo() .alsoDebug("df5 after second convert:") diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt index ddb89cdd59..674da50d5d 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt @@ -4,7 +4,9 @@ import io.kotest.matchers.shouldBe import org.jetbrains.kotlinx.dataframe.AnyRow import org.jetbrains.kotlinx.dataframe.ColumnsScope import org.jetbrains.kotlinx.dataframe.DataColumn +import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow +import org.jetbrains.kotlinx.dataframe.annotations.DataSchema import org.jetbrains.kotlinx.dataframe.api.columnOf import org.jetbrains.kotlinx.dataframe.api.dataFrameOf import org.jetbrains.kotlinx.dataframe.api.default @@ -16,6 +18,7 @@ import org.jetbrains.kotlinx.dataframe.api.groupBy import org.jetbrains.kotlinx.dataframe.api.into import org.jetbrains.kotlinx.dataframe.api.move import org.jetbrains.kotlinx.dataframe.api.pathOf +import org.jetbrains.kotlinx.dataframe.api.print import org.jetbrains.kotlinx.dataframe.api.schema import org.jetbrains.kotlinx.dataframe.api.toCodeString import org.jetbrains.kotlinx.dataframe.api.under @@ -23,9 +26,11 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGenerator import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl import org.jetbrains.kotlinx.dataframe.impl.toCamelCaseByDelimiters +import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema import org.jetbrains.kotlinx.dataframe.testSets.person.BaseTest import org.jetbrains.kotlinx.dataframe.testSets.person.Person import org.junit.Test +import kotlin.reflect.typeOf import kotlin.test.assertEquals class CodeGenerationTests : BaseTest() { @@ -570,6 +575,18 @@ class CodeGenerationTests : BaseTest() { assertEquals(expected, df.generateDataClasses().value) } + @DataSchema + class C(val i: Int) + + @DataSchema + class Schema(val df: DataFrame?) + + @Test + fun extractNullableDataFrameSchema() { + val schema = MarkersExtractor.get().schema + schema.columns["df"] shouldBe ColumnSchema.Value(typeOf?>()) + } + // region Tests for generateX functions @Test