Skip to content

Commit bb96ffd

Browse files
koperagenkoperagen
authored
Merge pull request #856 from Kotlin/compiler-plugin-bugfixes-1
Compiler plugin bugfixes
2 parents 0a5f526 + 993b5ad commit bb96ffd

File tree

7 files changed

+111
-43
lines changed

7 files changed

+111
-43
lines changed

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -251,11 +251,9 @@ internal fun convertToDataFrame(
251251
val kClass = returnType.classifier as KClass<*>
252252
val fieldKind = returnType.getFieldKind()
253253

254-
val shouldCreateValueCol = (
255-
maxDepth <= 0 &&
256-
!fieldKind.shouldBeConvertedToFrameColumn &&
257-
!fieldKind.shouldBeConvertedToColumnGroup
258-
) ||
254+
val keepSubtree =
255+
maxDepth <= 0 && !fieldKind.shouldBeConvertedToFrameColumn && !fieldKind.shouldBeConvertedToColumnGroup
256+
val shouldCreateValueCol = keepSubtree ||
259257
kClass == Any::class ||
260258
kClass in preserveClasses ||
261259
property in preserveProperties ||

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
66
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
77
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
88
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
9-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
109
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
10+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1111
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
1212
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
1313
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
1414
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
15+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
1516

1617
internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
1718
val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))
@@ -20,14 +21,21 @@ internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
2021
override val Arguments.startingSchema get() = receiver
2122

2223
override fun Arguments.interpret(): PluginDataFrameSchema {
23-
val columns = selector ?: TODO()
24+
val columns = selector ?: object : ColumnsResolver {
25+
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
26+
return df.flatten(includeFrames = false).filter {
27+
val column = it.column
28+
column is SimpleFrameColumn || column is SimpleDataColumn && column.type.isList()
29+
}
30+
}
31+
}
2432
return receiver.explodeImpl(dropEmpty, columns.resolve(receiver).map { ColumnPathApproximation(it.path.path) })
2533
}
2634
}
2735

28-
val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, selector: List<ColumnPathApproximation>?) -> PluginDataFrameSchema
36+
val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, selector: List<ColumnPathApproximation>) -> PluginDataFrameSchema
2937
get() = { dropEmpty, selector ->
30-
val columns = selector ?: TODO()
38+
val columns = selector
3139

3240
val selected: Set<List<String>> = columns.map { it.path }.toSet()
3341

@@ -36,9 +44,7 @@ val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, sel
3644
is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
3745
is SimpleFrameColumn -> column
3846
is SimpleDataColumn -> {
39-
// val nullable = if (dropEmpty) (column.type as TypeApproximationImpl).nullable else true
40-
41-
column.changeType(type = column.type.changeNullability { nullable -> if (dropEmpty) nullable else true })
47+
column.changeType(type = column.type.changeNullability { nullable -> selector.size > 1 || !dropEmpty || nullable })
4248
}
4349
}
4450
}
@@ -61,7 +67,7 @@ val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, sel
6167
column.type.isList() -> column.type.typeArgument()
6268
else -> column.type
6369
}
64-
SimpleDataColumn(column.name, newType)
70+
makeNullable(simpleColumnOf(column.name, newType.type))
6571
} else {
6672
column
6773
}

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

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

33
import org.jetbrains.kotlin.descriptors.EffectiveVisibility
4+
import org.jetbrains.kotlin.fir.FirSession
45
import org.jetbrains.kotlinx.dataframe.plugin.classId
56
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
67
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
8+
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
79
import org.jetbrains.kotlin.fir.declarations.utils.effectiveVisibility
810
import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
911
import org.jetbrains.kotlin.fir.declarations.utils.isStatic
@@ -37,13 +39,16 @@ import org.jetbrains.kotlin.fir.types.isStarProjection
3739
import org.jetbrains.kotlin.fir.types.isSubtypeOf
3840
import org.jetbrains.kotlin.fir.types.resolvedType
3941
import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
42+
import org.jetbrains.kotlin.fir.types.toSymbol
4043
import org.jetbrains.kotlin.fir.types.type
4144
import org.jetbrains.kotlin.fir.types.upperBoundIfFlexible
4245
import org.jetbrains.kotlin.fir.types.withArguments
4346
import org.jetbrains.kotlin.name.ClassId
4447
import org.jetbrains.kotlin.name.FqName
4548
import org.jetbrains.kotlin.name.Name
4649
import org.jetbrains.kotlin.name.StandardClassIds
50+
import org.jetbrains.kotlin.name.StandardClassIds.List
51+
import org.jetbrains.kotlinx.dataframe.codeGen.*
4752
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
4853
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
4954
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
@@ -58,6 +63,9 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
5863
import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
5964
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
6065
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
66+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DATA_ROW_CLASS_ID
67+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DATA_SCHEMA_CLASS_ID
68+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DF_CLASS_ID
6169
import java.util.*
6270

6371
class ToDataFrameDsl : AbstractSchemaModificationInterpreter() {
@@ -160,7 +168,7 @@ class Exclude1 : AbstractInterpreter<Unit>() {
160168
dsl.excludeProperties.addAll(properties.arguments.filterIsInstance<FirCallableReferenceAccess>())
161169
}
162170
}
163-
171+
@Suppress("INVISIBLE_MEMBER")
164172
@OptIn(SymbolInternals::class)
165173
internal fun KotlinTypeFacade.toDataFrame(
166174
maxDepth: Int,
@@ -228,17 +236,17 @@ internal fun KotlinTypeFacade.toDataFrame(
228236
.filterNot { excludedClasses.contains(it.first.resolvedReturnType) }
229237
.filter { it.first.effectiveVisibility == EffectiveVisibility.Public }
230238
.map { (it, name) ->
231-
var resolvedReturnType = it.fir.returnTypeRef.resolveIfJavaType(session, JavaTypeParameterStack.EMPTY, null)
239+
var returnType = it.fir.returnTypeRef.resolveIfJavaType(session, JavaTypeParameterStack.EMPTY, null)
232240
.coneType.upperBoundIfFlexible()
233241

234-
resolvedReturnType = if (resolvedReturnType is ConeTypeParameterType) {
235-
if (resolvedReturnType.canBeNull(session)) {
242+
returnType = if (returnType is ConeTypeParameterType) {
243+
if (returnType.canBeNull(session)) {
236244
session.builtinTypes.nullableAnyType.type
237245
} else {
238246
session.builtinTypes.anyType.type
239247
}
240248
} else {
241-
resolvedReturnType.withArguments {
249+
returnType.withArguments {
242250
val type = it.type
243251
if (type is ConeTypeParameterType) {
244252
session.builtinTypes.nullableAnyType.type
@@ -248,33 +256,34 @@ internal fun KotlinTypeFacade.toDataFrame(
248256
}
249257
}
250258

251-
if (depth >= maxDepth || resolvedReturnType.isValueType() || resolvedReturnType.classId in preserveClasses || it in preserveProperties ) {
252-
SimpleDataColumn(name,
253-
TypeApproximation(resolvedReturnType)
254-
)
259+
val fieldKind = returnType.getFieldKind(session)
260+
261+
val keepSubtree = depth >= maxDepth && !fieldKind.shouldBeConvertedToColumnGroup && !fieldKind.shouldBeConvertedToFrameColumn
262+
if (keepSubtree || returnType.isValueType() || returnType.classId in preserveClasses || it in preserveProperties) {
263+
SimpleDataColumn(name, TypeApproximation(returnType))
255264
} else if (
256-
resolvedReturnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection)), session) ||
257-
resolvedReturnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection), isNullable = true), session)
265+
returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection)), session) ||
266+
returnType.isSubtypeOf(StandardClassIds.Iterable.constructClassLikeType(arrayOf(ConeStarProjection), isNullable = true), session)
258267
) {
259-
val type: ConeKotlinType = when (val typeArgument = resolvedReturnType.typeArguments[0]) {
268+
val type: ConeKotlinType = when (val typeArgument = returnType.typeArguments[0]) {
260269
is ConeKotlinType -> typeArgument
261270
ConeStarProjection -> session.builtinTypes.nullableAnyType.type
262271
else -> session.builtinTypes.nullableAnyType.type
263272
}
264273
if (type.isValueType()) {
265274
SimpleDataColumn(name,
266275
TypeApproximation(
267-
StandardClassIds.List.constructClassLikeType(
276+
List.constructClassLikeType(
268277
arrayOf(type),
269-
resolvedReturnType.isNullable
278+
returnType.isNullable
270279
)
271280
)
272281
)
273282
} else {
274283
SimpleFrameColumn(name, convert(type, depth + 1))
275284
}
276285
} else {
277-
SimpleColumnGroup(name, convert(resolvedReturnType, depth + 1))
286+
SimpleColumnGroup(name, convert(returnType, depth + 1))
278287
}
279288
}
280289
}
@@ -291,6 +300,21 @@ internal fun KotlinTypeFacade.toDataFrame(
291300
}
292301
}
293302

303+
// org.jetbrains.kotlinx.dataframe.codeGen.getFieldKind
304+
@Suppress("INVISIBLE_MEMBER")
305+
private fun ConeKotlinType.getFieldKind(session: FirSession) = when {
306+
classId == DF_CLASS_ID -> Frame
307+
classId == List && typeArguments[0].type.hasAnnotation(DATA_SCHEMA_CLASS_ID, session) -> ListToFrame
308+
classId == DATA_ROW_CLASS_ID -> Group
309+
hasAnnotation(DATA_SCHEMA_CLASS_ID, session) -> ObjectToGroup
310+
else -> Default
311+
}
312+
313+
314+
private fun ConeKotlinType?.hasAnnotation(id: ClassId, session: FirSession) =
315+
this?.toSymbol(session)?.hasAnnotation(id, session) == true
316+
317+
294318
class CreateDataFrameDslImplApproximation {
295319
val configuration: CreateDataFrameConfiguration = CreateDataFrameConfiguration()
296320
val columns: MutableList<SimpleCol> = mutableListOf()

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -325,13 +325,13 @@ fun KotlinTypeFacade.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): P
325325
fun KotlinTypeFacade.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType): PluginDataFrameSchema {
326326
val symbol = coneClassLikeType.toSymbol(session) as? FirRegularClassSymbol ?: return PluginDataFrameSchema(emptyList())
327327
val declarationSymbols = if (symbol.isLocal && symbol.resolvedSuperTypes.firstOrNull() != session.builtinTypes.anyType.type) {
328-
val rootSchemaSymbol = symbol.resolvedSuperTypes.first().toSymbol(session) as FirRegularClassSymbol
329-
rootSchemaSymbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
328+
val rootSchemaSymbol = symbol.resolvedSuperTypes.first().toSymbol(session) as? FirRegularClassSymbol
329+
rootSchemaSymbol?.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
330330
} else {
331331
symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)
332332
}.let { scope ->
333-
val names = scope.getCallableNames()
334-
names.flatMap { scope.getProperties(it) }
333+
val names = scope?.getCallableNames() ?: emptySet()
334+
names.flatMap { scope?.getProperties(it) ?: emptyList() }
335335
}
336336

337337
val mapping = symbol.typeParameterSymbols

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

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,23 @@ import org.jetbrains.kotlinx.dataframe.annotations.*
33
import org.jetbrains.kotlinx.dataframe.api.*
44
import org.jetbrains.kotlinx.dataframe.io.*
55

6-
@DataSchema
7-
interface ExplodeSchema {
8-
val timestamps: List<Int>
9-
}
10-
11-
fun explode(df: DataFrame<ExplodeSchema>) {
12-
val res = df.explode { timestamps }
13-
val col: DataColumn<Int> = res.timestamps
14-
}
15-
166
fun box(): String {
17-
val df = dataFrameOf("timestamps")(listOf(100, 113, 140), listOf(400, 410, 453)).cast<ExplodeSchema>()
7+
val df = dataFrameOf("timestamps")(listOf(100, 113, 140), listOf(400, 410, 453))
188
val df1 = df.explode { timestamps }
199
val timestamps: DataColumn<Int> = df1.timestamps
2010
timestamps.print()
11+
12+
13+
val df2 = dataFrameOf("a", "b")(listOf(100, 113, 140), listOf(400, 410))
14+
val df3 = df2.explode { a and b }
15+
// exploding multiple columns will introduce nulls
16+
df3.print()
17+
// compiler needs to play safe and make both selected columns nullable
18+
df3.compileTimeSchema().columns.let {
19+
assert(it["a"]!!.nullable)
20+
assert(it["b"]!!.nullable)
21+
}
22+
// compile time schema is still compatible with runtime
23+
df3.compareSchemas()
2124
return "OK"
2225
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
@file:Suppress("UnusedImport")
2+
3+
import org.jetbrains.kotlinx.dataframe.*
4+
import org.jetbrains.kotlinx.dataframe.annotations.*
5+
import org.jetbrains.kotlinx.dataframe.api.*
6+
import org.jetbrains.kotlinx.dataframe.io.*
7+
8+
@DataSchema
9+
data class Person(val firstName: String, val lastName: String, val age: Int, val city: String?)
10+
11+
@DataSchema
12+
data class Group(val id: String, val participants: List<Person>)
13+
14+
fun box(): String {
15+
val df = listOf(
16+
Group("1", listOf(
17+
Person("Alice", "Cooper", 15, "London"),
18+
Person("Bob", "Dylan", 45, "Dubai")
19+
)),
20+
Group("2", listOf(
21+
Person("Charlie", "Daniels", 20, "Moscow"),
22+
Person("Charlie", "Chaplin", 40, "Milan"),
23+
)),
24+
).toDataFrame(maxDepth = 0)
25+
26+
val groups = df.add("groupSize") {
27+
// interestingly, this is still a DataFrame. Because @DataSchema Person
28+
participants.rowsCount()
29+
}
30+
return "OK"
31+
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,12 @@ public void testToDataFrame_column() {
388388
runTest("testData/box/toDataFrame_column.kt");
389389
}
390390

391+
@Test
392+
@TestMetadata("toDataFrame_dataSchema.kt")
393+
public void testToDataFrame_dataSchema() {
394+
runTest("testData/box/toDataFrame_dataSchema.kt");
395+
}
396+
391397
@Test
392398
@TestMetadata("toDataFrame_dsl.kt")
393399
public void testToDataFrame_dsl() {

0 commit comments

Comments
 (0)