Skip to content

Commit 423dca8

Browse files
koperagenkoperagen
committed
[Compiler plugin] Use annotation to filter out scope properties from schema classes
1 parent b154937 commit 423dca8

File tree

13 files changed

+175
-27
lines changed

13 files changed

+175
-27
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,18 @@ public annotation class Import
2828
@Target(AnnotationTarget.PROPERTY)
2929
public annotation class Order(val order: Int)
3030

31+
/**
32+
* For internal use
33+
* Compiler plugin materializes schemas as classes.
34+
* These classes have two kinds of properties:
35+
* 1. Scope properties that only serve as a reference for internal property resolution
36+
* 2. Schema properties that reflect dataframe structure
37+
* Scope properties need
38+
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
39+
* This annotation serves to distinguish between the two where needed
40+
*/
41+
@Target(AnnotationTarget.PROPERTY)
42+
public annotation class ScopeProperty
43+
3144
@Target(AnnotationTarget.FUNCTION)
3245
internal annotation class Check

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.DataRow
55
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
66
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
78
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
89
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
910
import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
5455

5556
private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
5657
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
57-
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
58+
val structuralProperties = markerClass.memberProperties.filter { it.findAnnotation<ScopeProperty>() == null }
59+
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
5860
val fieldName = ValidFieldName.of(it.name)
5961
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
6062
val type = it.returnType

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,18 @@ public annotation class Import
2828
@Target(AnnotationTarget.PROPERTY)
2929
public annotation class Order(val order: Int)
3030

31+
/**
32+
* For internal use
33+
* Compiler plugin materializes schemas as classes.
34+
* These classes have two kinds of properties:
35+
* 1. Scope properties that only serve as a reference for internal property resolution
36+
* 2. Schema properties that reflect dataframe structure
37+
* Scope properties need
38+
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
39+
* This annotation serves to distinguish between the two where needed
40+
*/
41+
@Target(AnnotationTarget.PROPERTY)
42+
public annotation class ScopeProperty
43+
3144
@Target(AnnotationTarget.FUNCTION)
3245
internal annotation class Check

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.DataRow
55
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
66
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
78
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
89
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
910
import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
5455

5556
private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
5657
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
57-
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
58+
val structuralProperties = markerClass.memberProperties.filter { it.hasAnnotation<ScopeProperty>() }
59+
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
5860
val fieldName = ValidFieldName.of(it.name)
5961
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
6062
val type = it.returnType

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
33
import org.jetbrains.kotlin.fir.FirSession
44
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
55
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
6+
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
67
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
78
import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
89
import org.jetbrains.kotlin.fir.scopes.collectAllProperties
@@ -21,7 +22,7 @@ class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolu
2122
val symbol = generatedTokenOrNull(functionCall) ?: return emptyList()
2223
return symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
2324
.filterIsInstance<FirPropertySymbol>()
24-
.filter { it.resolvedReturnType.classId?.shortClassName?.asString()?.startsWith("Scope") ?: false }
25+
.filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
2526
.map { it.resolvedReturnType }
2627
}
2728

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

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.caches.getValue
1212
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
1313
import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
1414
import org.jetbrains.kotlin.fir.declarations.FirProperty
15+
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
1516
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
1617
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
1718
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
@@ -51,8 +52,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
5152
}
5253
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
5354
val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
54-
// making them var appeared to be the easiest way to filter
55-
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isVal = false))
55+
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
5656
}
5757
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
5858
val propertyName = Name.identifier(schemaProperty.name)
@@ -89,7 +89,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
8989

9090
@OptIn(SymbolInternals::class)
9191
override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
92-
// maybe Init needed not for everything
9392
val destination = mutableSetOf<Name>()
9493
when (classSymbol.fir.callShapeData) {
9594
is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
@@ -110,28 +109,33 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
110109
resolvedTypeRef: FirResolvedTypeRef,
111110
propertyName: Name,
112111
k: FirClassSymbol<*>,
113-
isVal: Boolean = true,
112+
isScopeProperty: Boolean = false,
114113
order: Int? = null,
115114
): FirProperty {
116-
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type, isVal) {
115+
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
117116
modality = Modality.ABSTRACT
118117
visibility = Visibilities.Public
119118
}.apply {
119+
val annotations = mutableListOf<FirAnnotation>()
120120
if (order != null) {
121-
val orderAnnotation = buildAnnotation {
121+
annotations += buildAnnotation {
122122
annotationTypeRef = buildResolvedTypeRef {
123-
type = ConeClassLikeTypeImpl(
124-
ConeClassLikeLookupTagImpl(Names.ORDER_ANNOTATION),
125-
arrayOf(),
126-
isNullable = false
127-
)
123+
type = Names.ORDER_ANNOTATION.defaultType(emptyList())
128124
}
129125
argumentMapping = buildAnnotationArgumentMapping {
130126
mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
131127
}
132128
}
133-
replaceAnnotations(listOf(orderAnnotation))
134129
}
130+
if (isScopeProperty) {
131+
annotations += buildAnnotation {
132+
annotationTypeRef = buildResolvedTypeRef {
133+
type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
134+
}
135+
argumentMapping = buildAnnotationArgumentMapping()
136+
}
137+
}
138+
replaceAnnotations(annotations)
135139
}
136140
}
137141

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.name.ClassId
99
import org.jetbrains.kotlin.name.FqName
1010
import org.jetbrains.kotlin.name.Name
1111
import org.jetbrains.kotlinx.dataframe.annotations.Order
12+
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
1213
import kotlin.reflect.KClass
1314

1415
object Names {
@@ -32,10 +33,12 @@ object Names {
3233
get() = Name.identifier("org.jetbrains.kotlinx.dataframe.annotations")
3334
val INTERPRETABLE_FQNAME: FqName
3435
get() = FqName("org.jetbrains.kotlinx.dataframe.annotations.Interpretable")
35-
val ORDER_ANNOTATION = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier(Order::class.simpleName!!))
36+
private val annotationsPackage = FqName("org.jetbrains.kotlinx.dataframe.annotations")
37+
val ORDER_ANNOTATION = ClassId(annotationsPackage, Name.identifier(Order::class.simpleName!!))
3638
val ORDER_ARGUMENT = Name.identifier(Order::order.name)
39+
val SCOPE_PROPERTY_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ScopeProperty::class.simpleName!!))
3740

38-
val DATA_SCHEMA_CLASS_ID = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier("DataSchema"))
41+
val DATA_SCHEMA_CLASS_ID = ClassId(annotationsPackage, Name.identifier("DataSchema"))
3942
val LIST = ClassId(FqName("kotlin.collections"), Name.identifier("List"))
4043
val DURATION_CLASS_ID = kotlin.time.Duration::class.classId()
4144
val LOCAL_DATE_CLASS_ID = kotlinx.datetime.LocalDate::class.classId()

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import org.jetbrains.kotlinx.dataframe.annotations.*
2+
import org.jetbrains.kotlinx.dataframe.api.*
3+
import org.jetbrains.kotlinx.dataframe.io.*
4+
import org.jetbrains.kotlinx.dataframe.*
5+
6+
fun box(): String {
7+
val sample =
8+
@Import DataFrame.readCSV("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
9+
10+
val organizations = listOf("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
11+
organizations.forEach { organization ->
12+
val df = DataFrame.readCSV(organization).castTo(sample)
13+
println(organizations)
14+
println("Repositories: ${df.count()}")
15+
println("Top 10:")
16+
df.sortBy { stargazers_count.desc() }.take(10).print()
17+
}
18+
return "OK"
19+
}

‎plugins/kotlin-dataframe/testData/box/propertiesOrder.fir.txt

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
FILE: propertiesOrder.kt
2+
public final fun box(): R|kotlin/String| {
3+
lval df: R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Read_16>| = Q|org/jetbrains/kotlinx/dataframe/DataFrame|.R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/DataFrame.Companion|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Read_16>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/DataFrame.Companion|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Read_16>| <inline=Inline, kind=EXACTLY_ONCE> {
4+
local abstract class Read_16I : R|kotlin/Any| {
5+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val full_name: R|kotlin/String|
6+
public get(): R|kotlin/String|
7+
8+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(3)) public abstract val topics: R|kotlin/String|
9+
public get(): R|kotlin/String|
10+
11+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(4)) public abstract val watchers: R|kotlin/Int|
12+
public get(): R|kotlin/Int|
13+
14+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val stargazers_count: R|kotlin/Int|
15+
public get(): R|kotlin/Int|
16+
17+
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val html_url: R|java/net/URL|
18+
public get(): R|java/net/URL|
19+
20+
public constructor(): R|<local>/Read_16I|
21+
22+
}
23+
24+
local final class Scope0 : R|kotlin/Any| {
25+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.full_name: R|kotlin/String|
26+
public get(): R|kotlin/String|
27+
28+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.full_name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
29+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
30+
31+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.topics: R|kotlin/String|
32+
public get(): R|kotlin/String|
33+
34+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.topics: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
35+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
36+
37+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.watchers: R|kotlin/Int|
38+
public get(): R|kotlin/Int|
39+
40+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.watchers: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
41+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
42+
43+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.stargazers_count: R|kotlin/Int|
44+
public get(): R|kotlin/Int|
45+
46+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.stargazers_count: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
47+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
48+
49+
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.html_url: R|java/net/URL|
50+
public get(): R|java/net/URL|
51+
52+
public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.html_url: R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|
53+
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|
54+
55+
public constructor(): R|<local>/Scope0|
56+
57+
}
58+
59+
local abstract class Read_16 : R|<local>/Read_16I| {
60+
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
61+
public get(): R|<local>/Scope0|
62+
63+
public constructor(): R|<local>/Read_16|
64+
65+
}
66+
67+
^ @R|org/jetbrains/kotlinx/dataframe/annotations/Import|() R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/io/read|(String(https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv))
68+
}
69+
)
70+
(this@R|/box|, R|<local>/df|).R|<local>/Scope0.full_name|
71+
^box String(OK)
72+
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
// FIR_DUMP
2+
3+
import org.jetbrains.kotlinx.dataframe.*
4+
import org.jetbrains.kotlinx.dataframe.io.read
5+
import org.jetbrains.kotlinx.dataframe.api.*
6+
import org.jetbrains.kotlinx.dataframe.annotations.*
7+
8+
fun box(): String {
9+
val df = @Import DataFrame.read("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
10+
df.full_name
11+
return "OK"
12+
}

0 commit comments

Comments
 (0)