[Compiler plugin] Propagate nullability in toDataFrame tree conversion

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<SimpleCol> {
+    fun convert(classLike: ConeKotlinType, depth: Int, makeNullable: Boolean): List<SimpleCol> {
         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)
         }
     }

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"
+}

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<Int>,
+    val ld: List<D>,
+)
+
+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"
+}

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<Int>,
+    val ld: List<D>,
+)
+
+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"
+}

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() {