[Compiler plugin] Support addId

Author: koperagen

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

@@ -3,6 +3,8 @@ package org.jetbrains.kotlinx.dataframe.api
 import org.jetbrains.kotlinx.dataframe.AnyCol
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 
 // region DataColumn
@@ -15,6 +17,8 @@ public fun AnyCol.addId(columnName: String = "id"): AnyFrame = toDataFrame().add
 
 public fun <T> DataFrame<T>.addId(column: ColumnAccessor<Int>): DataFrame<T> = insert(column) { index() }.at(0)
 
+@Refine
+@Interpretable("AddId")
 public fun <T> DataFrame<T>.addId(columnName: String = "id"): DataFrame<T> = insert(columnName) { index() }.at(0)
 
 // endregion

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

@@ -0,0 +1,21 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+
+class AddId : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columnName: String by arg(defaultValue = Present("id"))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = buildList {
+            add(simpleColumnOf(columnName, session.builtinTypes.intType.type))
+            addAll(receiver.columns())
+        }
+        return PluginDataFrameSchema(columns)
+    }
+}

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

@@ -67,6 +67,7 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
@@ -194,6 +195,7 @@ internal inline fun <reified T> String.load(): T {
         "UpdateWith0" -> UpdateWith0()
         "Flatten0" -> Flatten0()
         "FlattenDefault" -> FlattenDefault()
+        "AddId" -> AddId()
         else -> error("$this")
     } as T
 }

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

@@ -0,0 +1,11 @@
+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")(1).addId()
+    val i: DataColumn<Int> = df.id
+    val i1: DataColumn<Int> = dataFrameOf("a")(1).addId("i").i
+    return "OK"
+}

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

@@ -17,6 +17,12 @@
 @TestMetadata("testData/box")
 @TestDataPath("$PROJECT_ROOT")
 public class DataFrameBlackBoxCodegenTestGenerated extends AbstractDataFrameBlackBoxCodegenTest {
+  @Test
+  @TestMetadata("addId.kt")
+  public void testAddId() {
+    runTest("testData/box/addId.kt");
+  }
+
   @Test
   public void testAllFilesPresentInBox() {
     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("testData/box"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);