Merge pull request #1035 from Kotlin/supported-apis-1

More variations of "move", AddDsl.group, CreateDataFrameDsl String.invoke
Kotlin · Jan 29, 2025 · 3448fc9 · 3448fc9
2 parents 57a9768 + 8f8d368
commit 3448fc9
Show file tree

Hide file tree

Showing 20 changed files with 731 additions and 1 deletion.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
@@ -225,6 +225,7 @@ public class AddDsl<T>(
     @AccessApiOverload
     public fun group(column: AnyColumnGroupAccessor, body: AddDsl<T>.() -> Unit): Unit = group(column.name(), body)
 
+    @Interpretable("AddDslNamedGroup")
     public fun group(name: String, body: AddDsl<T>.() -> Unit) {
         val dsl = AddDsl(df)
         body(dsl)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
@@ -83,6 +83,7 @@ public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T>
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)
 
+@AccessApiOverload
 public fun <T> InsertClause<T>.after(column: ColumnAccessor<*>): DataFrame<T> = after(column.path())
 
 @AccessApiOverload

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/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
@@ -202,6 +202,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {
 
     public inline fun <reified R> inferType(noinline expression: (T) -> R): InferType<T, R> = InferType(expression)
 
+    @Interpretable("ToDataFrameDslStringInvoke")
     public abstract operator fun String.invoke(builder: CreateDataFrameDsl<T>.() -> Unit)
 }
 

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/add.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt
@@ -68,3 +68,15 @@ class AddDslStringInvoke : AbstractInterpreter<Unit>() {
         dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
     }
 }
+
+class AddDslNamedGroup : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg()
+    val Arguments.name: String by arg()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(name, addDsl.columns))
+    }
+}
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/impl/api/toDataFrame.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt
@@ -131,6 +131,18 @@ class Properties0 : AbstractInterpreter<Unit>() {
     }
 }
 
+class ToDataFrameDslStringInvoke : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.builder by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = CreateDataFrameDslImplApproximation()
+        builder(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
+    }
+}
+
 class CreateDataFrameConfiguration {
     var maxDepth = DEFAULT_MAX_DEPTH
     var traverseConfiguration: TraverseConfiguration = TraverseConfiguration()

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
@@ -71,6 +71,7 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslNamedGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
@@ -88,6 +89,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
@@ -96,6 +104,7 @@ 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.ToDataFrameDslStringInvoke
 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.TrimMargin
@@ -235,6 +244,7 @@ internal inline fun <reified T> String.load(): T {
         "toDataFrameDsl" -> ToDataFrameDsl()
         "toDataFrame" -> ToDataFrame()
         "toDataFrameDefault" -> ToDataFrameDefault()
+        "ToDataFrameDslStringInvoke" -> ToDataFrameDslStringInvoke()
         "DataFrameOf0" -> DataFrameOf0()
         "DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
         "ToDataFrameColumn" -> ToDataFrameColumn()
@@ -246,6 +256,7 @@ internal inline fun <reified T> String.load(): T {
         "FlattenDefault" -> FlattenDefault()
         "AddId" -> AddId()
         "AddDslStringInvoke" -> AddDslStringInvoke()
+        "AddDslNamedGroup" -> AddDslNamedGroup()
         "MapToFrame" -> MapToFrame()
         "Move0" -> Move0()
         "ToTop" -> ToTop()
@@ -255,6 +266,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/addDsl.kt b/plugins/kotlin-dataframe/testData/box/addDsl.kt
@@ -9,8 +9,12 @@ fun box(): String {
         "group" {
             "a" from { it }
         }
+        group("group1") {
+            "b" from { it }
+        }
     }
 
     df.group.a
+    df.group1.b
     return "OK"
 }
-Original file line number
+Diff line change
@@ Expand Up @@
         public inline fun <reified R> inferType(noinline expression: (T) -> R): InferType<T, R> = InferType(expression)
+        @Interpretable("ToDataFrameDslStringInvoke")
         public abstract operator fun String.invoke(builder: CreateDataFrameDsl<T>.() -> Unit)
     }
@@ Expand Down @@