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 5458c9de40..cbfb847f53 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 @@ -1,5 +1,6 @@ package org.jetbrains.kotlinx.dataframe.codeGen +import com.squareup.kotlinpoet.asTypeName import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow import org.jetbrains.kotlinx.dataframe.annotations.ColumnName @@ -109,7 +110,11 @@ internal object MarkersExtractor { else -> { fieldType = FieldType.ValueFieldType( - if (nullableProperties) type.toString().toNullable() else type.toString(), + if (nullableProperties) { + type.asTypeName().toString().toNullable() + } else { + type.asTypeName().toString() + }, ) ColumnSchema.Value( if (nullableProperties) type.withNullability(true) else type, diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaProcessorImpl.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaProcessorImpl.kt index 9517fbe078..ab66221a23 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaProcessorImpl.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaProcessorImpl.kt @@ -1,5 +1,6 @@ package org.jetbrains.kotlinx.dataframe.impl.codeGen +import com.squareup.kotlinpoet.asTypeName import org.jetbrains.kotlinx.dataframe.codeGen.FieldType import org.jetbrains.kotlinx.dataframe.codeGen.GeneratedField import org.jetbrains.kotlinx.dataframe.codeGen.Marker @@ -66,7 +67,7 @@ internal class SchemaProcessorImpl( fun getFieldType(columnSchema: ColumnSchema): FieldType = when (columnSchema) { is ColumnSchema.Value -> - FieldType.ValueFieldType(columnSchema.type.toString()) + FieldType.ValueFieldType(columnSchema.type.asTypeName().toString()) is ColumnSchema.Group -> FieldType.GroupFieldType( diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt index 57e13be92c..047ddcd6c5 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt @@ -61,8 +61,8 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames { @Test fun `short functional types are not supported`() { fields.keys.asClue { - fields["d"]!!.renderAccessorFieldType() shouldBe "() -> kotlin.Unit" - fields["d"]!!.renderFieldType() shouldBe "() -> kotlin.Unit" + fields["d"]!!.renderAccessorFieldType() shouldBe "kotlin.Function0" + fields["d"]!!.renderFieldType() shouldBe "kotlin.Function0" } } @@ -112,7 +112,7 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames { @Test fun `functional type column`() { fields.keys.asClue { - fields["d"]!!.renderColumnType() shouldBe "DataColumn<() -> kotlin.Unit>" + fields["d"]!!.renderColumnType() shouldBe "DataColumn>" } } 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 4dada877d7..f059f940e5 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 0000000000..fe197b26b4 --- /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 0000000000..2666765cf5 --- /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 0000000000..eda97d9298 --- /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 2e5d5b5d22..5cc3ffcb5c 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 @@ -418,6 +418,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() {