[Compiler plugin] Add more variations of "move"

Kotlin · Jan 28, 2025 · aa50988 · aa50988
1 parent d12549b
commit aa50988
Show file tree

Hide file tree

Showing 13 changed files with 671 additions and 1 deletion.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt
@@ -53,10 +53,13 @@ public fun <T> DataFrame<T>.moveTo(newColumnIndex: Int, vararg columns: KPropert
 
 // region moveToLeft
 
+@Refine
+@Interpretable("MoveToLeft1")
 public fun <T> DataFrame<T>.moveToLeft(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toLeft()
 
 public fun <T> DataFrame<T>.moveToLeft(vararg columns: String): DataFrame<T> = moveToLeft { columns.toColumnSet() }
 
+@AccessApiOverload
 public fun <T> DataFrame<T>.moveToLeft(vararg columns: AnyColumnReference): DataFrame<T> =
     moveToLeft { columns.toColumnSet() }
 
@@ -68,13 +71,17 @@ public fun <T> DataFrame<T>.moveToLeft(vararg columns: KProperty<*>): DataFrame<
 
 // region moveToRight
 
+@Refine
+@Interpretable("MoveToRight1")
 public fun <T> DataFrame<T>.moveToRight(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toRight()
 
 public fun <T> DataFrame<T>.moveToRight(vararg columns: String): DataFrame<T> = moveToRight { columns.toColumnSet() }
 
+@AccessApiOverload
 public fun <T> DataFrame<T>.moveToRight(vararg columns: AnyColumnReference): DataFrame<T> =
     moveToRight { columns.toColumnSet() }
 
+@AccessApiOverload
 public fun <T> DataFrame<T>.moveToRight(vararg columns: KProperty<*>): DataFrame<T> =
     moveToRight { columns.toColumnSet() }
 
@@ -94,6 +101,11 @@ public fun <T, C> MoveClause<T, C>.into(
         newPathExpression = column,
     )
 
+/**
+ * Move a single selected column to top level with a new name
+ */
+@Refine
+@Interpretable("MoveInto0")
 public fun <T, C> MoveClause<T, C>.into(column: String): DataFrame<T> = pathOf(column).let { path -> into { path } }
 
 public fun <T, C> MoveClause<T, C>.intoIndexed(
@@ -109,11 +121,19 @@ public fun <T, C> MoveClause<T, C>.intoIndexed(
 
 // region under
 
+@Refine
+@Interpretable("MoveUnder0")
 public fun <T, C> MoveClause<T, C>.under(column: String): DataFrame<T> = pathOf(column).let { path -> under { path } }
 
+@AccessApiOverload
 public fun <T, C> MoveClause<T, C>.under(column: AnyColumnGroupAccessor): DataFrame<T> =
     column.path().let { path -> under { path } }
 
+/**
+ * Move selected columns under existing column group
+ */
+@Refine
+@Interpretable("MoveUnder1")
 public fun <T, C> MoveClause<T, C>.under(
     column: ColumnsSelectionDsl<T>.(ColumnWithPath<C>) -> AnyColumnReference,
 ): DataFrame<T> =
@@ -138,10 +158,13 @@ public fun <T, C> MoveClause<T, C>.toTop(
 
 // region after
 
+@Refine
+@Interpretable("MoveAfter0")
 public fun <T, C> MoveClause<T, C>.after(column: ColumnSelector<T, *>): DataFrame<T> = afterOrBefore(column, true)
 
 public fun <T, C> MoveClause<T, C>.after(column: String): DataFrame<T> = after { column.toColumnAccessor() }
 
+@AccessApiOverload
 public fun <T, C> MoveClause<T, C>.after(column: AnyColumnReference): DataFrame<T> = after { column }
 
 @AccessApiOverload
@@ -157,8 +180,12 @@ fun <T, C> MoveColsClause<T, C>.before(column: String) = before { column.toColum
 fun <T, C> MoveColsClause<T, C>.before(column: ColumnSelector<T, *>) = afterOrBefore(column, false)
 */
 
+@Refine
+@Interpretable("MoveToLeft0")
 public fun <T, C> MoveClause<T, C>.toLeft(): DataFrame<T> = to(0)
 
+@Refine
+@Interpretable("MoveToRight0")
 public fun <T, C> MoveClause<T, C>.toRight(): DataFrame<T> = to(df.ncol)
 
 public class MoveClause<T, C>(internal val df: DataFrame<T>, internal val columns: ColumnsSelector<T, C>) {

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt
@@ -795,4 +795,10 @@ class DataFrameTreeTests : BaseTest() {
         val df = typed2.convert { nameAndCity }.asFrame { it.remove { city } }
         df.nameAndCity.columns() shouldBe typed2.nameAndCity.remove { city }.columns()
     }
+
+    @Test
+    fun `move under existing group`() {
+        val df = typed2.move { age }.under { nameAndCity }
+        df.nameAndCity.columnNames() shouldBe listOf("name", "city", "age")
+    }
 }
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/move.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/move.kt
@@ -1,8 +1,13 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
+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.pathOf
+import org.jetbrains.kotlinx.dataframe.api.moveToLeft
+import org.jetbrains.kotlinx.dataframe.api.toLeft
+import org.jetbrains.kotlinx.dataframe.api.toRight
 import org.jetbrains.kotlinx.dataframe.api.toTop
+import org.jetbrains.kotlinx.dataframe.api.under
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
@@ -30,4 +35,74 @@ class ToTop : AbstractSchemaModificationInterpreter() {
     }
 }
 
+class MoveUnder0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.column: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.under(column).toPluginDataFrameSchema()
+    }
+}
+
+class MoveUnder1 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.under { column.col.path }.toPluginDataFrameSchema()
+    }
+}
+
+class MoveInto0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.column: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.into(column).toPluginDataFrameSchema()
+    }
+}
+
+class MoveToLeft0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.toLeft().toPluginDataFrameSchema()
+    }
+}
+
+class MoveToLeft1 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = columns.resolve(receiver).map { it.path }
+        return receiver.asDataFrame().moveToLeft { columns.toColumnSet() }.toPluginDataFrameSchema()
+    }
+}
+
+
+class MoveToRight0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.toRight().toPluginDataFrameSchema()
+    }
+}
+
+class MoveAfter0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: MoveClauseApproximation by arg()
+    val Arguments.column: SingleColumnApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = receiver.columns.resolve(receiver.df).map { it.path }
+        return receiver.df.asDataFrame().move { columns.toColumnSet() }.after { column.col.path }.toPluginDataFrameSchema()
+    }
+}
+
+
 class MoveClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -88,6 +88,13 @@ 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.MoveAfter0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveInto0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToLeft0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToLeft1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveToRight0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveUnder1
 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.ReadExcel
@@ -255,6 +262,13 @@ internal inline fun <reified T> String.load(): T {
         "ValueCounts" -> ValueCounts()
         "RenameToCamelCase" -> RenameToCamelCase()
         "RenameToCamelCaseClause" -> RenameToCamelCaseClause()
+        "MoveUnder0" -> MoveUnder0()
+        "MoveUnder1" -> MoveUnder1()
+        "MoveInto0" -> MoveInto0()
+        "MoveToLeft0" -> MoveToLeft0()
+        "MoveToLeft1" -> MoveToLeft1()
+        "MoveToRight0" -> MoveToRight0()
+        "MoveAfter0" -> MoveAfter0()
         else -> error("$this")
     } as T
 }
diff --git a/plugins/kotlin-dataframe/testData/box/moveAfter.fir.txt b/plugins/kotlin-dataframe/testData/box/moveAfter.fir.txt
@@ -0,0 +1,107 @@
+FILE: moveAfter.kt
+    public final fun box(): R|kotlin/String| {
+        lval df: R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_66>| = R|org/jetbrains/kotlinx/dataframe/api/dataFrameOf|(vararg(String(a), String(b), String(c))).R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_64>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_64>| <inline=Inline, kind=EXACTLY_ONCE>  {
+            local abstract class Invoke_64I : R|kotlin/Any| {
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val c: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val b: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val a: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public constructor(): R|<local>/Invoke_64I|
+
+            }
+
+            local final class Scope0 : R|kotlin/Any| {
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_64I>|.c: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_64I>|.c: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_64I>|.b: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_64I>|.b: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_64I>|.a: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_64I>|.a: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public constructor(): R|<local>/Scope0|
+
+            }
+
+            local abstract class Invoke_64 : R|<local>/Invoke_64I| {
+                @R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
+                    public get(): R|<local>/Scope0|
+
+                public constructor(): R|<local>/Invoke_64|
+
+            }
+
+            ^ R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/api/DataFrameBuilder.invoke|(vararg(Int(1), Int(2), Int(3)))
+        }
+        ).R|org/jetbrains/kotlinx/dataframe/api/move|<R|<local>/Invoke_64|, R|kotlin/Int|>(<L> = move@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_64>|.<anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_64>|): R|org/jetbrains/kotlinx/dataframe/columns/ColumnsResolver<kotlin/Int>| <inline=NoInline>  {
+            ^ (this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.a|
+        }
+        ).R|kotlin/let|<R|org/jetbrains/kotlinx/dataframe/api/MoveClause<<local>/Invoke_64, kotlin/Int>|, R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_66>|>(<L> = fun <anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/MoveClause<<local>/Invoke_64, kotlin/Int>|): R|org/jetbrains/kotlinx/dataframe/DataFrame<<local>/Invoke_66>| <inline=Inline, kind=EXACTLY_ONCE>  {
+            local abstract class Invoke_66I : R|kotlin/Any| {
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(2)) public abstract val c: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(0)) public abstract val b: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                @R|org/jetbrains/kotlinx/dataframe/annotations/Order|(order = Int(1)) public abstract val a: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public constructor(): R|<local>/Invoke_66I|
+
+            }
+
+            local final class Scope0 : R|kotlin/Any| {
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_66I>|.c: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_66I>|.c: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_66I>|.b: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_66I>|.b: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/DataRow<<local>/Invoke_66I>|.a: R|kotlin/Int|
+                    public get(): R|kotlin/Int|
+
+                public final val R|org/jetbrains/kotlinx/dataframe/ColumnsScope<<local>/Invoke_66I>|.a: R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+                    public get(): R|org/jetbrains/kotlinx/dataframe/DataColumn<kotlin/Int>|
+
+                public constructor(): R|<local>/Scope0|
+
+            }
+
+            local abstract class Invoke_66 : R|<local>/Invoke_66I| {
+                @R|org/jetbrains/kotlinx/dataframe/annotations/ScopeProperty|() public abstract val scope0: R|<local>/Scope0|
+                    public get(): R|<local>/Scope0|
+
+                public constructor(): R|<local>/Invoke_66|
+
+            }
+
+            ^ R|<local>/it|.R|org/jetbrains/kotlinx/dataframe/api/after|<R|<local>/Invoke_64|, R|kotlin/Int|>(<L> = after@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_64>|.<anonymous>(it: R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl<<local>/Invoke_64>|): R|org/jetbrains/kotlinx/dataframe/columns/SingleColumn<*>| <inline=NoInline>  {
+                ^ (this@R|/box|, this@R|special/anonymous|).R|<local>/Scope0.b|
+            }
+            )
+        }
+        )
+        ^box String(OK)
+    }
diff --git a/plugins/kotlin-dataframe/testData/box/moveAfter.kt b/plugins/kotlin-dataframe/testData/box/moveAfter.kt
@@ -0,0 +1,10 @@
+// FIR_DUMP
+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).move { a }.after { b }
+    return "OK"
+}
diff --git a/plugins/kotlin-dataframe/testData/box/moveInto.kt b/plugins/kotlin-dataframe/testData/box/moveInto.kt
@@ -0,0 +1,13 @@
+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)
+        .group { a and b }.into("d")
+        .move { d.a }.into("renamedA")
+    df.renamedA
+    df.compareSchemas(strict = true)
+    return "OK"
+}