Skip to content

Commit

Permalink
[Compiler plugin] Use annotation to filter out scope properties from …
Browse files Browse the repository at this point in the history
…schema classes
  • Loading branch information
koperagen committed Jun 27, 2024
1 parent b154937 commit 423dca8
Show file tree
Hide file tree
Showing 13 changed files with 175 additions and 27 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,18 @@ public annotation class Import
@Target(AnnotationTarget.PROPERTY)
public annotation class Order(val order: Int)

/**
* For internal use
* Compiler plugin materializes schemas as classes.
* These classes have two kinds of properties:
* 1. Scope properties that only serve as a reference for internal property resolution
* 2. Schema properties that reflect dataframe structure
* Scope properties need
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
* This annotation serves to distinguish between the two where needed
*/
@Target(AnnotationTarget.PROPERTY)
public annotation class ScopeProperty

@Target(AnnotationTarget.FUNCTION)
internal annotation class Check
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
import kotlin.reflect.KClass
Expand Down Expand Up @@ -54,7 +55,8 @@ internal object MarkersExtractor {

private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
val structuralProperties = markerClass.memberProperties.filter { it.findAnnotation<ScopeProperty>() == null }
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
val fieldName = ValidFieldName.of(it.name)
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
val type = it.returnType
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,18 @@ public annotation class Import
@Target(AnnotationTarget.PROPERTY)
public annotation class Order(val order: Int)

/**
* For internal use
* Compiler plugin materializes schemas as classes.
* These classes have two kinds of properties:
* 1. Scope properties that only serve as a reference for internal property resolution
* 2. Schema properties that reflect dataframe structure
* Scope properties need
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
* This annotation serves to distinguish between the two where needed
*/
@Target(AnnotationTarget.PROPERTY)
public annotation class ScopeProperty

@Target(AnnotationTarget.FUNCTION)
internal annotation class Check
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
import kotlin.reflect.KClass
Expand Down Expand Up @@ -54,7 +55,8 @@ internal object MarkersExtractor {

private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
val structuralProperties = markerClass.memberProperties.filter { it.hasAnnotation<ScopeProperty>() }
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
val fieldName = ValidFieldName.of(it.name)
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
val type = it.returnType
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
import org.jetbrains.kotlin.fir.scopes.collectAllProperties
Expand All @@ -21,7 +22,7 @@ class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolu
val symbol = generatedTokenOrNull(functionCall) ?: return emptyList()
return symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
.filterIsInstance<FirPropertySymbol>()
.filter { it.resolvedReturnType.classId?.shortClassName?.asString()?.startsWith("Scope") ?: false }
.filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
.map { it.resolvedReturnType }
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.caches.getValue
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
import org.jetbrains.kotlin.fir.declarations.FirProperty
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
Expand Down Expand Up @@ -51,8 +52,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
}
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
// making them var appeared to be the easiest way to filter
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isVal = false))
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
}
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
val propertyName = Name.identifier(schemaProperty.name)
Expand Down Expand Up @@ -89,7 +89,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se

@OptIn(SymbolInternals::class)
override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
// maybe Init needed not for everything
val destination = mutableSetOf<Name>()
when (classSymbol.fir.callShapeData) {
is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
Expand All @@ -110,28 +109,33 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
resolvedTypeRef: FirResolvedTypeRef,
propertyName: Name,
k: FirClassSymbol<*>,
isVal: Boolean = true,
isScopeProperty: Boolean = false,
order: Int? = null,
): FirProperty {
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type, isVal) {
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
modality = Modality.ABSTRACT
visibility = Visibilities.Public
}.apply {
val annotations = mutableListOf<FirAnnotation>()
if (order != null) {
val orderAnnotation = buildAnnotation {
annotations += buildAnnotation {
annotationTypeRef = buildResolvedTypeRef {
type = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.ORDER_ANNOTATION),
arrayOf(),
isNullable = false
)
type = Names.ORDER_ANNOTATION.defaultType(emptyList())
}
argumentMapping = buildAnnotationArgumentMapping {
mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
}
}
replaceAnnotations(listOf(orderAnnotation))
}
if (isScopeProperty) {
annotations += buildAnnotation {
annotationTypeRef = buildResolvedTypeRef {
type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
}
argumentMapping = buildAnnotationArgumentMapping()
}
}
replaceAnnotations(annotations)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.annotations.Order
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
import kotlin.reflect.KClass

object Names {
Expand All @@ -32,10 +33,12 @@ object Names {
get() = Name.identifier("org.jetbrains.kotlinx.dataframe.annotations")
val INTERPRETABLE_FQNAME: FqName
get() = FqName("org.jetbrains.kotlinx.dataframe.annotations.Interpretable")
val ORDER_ANNOTATION = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier(Order::class.simpleName!!))
private val annotationsPackage = FqName("org.jetbrains.kotlinx.dataframe.annotations")
val ORDER_ANNOTATION = ClassId(annotationsPackage, Name.identifier(Order::class.simpleName!!))
val ORDER_ARGUMENT = Name.identifier(Order::order.name)
val SCOPE_PROPERTY_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ScopeProperty::class.simpleName!!))

val DATA_SCHEMA_CLASS_ID = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier("DataSchema"))
val DATA_SCHEMA_CLASS_ID = ClassId(annotationsPackage, Name.identifier("DataSchema"))
val LIST = ClassId(FqName("kotlin.collections"), Name.identifier("List"))
val DURATION_CLASS_ID = kotlin.time.Duration::class.classId()
val LOCAL_DATE_CLASS_ID = kotlinx.datetime.LocalDate::class.classId()
Expand Down
19 changes: 19 additions & 0 deletions plugins/kotlin-dataframe/testData/box/castTo.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
import org.jetbrains.kotlinx.dataframe.annotations.*
import org.jetbrains.kotlinx.dataframe.api.*
import org.jetbrains.kotlinx.dataframe.io.*
import org.jetbrains.kotlinx.dataframe.*

fun box(): String {
val sample =
@Import DataFrame.readCSV("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")

val organizations = listOf("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
organizations.forEach { organization ->
val df = DataFrame.readCSV(organization).castTo(sample)
println(organizations)
println("Repositories: ${df.count()}")
println("Top 10:")
df.sortBy { stargazers_count.desc() }.take(10).print()
}
return "OK"
}
72 changes: 72 additions & 0 deletions plugins/kotlin-dataframe/testData/box/propertiesOrder.fir.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
FILE: propertiesOrder.kt
public final fun box(): R|kotlin/String| {
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> {
local abstract class Read_16I : R|kotlin/Any| {
@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val full_name: R|kotlin/String|
public get(): R|kotlin/String|

@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(3)) public abstract val topics: R|kotlin/String|
public get(): R|kotlin/String|

@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(4)) public abstract val watchers: R|kotlin/Int|
public get(): R|kotlin/Int|

@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val stargazers_count: R|kotlin/Int|
public get(): R|kotlin/Int|

@R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val html_url: R|java/net/URL|
public get(): R|java/net/URL|

public constructor(): R|<local>/Read_16I|

}

local final class Scope0 : R|kotlin/Any| {
public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.full_name: R|kotlin/String|
public get(): R|kotlin/String|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.full_name: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.topics: R|kotlin/String|
public get(): R|kotlin/String|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.topics: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/String>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.watchers: R|kotlin/Int|
public get(): R|kotlin/Int|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.watchers: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.stargazers_count: R|kotlin/Int|
public get(): R|kotlin/Int|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.stargazers_count: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|

public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Read_16I>|.html_url: R|java/net/URL|
public get(): R|java/net/URL|

public final val R|org/jetbrains/kotlinx/dataframe/ColumnsContainer<<local>/Read_16I>|.html_url: R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|
public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<java/net/URL>|

public constructor(): R|<local>/Scope0|

}

local abstract class Read_16 : R|<local>/Read_16I| {
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
public get(): R|<local>/Scope0|

public constructor(): R|<local>/Read_16|

}

^ @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))
}
)
(this@R|/box|, R|<local>/df|).R|<local>/Scope0.full_name|
^box String(OK)
}
12 changes: 12 additions & 0 deletions plugins/kotlin-dataframe/testData/box/propertiesOrder.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
// FIR_DUMP

import org.jetbrains.kotlinx.dataframe.*
import org.jetbrains.kotlinx.dataframe.io.read
import org.jetbrains.kotlinx.dataframe.api.*
import org.jetbrains.kotlinx.dataframe.annotations.*

fun box(): String {
val df = @Import DataFrame.read("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
df.full_name
return "OK"
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,8 @@ FILE: selectDuringTyping.kt
}

local abstract class ExplodeSchema_94 : R|<local>/ExplodeSchema_94I| {
public abstract var scope0: R|<local>/Scope0|
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
public get(): R|<local>/Scope0|
public set(value: R|<local>/Scope0|): R|kotlin/Unit|

public constructor(): R|<local>/ExplodeSchema_94|

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,21 +87,17 @@ FILE: test.kt
}

local abstract class S_43 : R|<local>/S_43I| {
public abstract var scope0: R|<local>/Scope0|
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
public get(): R|<local>/Scope0|
public set(value: R|<local>/Scope0|): R|kotlin/Unit|

public abstract var scope3: R|<local>/Scope3|
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope3: R|<local>/Scope3|
public get(): R|<local>/Scope3|
public set(value: R|<local>/Scope3|): R|kotlin/Unit|

public abstract var scope2: R|<local>/Scope2|
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope2: R|<local>/Scope2|
public get(): R|<local>/Scope2|
public set(value: R|<local>/Scope2|): R|kotlin/Unit|

public abstract var scope1: R|<local>/Scope1|
@R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope1: R|<local>/Scope1|
public get(): R|<local>/Scope1|
public set(value: R|<local>/Scope1|): R|kotlin/Unit|

public constructor(): R|<local>/S_43|

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,12 @@ public void testAllFilesPresentInBox() {
KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("testData/box"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
}

@Test
@TestMetadata("castTo.kt")
public void testCastTo() {
runTest("testData/box/castTo.kt");
}

@Test
@TestMetadata("columnGroupApi.kt")
public void testColumnGroupApi() {
Expand Down Expand Up @@ -214,6 +220,12 @@ public void testPlatformType() {
runTest("testData/box/platformType.kt");
}

@Test
@TestMetadata("propertiesOrder.kt")
public void testPropertiesOrder() {
runTest("testData/box/propertiesOrder.kt");
}

@Test
@TestMetadata("read.kt")
public void testRead() {
Expand Down

0 comments on commit 423dca8

Please sign in to comment.