[Compiler plugin] DataFrame.reorderColumnsByName

Author: koperagen

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

@@ -6,6 +6,8 @@ import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
+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.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.api.reorderImpl
@@ -52,6 +54,8 @@ public fun <T, V : Comparable<V>> DataFrame<T>.reorderColumnsBy(
         inFrameColumns = atAnyDepth,
     ).reorderImpl(desc, expression)
 
+@Refine
+@Interpretable("ReorderColumnsByName")
 public fun <T> DataFrame<T>.reorderColumnsByName(atAnyDepth: Boolean = true, desc: Boolean = false): DataFrame<T> =
     reorderColumnsBy(atAnyDepth, desc) { name() }
 

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

@@ -0,0 +1,20 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.reorderColumnsByName
+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.Present
+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 ReorderColumnsByName : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.atAnyDepth: Boolean by arg(defaultValue = Present(true))
+    val Arguments.desc: Boolean by arg(defaultValue = Present(false))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.asDataFrame().reorderColumnsByName(atAnyDepth, desc).toPluginDataFrameSchema()
+    }
+}

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

@@ -123,6 +123,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCase
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameToCamelCaseClause
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReorderColumnsByName
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
@@ -293,6 +294,7 @@ internal inline fun <reified T> String.load(): T {
         "MergeId" -> MergeId()
         "MergeBy0" -> MergeBy0()
         "MergeBy1" -> MergeBy1()
+        "ReorderColumnsByName" -> ReorderColumnsByName()
         else -> error("$this")
     } as T
 }