[Compiler plugin] Support move to and insert at

Author: koperagen

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

@@ -94,6 +94,8 @@ public fun <T> InsertClause<T>.after(columnPath: ColumnPath): DataFrame<T> {
 
 // region at
 
+@Refine
+@Interpretable("InsertAt")
 public fun <T> InsertClause<T>.at(position: Int): DataFrame<T> = df.add(column).move(column).to(position)
 
 // endregion

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

@@ -520,6 +520,8 @@ public fun <T, C> MoveClause<T, C>.under(
  * @param [columnIndex] The index specifying the position in the [DataFrame] columns
  *  * where the selected columns will be moved.
  */
+@Refine
+@Interpretable("MoveTo")
 public fun <T, C> MoveClause<T, C>.to(columnIndex: Int): DataFrame<T> = moveTo(columnIndex)
 
 /**

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt

@@ -2340,6 +2340,18 @@ class DataFrameTests : BaseTest() {
                 .move { city }.to(1)
     }
 
+    @Test
+    fun `move beyond range of existing column indices`() {
+        val res = typed.move { city }.to(11)
+        res.columnNames() shouldBe listOf("name", "age", "weight", "city")
+    }
+
+    @Test
+    fun `move multiple columns beyond range of existing column indices`() {
+        val res = typed.move { city and name }.to(11)
+        res.columnNames() shouldBe listOf("age", "weight", "city", "name")
+    }
+
     @Test
     fun splitIntoThisAndNewColumn() {
         val split = typed.split { name }.by { listOf(it.dropLast(1), it.last()) }.into("name", "lastChar")

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.after
+import org.jetbrains.kotlinx.dataframe.api.at
 import org.jetbrains.kotlinx.dataframe.api.insert
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.under
@@ -86,3 +87,14 @@ internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
             .toPluginDataFrameSchema()
     }
 }
+
+internal class InsertAt : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.receiver: InsertClauseApproximation by arg()
+    val Arguments.position: Int by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).at(position)
+            .toPluginDataFrameSchema()
+    }
+}

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

@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.api.after
 import org.jetbrains.kotlinx.dataframe.api.into
 import org.jetbrains.kotlinx.dataframe.api.move
 import org.jetbrains.kotlinx.dataframe.api.moveToStart
+import org.jetbrains.kotlinx.dataframe.api.to
 import org.jetbrains.kotlinx.dataframe.api.toStart
 import org.jetbrains.kotlinx.dataframe.api.toEnd
 import org.jetbrains.kotlinx.dataframe.api.toTop
@@ -104,5 +105,13 @@ class MoveAfter0 : AbstractSchemaModificationInterpreter() {
     }
 }
 
+class MoveTo : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.columnIndex: Int by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame().move { receiver.columns }.to(columnIndex).toPluginDataFrameSchema()
+    }
+}
 
 class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)

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

@@ -148,6 +148,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByStdOf
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySum1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBySumOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAt
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Merge0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
@@ -157,6 +158,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeInto0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveAfter0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveInto0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveTo
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToStart1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToEnd0
@@ -307,6 +309,7 @@ internal inline fun <reified T> String.load(): T {
         "Under1" -> Under1()
         "Under4" -> Under4()
         "InsertAfter0" -> InsertAfter0()
+        "InsertAt" -> InsertAt()
         "Join0" -> Join0()
         "Match0" -> Match0()
         "Rename" -> Rename()
@@ -431,6 +434,7 @@ internal inline fun <reified T> String.load(): T {
         "MoveToStart1" -> MoveToStart1()
         "MoveToEnd0" -> MoveToEnd0()
         "MoveAfter0" -> MoveAfter0()
+        "MoveTo" -> MoveTo()
         "GroupByAdd" -> GroupByAdd()
         "Merge0" -> Merge0()
         "MergeInto0" -> MergeInto0()

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

@@ -0,0 +1,15 @@
+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", "b", "c")(1, 2, 3, 4, 5, 6)
+    val dCol: DataColumn<String> = df.insert("d") { (b * c).toString() }.at(2)
+        // weird way to check order, but ok?..
+        .select { drop(2) }
+        .select { take(1) }
+        .select { d }.d
+
+    return "OK"
+}

plugins/kotlin-dataframe/testData/box/moveTo.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", "b", "c")(1, 2, 3, 4, 5, 6)
+    val res = df.move { c and a }.to(3).select { drop(1) }
+    res.compareSchemas(strict = true)
+    return "OK"
+}

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

@@ -352,6 +352,12 @@ public void testInsertAfterNested() {
     runTest("testData/box/insertAfterNested.kt");
   }
 
+  @Test
+  @TestMetadata("insertAt.kt")
+  public void testInsertAt() {
+    runTest("testData/box/insertAt.kt");
+  }
+
   @Test
   @TestMetadata("inventNamesForLocalClasses.kt")
   public void testInventNamesForLocalClasses() {
@@ -424,6 +430,12 @@ public void testMoveInto() {
     runTest("testData/box/moveInto.kt");
   }
 
+  @Test
+  @TestMetadata("moveTo.kt")
+  public void testMoveTo() {
+    runTest("testData/box/moveTo.kt");
+  }
+
   @Test
   @TestMetadata("moveToEnd.kt")
   public void testMoveToEnd() {