[KDF] Consider ColumnName when a generated extension property is used in CS DSL

Author: koperagen

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

@@ -8,8 +8,10 @@ package org.jetbrains.kotlinx.dataframe.plugin
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.declaredProperties
 import org.jetbrains.kotlin.fir.declarations.findArgumentByName
 import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.declarations.getAnnotationWithResolvedArgumentsByClassId
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.expressions.*
@@ -105,7 +107,9 @@ fun <T> KotlinTypeFacade.interpret(
         val expectedReturnType = expectedArgument.klass
         val value: Interpreter.Success<Any?>? = when (expectedArgument.lens) {
             is Interpreter.Value -> {
-                extractValue(it.expression, reporter)
+                context(session) {
+                    extractValue(it.expression, reporter)
+                }
             }
 
             is Interpreter.ReturnType -> {
@@ -188,6 +192,7 @@ fun <T> KotlinTypeFacade.interpret(
     }
 }
 
+context(session: FirSession)
 private fun KotlinTypeFacade.extractValue(
     expression: FirExpression?,
     reporter: InterpretationErrorReporter,
@@ -353,20 +358,21 @@ private fun List<FirPropertySymbol>.sortPropertiesByOrderAnnotation(sessionConte
     return result
 }
 
-private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAccessExpression): ColumnsResolver {
-    return result.resolvedType.let {
+context(session: FirSession)
+private fun KotlinTypeFacade.columnWithPathApproximations(propertyAccess: FirPropertyAccessExpression): ColumnsResolver {
+    return propertyAccess.resolvedType.let {
         val column = when (it.classId) {
             Names.DATA_COLUMN_CLASS_ID -> {
                 val type = when (val arg = it.typeArguments.single()) {
                     is ConeStarProjection -> session.builtinTypes.nullableAnyType.coneType
                     else -> arg as ConeClassLikeType
                 }
-                simpleColumnOf(f(result), type)
+                simpleColumnOf(propertyAccess.columnName(), type)
             }
             Names.COLUM_GROUP_CLASS_ID -> {
                 val arg = it.typeArguments.single()
-                val path = f(result)
-                SimpleColumnGroup(path, pluginDataFrameSchema(arg).columns())
+                val name = propertyAccess.columnName()
+                SimpleColumnGroup(name, pluginDataFrameSchema(arg).columns())
             }
             else -> return object : ColumnsResolver {
                 override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
@@ -376,7 +382,7 @@ private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAcc
         }
         SingleColumnApproximation(
             ColumnWithPathApproximation(
-                path = ColumnPathApproximation(path(result)),
+                path = ColumnPathApproximation(path(propertyAccess)),
                 column
             )
         )
@@ -450,8 +456,9 @@ private fun SessionContext.shouldBeConvertedToFrameColumn(it: FirPropertySymbol)
 private fun isDataFrame(it: FirPropertySymbol) =
     it.resolvedReturnType.classId == Names.DF_CLASS_ID
 
+context(session: FirSession)
 fun path(propertyAccessExpression: FirPropertyAccessExpression): List<String> {
-    val colName = f(propertyAccessExpression)
+    val colName = propertyAccessExpression.columnName()
     val typeRef = propertyAccessExpression.dispatchReceiver?.resolvedType
     val joinDsl = ClassId(FqName("org.jetbrains.kotlinx.dataframe.api"), Name.identifier("JoinDsl"))
     if (typeRef?.classId?.equals(joinDsl) == true && colName == "right") {
@@ -470,8 +477,19 @@ fun path(propertyAccessExpression: FirPropertyAccessExpression): List<String> {
     }
 }
 
-fun f(propertyAccessExpression: FirPropertyAccessExpression): String {
-    return propertyAccessExpression.calleeReference.resolved!!.name.identifier
+context(session: FirSession)
+fun FirPropertyAccessExpression.columnName(): String {
+    val name = toResolvedCallableSymbol()?.name
+    val columnName =
+        extensionReceiver?.resolvedType?.typeArguments?.getOrNull(0)?.type?.toRegularClassSymbol(session)
+            ?.declaredProperties(session)
+            ?.firstOrNull { it.name == name }
+            ?.let {
+                val expression = it.getAnnotationWithResolvedArgumentsByClassId(Names.COLUMN_NAME_ANNOTATION, session)
+                    ?.argumentMapping?.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression
+                expression?.value as? String
+            }
+    return columnName ?: calleeReference.resolved!!.name.identifier
 }
 
 internal fun FirFunctionCall.collectArgumentExpressions(): RefinedArguments {

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

@@ -0,0 +1,19 @@
+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 A(
+    @ColumnName("hello_world")
+    val helloWorld: Int
+)
+
+fun box(): String {
+    val df = dataFrameOf(A(123)).convert { helloWorld }.with { "test" }
+    val col: DataColumn<String> = df.hello_world
+
+    val df1 = dataFrameOf(A(123)).rename { helloWorld }.into("hiWorld")
+    val col1: DataColumn<Int> = df1.hiWorld
+    return "OK"
+}