[Compiler plugin] Support "df.rename(Pair<String, String>)" operation

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/rename.kt

@@ -24,6 +24,8 @@ import kotlin.reflect.KProperty
 
 // region DataFrame
 
+@Refine
+@Interpretable("RenameMapping")
 public fun <T> DataFrame<T>.rename(vararg mappings: Pair<String, String>): DataFrame<T> =
     rename { mappings.map { it.first.toColumnAccessor() }.toColumnSet() }
         .into(*mappings.map { it.second }.toTypedArray())

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt

@@ -1,12 +1,14 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
+import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
 import org.jetbrains.kotlinx.dataframe.api.rename
 import org.jetbrains.kotlinx.dataframe.api.renameToCamelCase
 import org.jetbrains.kotlinx.dataframe.api.toCamelCase
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
@@ -40,6 +42,22 @@ class RenameInto : AbstractSchemaModificationInterpreter() {
     }
 }
 
+class RenameMapping : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.mappings: List<Interpreter.Success<Pair<*, *>>> by arg()
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val mappings = mappings.map {
+            val it = it.value
+            val name = (it.first as? FirLiteralExpression)?.value as? String
+            val newName = (it.second as? FirLiteralExpression)?.value  as? String
+            if (name == null || newName == null) return PluginDataFrameSchema(emptyList())
+            name to newName
+        }
+
+        return receiver.asDataFrame().rename(*mappings.toTypedArray()).toPluginDataFrameSchema()
+    }
+}
+
 internal fun PluginDataFrameSchema.map(selected: ColumnsSet, nextName: () -> String): PluginDataFrameSchema =
     PluginDataFrameSchema(
         f(columns(), nextName, selected, emptyList()),

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

@@ -107,6 +107,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairToConstructor
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PerRowCol
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameMapping
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.StringColumnsConstructor
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
@@ -224,6 +225,7 @@ internal inline fun <reified T> String.load(): T {
         "ReadJson0" -> ReadJson0()
         "ReadCSV0" -> ReadCSV0()
         "Rename" -> Rename()
+        "RenameMapping" -> RenameMapping()
         "Select0" -> Select0()
         "Expr0" -> Expr0()
         "And0" -> And0()

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

@@ -0,0 +1,10 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("a")(123).rename("a" to "c")
+    val value: Int = df.c[0]
+    return "OK"
+}

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

@@ -478,6 +478,12 @@ public void testRename() {
     runTest("testData/box/rename.kt");
   }
 
+  @Test
+  @TestMetadata("renameMapping.kt")
+  public void testRenameMapping() {
+    runTest("testData/box/renameMapping.kt");
+  }
+
   @Test
   @TestMetadata("renameToCamelCase.kt")
   public void testRenameToCamelCase() {