Merge pull request #1474 from CarloMariaProietti/add_move_before

Jolanrensen · web-flow · commit fbcccc4aea8a · 2025-10-07T14:32:58.000+02:00
Add move before
diff --git a/core/api/core.api b/core/api/core.api
@@ -3414,6 +3414,8 @@ public final class org/jetbrains/kotlinx/dataframe/api/MoveKt {
 	public static final fun after (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun after (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lkotlin/reflect/KProperty;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun after (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
+	public static final fun before (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
+	public static final fun before (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun into (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun into (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun intoIndexed (Lorg/jetbrains/kotlinx/dataframe/api/MoveClause;Lkotlin/jvm/functions/Function3;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
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
@@ -621,13 +621,54 @@ public fun <T, C> MoveClause<T, C>.after(column: KProperty<*>): DataFrame<T> = a
 
 // endregion
 
-/* TODO: implement 'before'
-fun <T, C> MoveColsClause<T, C>.before(columnPath: ColumnPath) = before { columnPath.toColumnDef() }
-fun <T, C> MoveColsClause<T, C>.before(column: Column) = before { column }
-fun <T, C> MoveColsClause<T, C>.before(column: KProperty<*>) = before { column.toColumnDef() }
-fun <T, C> MoveColsClause<T, C>.before(column: String) = before { column.toColumnDef() }
-fun <T, C> MoveColsClause<T, C>.before(column: ColumnSelector<T, *>) = afterOrBefore(column, false)
-*/
+// region before
+
+/**
+ * Moves columns, previously selected with [move] to the position before the
+ * specified [column] within the [DataFrame].
+ *
+ * Returns a new [DataFrame] with updated columns.
+ *
+ * See [Selecting Columns][SelectingColumns].
+ *
+ * For more information: {@include [DocumentationUrls.Move]}
+ *
+ * ### This Before Overload
+ */
+@ExcludeFromSources
+internal interface MoveBefore
+
+/**
+ * {@include [MoveBefore]}
+ * @include [SelectingColumns.Dsl]
+ *
+ * ### Examples:
+ * ```kotlin
+ * df.move { age and weight }.before { surname }
+ * df.move { cols(3..5) }.before { col(2) }
+ * ```
+ *
+ * @param [column] A [ColumnSelector] specifying the column
+ * before which the selected columns will be placed.
+ */
+@Refine
+@Interpretable("MoveBefore0")
+public fun <T, C> MoveClause<T, C>.before(column: ColumnSelector<T, *>): DataFrame<T> = afterOrBefore(column, false)
+
+/**
+ * {@include [MoveBefore]}
+ * @include [SelectingColumns.ColumnNames]
+ *
+ * ### Examples:
+ * ```kotlin
+ * df.move("age", "weight").before("surname")
+ * ```
+ * @param [column] The [Column Name][String] specifying the column
+ * before which the selected columns will be placed.
+ */
+public fun <T, C> MoveClause<T, C>.before(column: String): DataFrame<T> = before { column.toColumnAccessor() }
+
+// endregion
 
 @Deprecated(TO_LEFT, ReplaceWith(TO_LEFT_REPLACE), DeprecationLevel.ERROR)
 public fun <T, C> MoveClause<T, C>.toLeft(): DataFrame<T> = to(0)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt
@@ -6,21 +6,25 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
 import org.jetbrains.kotlinx.dataframe.api.MoveClause
+import org.jetbrains.kotlinx.dataframe.api.after
+import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.getColumn
 import org.jetbrains.kotlinx.dataframe.api.getColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.getColumnWithPath
+import org.jetbrains.kotlinx.dataframe.api.move
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.DataFrameReceiver
 import org.jetbrains.kotlinx.dataframe.impl.asList
 import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnWithPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ColumnPosition
 import org.jetbrains.kotlinx.dataframe.impl.columns.tree.getOrPut
+import org.jetbrains.kotlinx.dataframe.path
 
-// TODO: support 'before' mode
 internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>, isAfter: Boolean): DataFrame<T> {
     val removeResult = df.removeImpl(columns = columns)
 
@@ -35,8 +39,9 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
         if (sourceSegments.size <= targetSegments.size &&
             sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }
         ) {
+            val afterOrBefore = if (isAfter) "after" else "before"
             throw IllegalArgumentException(
-                "Cannot move column '${sourcePath.joinToString()}' after its own child column '${targetPath.joinToString()}'",
+                "Cannot move column '${sourcePath.joinToString()}' $afterOrBefore its own child column '${targetPath.joinToString()}'",
             )
         }
     }
@@ -78,7 +83,16 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
         }
         ColumnToInsert(path, sourceCol.data, refNode)
     }
-    return removeResult.df.insertImpl(toInsert)
+    if (isAfter) {
+        return removeResult.df.insertImpl(toInsert)
+    }
+
+    //  Move the target column after the removed/inserted columns
+    val logicOfAfter = removeResult.df.insertImpl(toInsert)
+    val lastOfInsertedCols = toInsert.last().insertionPath
+    val siblingsOfTargetAndTarget = removeResult.df[parentPath].asColumnGroup().columns().map { parentPath + it.path }
+    val target = siblingsOfTargetAndTarget.filter { it.last() == targetPath.last() }
+    return logicOfAfter.move { target.toColumnSet() }.after { lastOfInsertedCols }
 }
 
 internal fun <T, C> MoveClause<T, C>.moveImpl(
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt
@@ -152,4 +152,72 @@ class MoveTests {
             grouped.move { "a"["b"] }.after { "a"["b"] }
         }.message shouldBe "Cannot move column 'a/b' after its own child column 'a/b'"
     }
+
+    @Test
+    fun `move before first`() {
+        val df = dataFrameOf("1", "2")(1, 2)
+        shouldNotThrowAny {
+            df.move("2").before("1") shouldBe dataFrameOf("2", "1")(2, 1)
+        }
+    }
+
+    @Test
+    fun `move before in nested structure`() {
+        val df = grouped.move { "a"["b"] }
+            .before { "a"["c"]["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
+        df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
+        df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("b", "d")
+    }
+
+    @Test
+    fun `move before multiple columns`() {
+        val df = grouped.move { "a"["b"] and "b"["c"] }
+            .before { "a"["c"]["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
+        df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
+        df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("b", "c", "d")
+        df["b"].asColumnGroup().columnNames() shouldBe listOf("d")
+    }
+
+    @Test
+    fun `move before with column selector`() {
+        val df = grouped.move { colsAtAnyDepth().filter { it.name == "r" || it.name == "w" } }
+            .before { "a"["c"]["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "e")
+        df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("w", "r", "d")
+    }
+
+    @Test
+    fun `move before between groups`() {
+        val df = grouped.move { "a"["b"] }.before { "b"["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
+        df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
+        df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "b", "d")
+    }
+
+    @Test
+    fun `should throw when moving parent before child`() {
+        // Simple case: direct parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").before { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a' before its own child column 'a/b'"
+
+        // Nested case: deeper parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").before { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a' before its own child column 'a/c/d'"
+
+        // Group case: moving group after its nested column
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["c"] }.before { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a/c' before its own child column 'a/c/d'"
+    }
+
+    @Test
+    fun `should throw when moving column before itself`() {
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["b"] }.before { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a/b' before its own child column 'a/b'"
+    }
 }
