From 9839b5c94bf221df551857f2e134c373ae31d9cd Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Fri, 4 Oct 2024 21:30:58 +0300 Subject: [PATCH] [Compiler plugin] Propagate nullability in toDataFrame tree conversion --- .../dataframe/plugin/impl/api/toDataFrame.kt | 34 +++++++++++-------- .../testData/box/toDataFrame_nullableList.kt | 16 +++++++++ .../box/toDataFrame_nullableListSubtree.kt | 27 +++++++++++++++ .../box/toDataFrame_nullableSubtree.kt | 27 +++++++++++++++ ...DataFrameBlackBoxCodegenTestGenerated.java | 18 ++++++++++ 5 files changed, 108 insertions(+), 14 deletions(-) create mode 100644 plugins/kotlin-dataframe/testData/box/toDataFrame_nullableList.kt create mode 100644 plugins/kotlin-dataframe/testData/box/toDataFrame_nullableListSubtree.kt create mode 100644 plugins/kotlin-dataframe/testData/box/toDataFrame_nullableSubtree.kt diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt index 4dada877d..f059f940e 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt @@ -25,7 +25,9 @@ import org.jetbrains.kotlin.fir.symbols.SymbolInternals import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol import org.jetbrains.kotlin.fir.types.ConeClassLikeType +import org.jetbrains.kotlin.fir.types.ConeFlexibleType import org.jetbrains.kotlin.fir.types.ConeKotlinType +import org.jetbrains.kotlin.fir.types.ConeNullability import org.jetbrains.kotlin.fir.types.ConeStarProjection import org.jetbrains.kotlin.fir.types.ConeTypeParameterType import org.jetbrains.kotlin.fir.types.canBeNull @@ -41,15 +43,19 @@ import org.jetbrains.kotlin.fir.types.resolvedType import org.jetbrains.kotlin.fir.types.toRegularClassSymbol import org.jetbrains.kotlin.fir.types.toSymbol import org.jetbrains.kotlin.fir.types.type +import org.jetbrains.kotlin.fir.types.typeContext import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible import org.jetbrains.kotlin.fir.types.withArguments +import org.jetbrains.kotlin.fir.types.withNullability import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.name.StandardClassIds import org.jetbrains.kotlin.name.StandardClassIds.List +import org.jetbrains.kotlin.types.checker.SimpleClassicTypeSystemContext.withNullability import org.jetbrains.kotlinx.dataframe.codeGen.* import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade +import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments @@ -197,7 +203,7 @@ internal fun KotlinTypeFacade.toDataFrame( val preserveClasses = traverseConfiguration.preserveClasses.mapNotNullTo(mutableSetOf()) { it.classId } val preserveProperties = traverseConfiguration.preserveProperties.mapNotNullTo(mutableSetOf()) { it.calleeReference.toResolvedPropertySymbol() } - fun convert(classLike: ConeKotlinType, depth: Int): List { + fun convert(classLike: ConeKotlinType, depth: Int, makeNullable: Boolean): List { val symbol = classLike.toRegularClassSymbol(session) ?: return emptyList() val scope = symbol.unsubstitutedScope(session, ScopeSession(), false, FirResolvePhase.STATUS) val declarations = if (symbol.fir is FirJavaClass) { @@ -260,7 +266,7 @@ internal fun KotlinTypeFacade.toDataFrame( val keepSubtree = depth >= maxDepth && !fieldKind.shouldBeConvertedToColumnGroup && !fieldKind.shouldBeConvertedToFrameColumn if (keepSubtree || returnType.isValueType() || returnType.classId in preserveClasses || it in preserveProperties) { - SimpleDataColumn(name, TypeApproximation(returnType)) + SimpleDataColumn(name, TypeApproximation(returnType.withNullability(ConeNullability.create(makeNullable), session.typeContext))) } else if ( returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection)), session) || returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection), isNullable = true), session) @@ -271,19 +277,15 @@ internal fun KotlinTypeFacade.toDataFrame( else -> session.builtinTypes.nullableAnyType.type } if (type.isValueType()) { - SimpleDataColumn(name, - TypeApproximation( - List.constructClassLikeType( - arrayOf(type), - returnType.isNullable - ) - ) - ) + val columnType = List.constructClassLikeType(arrayOf(type), returnType.isNullable) + .withNullability(ConeNullability.create(makeNullable), session.typeContext) + .wrap() + SimpleDataColumn(name, columnType) } else { - SimpleFrameColumn(name, convert(type, depth + 1)) + SimpleFrameColumn(name, convert(type, depth + 1, makeNullable = false)) } } else { - SimpleColumnGroup(name, convert(returnType, depth + 1)) + SimpleColumnGroup(name, convert(returnType, depth + 1, returnType.isNullable || makeNullable)) } } } @@ -293,8 +295,12 @@ internal fun KotlinTypeFacade.toDataFrame( return when { arg.isStarProjection -> PluginDataFrameSchema.EMPTY else -> { - val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema.EMPTY - val columns = convert(classLike, 0) + val classLike = when (val type = arg.type) { + is ConeClassLikeType -> type + is ConeFlexibleType -> type.upperBound + else -> null + } ?: return PluginDataFrameSchema.EMPTY + val columns = convert(classLike, 0, makeNullable = classLike.isNullable) PluginDataFrameSchema(columns) } } diff --git a/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableList.kt b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableList.kt new file mode 100644 index 000000000..fe197b26b --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableList.kt @@ -0,0 +1,16 @@ +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 D( + val s: String +) + +fun box(): String { + val df1 = listOf(D("bb"), null).toDataFrame() + df1.schema().print() + df1.compileTimeSchema().print() + return "OK" +} diff --git a/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableListSubtree.kt b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableListSubtree.kt new file mode 100644 index 000000000..2666765cf --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableListSubtree.kt @@ -0,0 +1,27 @@ +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 D( + val s: String +) + +class Subtree( + val p: Int, + val l: List, + val ld: List, +) + +class Root(val a: Subtree) + +fun box(): String { + val l = listOf( + Root(Subtree(123, listOf(1), listOf(D("ff")))), + null + ) + val df = l.toDataFrame(maxDepth = 2) + df.compareSchemas(strict = true) + return "OK" +} diff --git a/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableSubtree.kt b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableSubtree.kt new file mode 100644 index 000000000..eda97d929 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/toDataFrame_nullableSubtree.kt @@ -0,0 +1,27 @@ +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 D( + val s: String +) + +class Subtree( + val p: Int, + val l: List, + val ld: List, +) + +class Root(val a: Subtree?) + +fun box(): String { + val l = listOf( + Root(Subtree(123, listOf(1), listOf(D("ff")))), + Root(null) + ) + val df = l.toDataFrame(maxDepth = 2) + df.compareSchemas(strict = true) + return "OK" +} 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 9d35da9f2..b98c00e96 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 @@ -436,6 +436,24 @@ public void testToDataFrame_from() { runTest("testData/box/toDataFrame_from.kt"); } + @Test + @TestMetadata("toDataFrame_nullableList.kt") + public void testToDataFrame_nullableList() { + runTest("testData/box/toDataFrame_nullableList.kt"); + } + + @Test + @TestMetadata("toDataFrame_nullableListSubtree.kt") + public void testToDataFrame_nullableListSubtree() { + runTest("testData/box/toDataFrame_nullableListSubtree.kt"); + } + + @Test + @TestMetadata("toDataFrame_nullableSubtree.kt") + public void testToDataFrame_nullableSubtree() { + runTest("testData/box/toDataFrame_nullableSubtree.kt"); + } + @Test @TestMetadata("toDataFrame_superType.kt") public void testToDataFrame_superType() {