[Compiler plugin] Support move { }.toTop()

Only works for default lambda

Author: koperagen

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

@@ -5,6 +5,8 @@ import org.jetbrains.kotlinx.dataframe.AnyColumnReference
 import org.jetbrains.kotlinx.dataframe.ColumnSelector
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 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.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
@@ -18,6 +20,7 @@ import kotlin.reflect.KProperty
 
 // region move
 
+@Interpretable("Move0")
 public fun <T, C> DataFrame<T>.move(columns: ColumnsSelector<T, C>): MoveClause<T, C> = MoveClause(this, columns)
 
 public fun <T> DataFrame<T>.move(vararg cols: String): MoveClause<T, Any?> = move { cols.toColumnSet() }
@@ -120,6 +123,8 @@ public fun <T, C> MoveClause<T, C>.under(
 
 public fun <T, C> MoveClause<T, C>.to(columnIndex: Int): DataFrame<T> = moveTo(columnIndex)
 
+@Refine
+@Interpretable("ToTop")
 public fun <T, C> MoveClause<T, C>.toTop(
     newColumnName: ColumnsSelectionDsl<T>.(ColumnWithPath<C>) -> String = { it.name() },
 ): DataFrame<T> = into { newColumnName(it).toColumnAccessor() }

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

@@ -0,0 +1,33 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.move
+import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.api.toTop
+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.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class Move0 : AbstractInterpreter<MoveClauseApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): MoveClauseApproximation {
+        return MoveClauseApproximation(receiver, columns)
+    }
+}
+
+class ToTop : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { pathOf(*it.path.path.toTypedArray()) }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.toTop().toPluginDataFrameSchema()
+    }
+}
+
+class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)

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

@@ -80,12 +80,14 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
@@ -200,6 +202,8 @@ internal inline fun <reified T> String.load(): T {
         "AddId" -> AddId()
         "AddDslStringInvoke" -> AddDslStringInvoke()
         "MapToFrame" -> MapToFrame()
+        "Move0" -> Move0()
+        "ToTop" -> ToTop()
         else -> error("$this")
     } as T
 }

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

@@ -0,0 +1,12 @@
+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")(1, 2).group { a and b }.into("c").move { c.a }.toTop()
+    df.a
+    df.c.b
+    df.compareSchemas()
+    return "OK"
+}

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

@@ -250,6 +250,12 @@ public void testMapToFrame() {
     runTest("testData/box/mapToFrame.kt");
   }
 
+  @Test
+  @TestMetadata("moveToTop.kt")
+  public void testMoveToTop() {
+    runTest("testData/box/moveToTop.kt");
+  }
+
   @Test
   @TestMetadata("nestedDataSchemaCodegen.kt")
   public void testNestedDataSchemaCodegen() {