Skip to content

Commit b34cca0

Browse files
koperagenkoperagen
authored
Merge ebcb374 into fea737e 
2 parents fea737e + ebcb374 commit b34cca0

File tree

8 files changed

+119
-19
lines changed

8 files changed

+119
-19
lines changed

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.codeGen
22

3+
import com.squareup.kotlinpoet.asTypeName
34
import org.jetbrains.kotlinx.dataframe.DataFrame
45
import org.jetbrains.kotlinx.dataframe.DataRow
56
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
@@ -109,7 +110,11 @@ internal object MarkersExtractor {
109110

110111
else -> {
111112
fieldType = FieldType.ValueFieldType(
112-
if (nullableProperties) type.toString().toNullable() else type.toString(),
113+
if (nullableProperties) {
114+
type.asTypeName().toString().toNullable()
115+
} else {
116+
type.asTypeName().toString()
117+
},
113118
)
114119
ColumnSchema.Value(
115120
if (nullableProperties) type.withNullability(true) else type,

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaProcessorImpl.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.impl.codeGen
22

3+
import com.squareup.kotlinpoet.asTypeName
34
import org.jetbrains.kotlinx.dataframe.codeGen.FieldType
45
import org.jetbrains.kotlinx.dataframe.codeGen.GeneratedField
56
import org.jetbrains.kotlinx.dataframe.codeGen.Marker
@@ -66,7 +67,7 @@ internal class SchemaProcessorImpl(
6667
fun getFieldType(columnSchema: ColumnSchema): FieldType =
6768
when (columnSchema) {
6869
is ColumnSchema.Value ->
69-
FieldType.ValueFieldType(columnSchema.type.toString())
70+
FieldType.ValueFieldType(columnSchema.type.asTypeName().toString())
7071

7172
is ColumnSchema.Group ->
7273
FieldType.GroupFieldType(

‎core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,8 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
6161
@Test
6262
fun `short functional types are not supported`() {
6363
fields.keys.asClue {
64-
fields["d"]!!.renderAccessorFieldType() shouldBe "() -> kotlin.Unit"
65-
fields["d"]!!.renderFieldType() shouldBe "() -> kotlin.Unit"
64+
fields["d"]!!.renderAccessorFieldType() shouldBe "kotlin.Function0<kotlin.Unit>"
65+
fields["d"]!!.renderFieldType() shouldBe "kotlin.Function0<kotlin.Unit>"
6666
}
6767
}
6868

@@ -112,7 +112,7 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
112112
@Test
113113
fun `functional type column`() {
114114
fields.keys.asClue {
115-
fields["d"]!!.renderColumnType() shouldBe "DataColumn<() -> kotlin.Unit>"
115+
fields["d"]!!.renderColumnType() shouldBe "DataColumn<kotlin.Function0<kotlin.Unit>>"
116116
}
117117
}
118118

‎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
@@ -418,6 +418,24 @@ public void testToDataFrame_from() {
418418
runTest("testData/box/toDataFrame_from.kt");
419419
}
420420

421+
@Test
422+
@TestMetadata("toDataFrame_nullableList.kt")
423+
public void testToDataFrame_nullableList() {
424+
runTest("testData/box/toDataFrame_nullableList.kt");
425+
}
426+
427+
@Test
428+
@TestMetadata("toDataFrame_nullableListSubtree.kt")
429+
public void testToDataFrame_nullableListSubtree() {
430+
runTest("testData/box/toDataFrame_nullableListSubtree.kt");
431+
}
432+
433+
@Test
434+
@TestMetadata("toDataFrame_nullableSubtree.kt")
435+
public void testToDataFrame_nullableSubtree() {
436+
runTest("testData/box/toDataFrame_nullableSubtree.kt");
437+
}
438+
421439
@Test
422440
@TestMetadata("toDataFrame_superType.kt")
423441
public void testToDataFrame_superType() {

0 commit comments

Comments
 (0)