insert after plugin support

Author: AndreiKingsley

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

@@ -79,6 +79,8 @@ public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathO
 
 // region after
 
+@Refine
+@Interpretable("InsertAfter0")
 public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T> = after(df.getColumnPath(column))
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)

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.Infer
+import org.jetbrains.kotlinx.dataframe.api.after
 import org.jetbrains.kotlinx.dataframe.api.insert
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.under
@@ -122,3 +123,14 @@ internal class Under4 : AbstractInterpreter<PluginDataFrameSchema>() {
             .toPluginDataFrameSchema()
     }
 }
+
+internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.column: SingleColumnApproximation by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).after(column.col.path)
+            .toPluginDataFrameSchema()
+    }
+}

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

@@ -128,6 +128,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceExpression
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceInto
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReducePredicate
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByXs
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAfter0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last2
@@ -286,6 +287,7 @@ internal inline fun <reified T> String.load(): T {
         "Under2" -> Under2()
         "Under3" -> Under3()
         "Under4" -> Under4()
+        "InsertAfter0" -> InsertAfter0()
         "Join0" -> Join0()
         "Match0" -> Match0()
         "Rename" -> Rename()

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

@@ -0,0 +1,12 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.api.*
+
+fun box(): String {
+    val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6)
+
+    val dfWithD = df.insert("d") { b  * c }.after { a }
+
+    val dCol: DataColumn<Int> = dfWithD.d
+
+    return "OK"
+}

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

@@ -0,0 +1,13 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.api.*
+
+fun box(): String {
+    val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6)
+    val grouped = df.group { a and b and c }.into("g")
+
+    val dfWithD = grouped.insert("d") { g.b  * g.c }.after { g.b }
+
+    val dCol: DataColumn<Int> = dfWithD.g.d
+
+    return "OK"
+}

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

@@ -0,0 +1,13 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.api.*
+
+fun box(): String {
+    val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6)
+    val grouped = df.group { a and b }.into("g")
+    grouped.move { g.a }.after { g.b }
+
+    val dfMovedAfterNested = grouped.move { c }.after { g.a }
+    val cNested: DataColumn<Int> = dfMovedAfterNested.g.c
+
+    return "OK"
+}

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

@@ -286,6 +286,18 @@ public void testInsert() {
     runTest("testData/box/insert.kt");
   }
 
+  @Test
+  @TestMetadata("insertAfter.kt")
+  public void testInsertAfter() {
+    runTest("testData/box/insertAfter.kt");
+  }
+
+  @Test
+  @TestMetadata("insertAfterNested.kt")
+  public void testInsertAfterNested() {
+    runTest("testData/box/insertAfterNested.kt");
+  }
+
   @Test
   @TestMetadata("inventNamesForLocalClasses.kt")
   public void testInventNamesForLocalClasses() {
@@ -346,6 +358,12 @@ public void testMoveAfter() {
     runTest("testData/box/moveAfter.kt");
   }
 
+  @Test
+  @TestMetadata("moveAfterNested.kt")
+  public void testMoveAfterNested() {
+    runTest("testData/box/moveAfterNested.kt");
+  }
+
   @Test
   @TestMetadata("moveInto.kt")
   public void testMoveInto() {