Skip to content

Commit 314ff41

Browse files
koperagenkoperagen
committed
[Compiler plugin] Sync toDataFrame implementations
1 parent 8e1084d commit 314ff41

File tree

4 files changed

+79
-20
lines changed

4 files changed

+79
-20
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/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/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)