[Compiler plugin] Use name from ColumnName in user declared data schemas to get columns

Author: koperagen

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

@@ -2,4 +2,4 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
 
 import org.jetbrains.kotlin.GeneratedDeclarationKey
 
-object DataFramePlugin : GeneratedDeclarationKey()
+class DataFramePlugin(val columnName: String?) : GeneratedDeclarationKey()

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

@@ -2,11 +2,13 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
 import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
 import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
@@ -27,6 +29,7 @@ import org.jetbrains.kotlin.fir.types.toSymbol
 import org.jetbrains.kotlin.fir.types.toTypeProjection
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 
@@ -77,7 +80,9 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
         return when (owner) {
             null -> fields.filter { it.callableId == callableId }.flatMap { (owner, property, callableId) ->
                 var resolvedReturnTypeRef = property.resolvedReturnTypeRef
-
+                val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation ->
+                    (annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
+                }
                 val propertyName = property.name
                 val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT)
 
@@ -109,7 +114,9 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
                         ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
                         typeArguments = arrayOf(marker),
                         isNullable = false
-                    ), propertyName = propertyName,
+                    ),
+                    propertyName = propertyName,
+                    columnName = columnName,
                     returnTypeRef = resolvedReturnTypeRef
                 )
 
@@ -132,6 +139,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
                         isNullable = false
                     ),
                     propertyName = propertyName,
+                    columnName = columnName,
                     returnTypeRef = columnReturnType
                 )
                 listOf(rowExtension.symbol, columnsContainerExtension.symbol)

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

@@ -134,7 +134,7 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
 
     override fun visitClass(declaration: IrClass): IrStatement {
         val origin = declaration.origin
-        return if (origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey == DataFramePlugin) {
+        return if (origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin) {
             declaration.transformChildren(this, null)
             declaration
         } else {
@@ -144,7 +144,7 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
 
     override fun visitConstructor(declaration: IrConstructor): IrStatement {
         val origin = declaration.origin
-        if (!(origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey == DataFramePlugin)) return declaration
+        if (!(origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is TokenGenerator.Key)) return declaration
         declaration.body = generateBodyForDefaultConstructor(declaration)
         return declaration
     }
@@ -176,7 +176,8 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
     @OptIn(UnsafeDuringIrConstructionAPI::class)
     override fun visitProperty(declaration: IrProperty): IrStatement {
         val origin = declaration.origin
-        if (!(origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey == DataFramePlugin)) {
+        val pluginKey = (origin as? IrDeclarationOrigin.GeneratedByPlugin)?.pluginKey as? DataFramePlugin
+        if (pluginKey == null) {
             declaration.transformChildren(this, null)
             return declaration
         }
@@ -214,7 +215,8 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
         val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also {
             val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!!
             it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol)
-            it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, declaration.name.identifier))
+            val columName = pluginKey.columnName ?: declaration.name.identifier
+            it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName))
         }
 
         val typeOp = IrTypeOperatorCallImpl(-1, -1, returnType, IrTypeOperator.CAST, returnType, call)
@@ -229,7 +231,7 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
     @OptIn(UnsafeDuringIrConstructionAPI::class)
     override fun visitErrorCallExpression(expression: IrErrorCallExpression): IrExpression {
         val origin = (expression.type.classifierOrNull?.owner as? IrClass)?.origin ?: return expression
-        val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey == DataFramePlugin
+        val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin
         val scopeReference = expression.type.classFqName?.shortName()?.asString()?.startsWith("Scope") ?: false
         if (!(fromPlugin || scopeReference)) {
             return expression

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

@@ -140,6 +140,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
     }
 
     override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        return listOf(createConstructor(context.owner, DataFramePlugin, isPrimary = true).symbol)
+        return listOf(createConstructor(context.owner, Key, isPrimary = true).symbol)
     }
 }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.utils
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
 import org.jetbrains.kotlinx.dataframe.annotations.Order
 import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
 import kotlin.reflect.KClass
@@ -37,6 +38,8 @@ object Names {
     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 COLUMN_NAME_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ColumnName::class.simpleName!!))
+    val COLUMN_NAME_ARGUMENT = Name.identifier(ColumnName::name.name)
 
     val DATA_SCHEMA_CLASS_ID = ClassId(annotationsPackage, Name.identifier("DataSchema"))
     val LIST = ClassId(FqName("kotlin.collections"), Name.identifier("List"))

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

@@ -29,14 +29,15 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
     receiverType: ConeClassLikeTypeImpl,
     propertyName: Name,
     returnTypeRef: FirResolvedTypeRef,
+    columnName: String? = null,
     symbol: FirClassSymbol<*>? = null,
     effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public
 ): FirProperty {
     val firPropertySymbol = FirPropertySymbol(callableId)
     return buildProperty {
         moduleData = session.moduleData
         resolvePhase = FirResolvePhase.BODY_RESOLVE
-        origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
+        origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
         status = FirResolvedDeclarationStatusImpl(
                 Visibilities.Public,
                 Modality.FINAL,
@@ -66,7 +67,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
         getter = buildPropertyAccessor {
             moduleData = session.moduleData
             resolvePhase = FirResolvePhase.BODY_RESOLVE
-            origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
+            origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
             this.returnTypeRef = returnTypeRef
             dispatchReceiverType = receiverType
             this.symbol = firPropertyAccessorSymbol

plugins/kotlin-dataframe/testData/box/columnName.kt

@@ -0,0 +1,16 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class Record(
+    @ColumnName("a")
+    val abc: String,
+)
+
+fun box(): String {
+    val df = dataFrameOf("a")(1).cast<Record>()
+    df.abc
+    return "OK"
+}

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

@@ -34,6 +34,12 @@ public void testColumnGroupApi() {
     runTest("testData/box/columnGroupApi.kt");
   }
 
+  @Test
+  @TestMetadata("columnName.kt")
+  public void testColumnName() {
+    runTest("testData/box/columnName.kt");
+  }
+
   @Test
   @TestMetadata("columnWithStarProjection.kt")
   public void testColumnWithStarProjection() {