Skip to content

Commit 9839b5c

Browse files
koperagenkoperagen
committed
[Compiler plugin] Propagate nullability in toDataFrame tree conversion
1 parent 439f65d commit 9839b5c

File tree

5 files changed

+108
-14
lines changed

5 files changed

+108
-14
lines changed

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@ import org.jetbrains.kotlin.fir.symbols.SymbolInternals
2525
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
2626
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
2727
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
28+
import org.jetbrains.kotlin.fir.types.ConeFlexibleType
2829
import org.jetbrains.kotlin.fir.types.ConeKotlinType
30+
import org.jetbrains.kotlin.fir.types.ConeNullability
2931
import org.jetbrains.kotlin.fir.types.ConeStarProjection
3032
import org.jetbrains.kotlin.fir.types.ConeTypeParameterType
3133
import org.jetbrains.kotlin.fir.types.canBeNull
@@ -41,15 +43,19 @@ import org.jetbrains.kotlin.fir.types.resolvedType
4143
import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
4244
import org.jetbrains.kotlin.fir.types.toSymbol
4345
import org.jetbrains.kotlin.fir.types.type
46+
import org.jetbrains.kotlin.fir.types.typeContext
4447
import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
4548
import org.jetbrains.kotlin.fir.types.withArguments
49+
import org.jetbrains.kotlin.fir.types.withNullability
4650
import org.jetbrains.kotlin.name.ClassId
4751
import org.jetbrains.kotlin.name.FqName
4852
import org.jetbrains.kotlin.name.Name
4953
import org.jetbrains.kotlin.name.StandardClassIds
5054
import org.jetbrains.kotlin.name.StandardClassIds.List
55+
import org.jetbrains.kotlin.types.checker.SimpleClassicTypeSystemContext.withNullability
5156
import org.jetbrains.kotlinx.dataframe.codeGen.*
5257
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
58+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
5359
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
5460
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
5561
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
@@ -197,7 +203,7 @@ internal fun KotlinTypeFacade.toDataFrame(
197203
val preserveClasses = traverseConfiguration.preserveClasses.mapNotNullTo(mutableSetOf()) { it.classId }
198204
val preserveProperties = traverseConfiguration.preserveProperties.mapNotNullTo(mutableSetOf()) { it.calleeReference.toResolvedPropertySymbol() }
199205

200-
fun convert(classLike: ConeKotlinType, depth: Int): List<SimpleCol> {
206+
fun convert(classLike: ConeKotlinType, depth: Int, makeNullable: Boolean): List<SimpleCol> {
201207
val symbol = classLike.toRegularClassSymbol(session) ?: return emptyList()
202208
val scope = symbol.unsubstitutedScope(session, ScopeSession(), false, FirResolvePhase.STATUS)
203209
val declarations = if (symbol.fir is FirJavaClass) {
@@ -260,7 +266,7 @@ internal fun KotlinTypeFacade.toDataFrame(
260266

261267
val keepSubtree = depth >= maxDepth && !fieldKind.shouldBeConvertedToColumnGroup && !fieldKind.shouldBeConvertedToFrameColumn
262268
if (keepSubtree || returnType.isValueType() || returnType.classId in preserveClasses || it in preserveProperties) {
263-
SimpleDataColumn(name, TypeApproximation(returnType))
269+
SimpleDataColumn(name, TypeApproximation(returnType.withNullability(ConeNullability.create(makeNullable), session.typeContext)))
264270
} else if (
265271
returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection)), session) ||
266272
returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection), isNullable = true), session)
@@ -271,19 +277,15 @@ internal fun KotlinTypeFacade.toDataFrame(
271277
else -> session.builtinTypes.nullableAnyType.type
272278
}
273279
if (type.isValueType()) {
274-
SimpleDataColumn(name,
275-
TypeApproximation(
276-
List.constructClassLikeType(
277-
arrayOf(type),
278-
returnType.isNullable
279-
)
280-
)
281-
)
280+
val columnType = List.constructClassLikeType(arrayOf(type), returnType.isNullable)
281+
.withNullability(ConeNullability.create(makeNullable), session.typeContext)
282+
.wrap()
283+
SimpleDataColumn(name, columnType)
282284
} else {
283-
SimpleFrameColumn(name, convert(type, depth + 1))
285+
SimpleFrameColumn(name, convert(type, depth + 1, makeNullable = false))
284286
}
285287
} else {
286-
SimpleColumnGroup(name, convert(returnType, depth + 1))
288+
SimpleColumnGroup(name, convert(returnType, depth + 1, returnType.isNullable || makeNullable))
287289
}
288290
}
289291
}
@@ -293,8 +295,12 @@ internal fun KotlinTypeFacade.toDataFrame(
293295
return when {
294296
arg.isStarProjection -> PluginDataFrameSchema.EMPTY
295297
else -> {
296-
val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema.EMPTY
297-
val columns = convert(classLike, 0)
298+
val classLike = when (val type = arg.type) {
299+
is ConeClassLikeType -> type
300+
is ConeFlexibleType -> type.upperBound
301+
else -> null
302+
} ?: return PluginDataFrameSchema.EMPTY
303+
val columns = convert(classLike, 0, makeNullable = classLike.isNullable)
298304
PluginDataFrameSchema(columns)
299305
}
300306
}

‎plugins/kotlin-dataframe/testData/box/toDataFrame_nullableList.kt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
@DataSchema
7+
data class D(
8+
val s: String
9+
)
10+
11+
fun box(): String {
12+
val df1 = listOf(D("bb"), null).toDataFrame()
13+
df1.schema().print()
14+
df1.compileTimeSchema().print()
15+
return "OK"
16+
}

‎plugins/kotlin-dataframe/testData/box/toDataFrame_nullableListSubtree.kt

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
@DataSchema
7+
data class D(
8+
val s: String
9+
)
10+
11+
class Subtree(
12+
val p: Int,
13+
val l: List<Int>,
14+
val ld: List<D>,
15+
)
16+
17+
class Root(val a: Subtree)
18+
19+
fun box(): String {
20+
val l = listOf(
21+
Root(Subtree(123, listOf(1), listOf(D("ff")))),
22+
null
23+
)
24+
val df = l.toDataFrame(maxDepth = 2)
25+
df.compareSchemas(strict = true)
26+
return "OK"
27+
}

‎plugins/kotlin-dataframe/testData/box/toDataFrame_nullableSubtree.kt

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
@DataSchema
7+
data class D(
8+
val s: String
9+
)
10+
11+
class Subtree(
12+
val p: Int,
13+
val l: List<Int>,
14+
val ld: List<D>,
15+
)
16+
17+
class Root(val a: Subtree?)
18+
19+
fun box(): String {
20+
val l = listOf(
21+
Root(Subtree(123, listOf(1), listOf(D("ff")))),
22+
Root(null)
23+
)
24+
val df = l.toDataFrame(maxDepth = 2)
25+
df.compareSchemas(strict = true)
26+
return "OK"
27+
}

‎plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,24 @@ public void testToDataFrame_from() {
436436
runTest("testData/box/toDataFrame_from.kt");
437437
}
438438

439+
@Test
440+
@TestMetadata("toDataFrame_nullableList.kt")
441+
public void testToDataFrame_nullableList() {
442+
runTest("testData/box/toDataFrame_nullableList.kt");
443+
}
444+
445+
@Test
446+
@TestMetadata("toDataFrame_nullableListSubtree.kt")
447+
public void testToDataFrame_nullableListSubtree() {
448+
runTest("testData/box/toDataFrame_nullableListSubtree.kt");
449+
}
450+
451+
@Test
452+
@TestMetadata("toDataFrame_nullableSubtree.kt")
453+
public void testToDataFrame_nullableSubtree() {
454+
runTest("testData/box/toDataFrame_nullableSubtree.kt");
455+
}
456+
439457
@Test
440458
@TestMetadata("toDataFrame_superType.kt")
441459
public void testToDataFrame_superType() {

0 commit comments

Comments
 (0)