Move/Insert After fix

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt

@@ -79,6 +79,8 @@ public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathO
 
 // region after
 
+@Refine
+@Interpretable("InsertAfter0")
 public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T> = after(df.getColumnPath(column))
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt

@@ -11,6 +11,7 @@ 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.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.impl.DataFrameReceiver
@@ -23,20 +24,70 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
     val removeResult = df.removeImpl(columns = columns)
 
     val targetPath = df.getColumnWithPath(column).path
+    val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath<C>().path }
+
+    // Check if any source path is a prefix of the target path
+    sourcePaths.forEach { sourcePath ->
+        val sourceSegments = sourcePath.toList()
+        val targetSegments = targetPath.toList()
+
+        if (sourceSegments.size <= targetSegments.size &&
+            sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }
+        ) {
+            throw IllegalArgumentException(
+                "Cannot move column '${sourcePath.joinToString()}' after its child column '${targetPath.joinToString()}'",
+            )
+        }
+    }
+
     val removeRoot = removeResult.removedColumns.first().getRoot()
 
     val refNode = removeRoot.getOrPut(targetPath) {
-        val parentPath = it.dropLast(1)
-        val parent = if (parentPath.isEmpty()) df else df.getColumnGroup(parentPath)
-        val index = parent.getColumnIndex(it.last())
-        val col = df.getColumn(index)
+        val path = it.toList()
+
+        // Find the group reference (if any) and adjust the path
+        val groupRefIndex = path.indexOfFirst { it.isEmpty() }
+
+        // Calculate effective path and column name based on group reference
+        val effectivePath = if (groupRefIndex >= 0) {
+            // Nested column reference
+            path.take(groupRefIndex)
+        } else {
+            path.dropLast(1)
+        }
+
+        val columnName = if (groupRefIndex >= 0) {
+            // Use the next segment after group reference, or previous if no next segment
+            path.getOrNull(groupRefIndex + 1) ?: path[groupRefIndex - 1]
+        } else {
+            path.last()
+        }
+
+        // Get the parent group using ColumnPath
+        val parent = if (effectivePath.isEmpty()) {
+            df
+        } else {
+            df.getColumnGroup(ColumnPath(effectivePath))
+        }
+
+        // Get the column index and the column itself
+        val index = parent.getColumnIndex(columnName)
+        val col = parent.getColumn(index)
+
         ColumnPosition(index, false, col)
     }
 
     val parentPath = targetPath.dropLast(1)
     val toInsert = removeResult.removedColumns.map {
-        val path = parentPath + it.name
-        ColumnToInsert(path, it.toColumnWithPath<C>().data, refNode)
+        val sourceCol = it.toColumnWithPath<C>()
+        val sourcePath = sourceCol.path
+        val path = if (sourcePath.size > 1) {
+            // If source is nested, preserve its structure under the target parent
+            parentPath + sourcePath.last()
+        } else {
+            parentPath + sourceCol.name()
+        }
+        ColumnToInsert(path, sourceCol.data, refNode)
     }
     return removeResult.df.insertImpl(toInsert)
 }

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.api
 
 import io.kotest.assertions.throwables.shouldNotThrowAny
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.junit.Test
@@ -91,4 +92,64 @@ class MoveTests {
             df.move("1").after("2") shouldBe dataFrameOf("2", "1")(2, 1)
         }
     }
+
+    @Test
+    fun `move after in nested structure`() {
+        val df = grouped.move { "a"["b"] }
+            .after { "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("d", "b")
+    }
+
+    @Test
+    fun `move after multiple columns`() {
+        val df = grouped.move { "a"["b"] and "b"["c"] }
+            .after { "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("d", "b", "c")
+        df["b"].asColumnGroup().columnNames() shouldBe listOf("d")
+    }
+
+    @Test
+    fun `move after with column selector`() {
+        val df = grouped.move { colsAtAnyDepth { it.name == "r" || it.name == "w" } }
+            .after { "a"["c"]["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "e")
+        df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "w", "r")
+    }
+
+    @Test
+    fun `move after between groups`() {
+        val df = grouped.move { "a"["b"] }.after { "b"["c"] }
+        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 after child`() {
+        // Simple case: direct parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").after { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a' after its child column 'a/b'"
+
+        // Nested case: deeper parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").after { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a' after its child column 'a/c/d'"
+
+        // Group case: moving group after its nested column
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["c"] }.after { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a/c' after its child column 'a/c/d'"
+    }
+
+    @Test
+    fun `should throw when moving column after itself`() {
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["b"] }.after { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a/b' after its child column 'a/b'"
+    }
 }

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

@@ -527,6 +527,7 @@ class DataFrameTreeTests : BaseTest() {
     @Test
     fun moveAfter() {
         val moved = typed2.move { age }.after { nameAndCity.name }
+        println(moved)
         moved.columnsCount() shouldBe 2
         moved.nameAndCity.columnsCount() shouldBe 3
         moved.nameAndCity.select { all() } shouldBe dataFrameOf(
@@ -544,6 +545,17 @@ class DataFrameTreeTests : BaseTest() {
         moved.remove { nameAndCity } shouldBe typed2.select { age and nameAndCity.name and weight }
     }
 
+    @Test
+    fun `move nested column after nested column`() {
+        val moved = typed2.move { nameAndCity.name }.after { nameAndCity.city }
+        moved.columnsCount() shouldBe 3
+        moved.nameAndCity.columnsCount() shouldBe 2
+        moved.nameAndCity.select { all() } shouldBe dataFrameOf(
+            typed2.nameAndCity.city,
+            typed2.nameAndCity.name,
+        )
+    }
+
     @Test
     fun splitFrameColumnsIntoRows() {
         val grouped = typed.groupBy { city }

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

@@ -79,6 +79,8 @@ public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathO
 
 // region after
 
+@Refine
+@Interpretable("InsertAfter0")
 public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T> = after(df.getColumnPath(column))
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)

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

@@ -11,9 +11,11 @@ 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.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.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
@@ -23,20 +25,58 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
     val removeResult = df.removeImpl(columns = columns)
 
     val targetPath = df.getColumnWithPath(column).path
+    val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath<C>().path }
+
+    // Check if any source path is a prefix of the target path
+    sourcePaths.forEach { sourcePath ->
+        val sourceSegments = sourcePath.toList()
+        val targetSegments = targetPath.toList()
+
+        if (sourceSegments.size <= targetSegments.size &&
+            sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }
+        ) {
+            throw IllegalArgumentException(
+                "Cannot move column '${sourcePath.joinToString()}' after its own child column '${targetPath.joinToString()}'",
+            )
+        }
+    }
+
     val removeRoot = removeResult.removedColumns.first().getRoot()
 
     val refNode = removeRoot.getOrPut(targetPath) {
-        val parentPath = it.dropLast(1)
-        val parent = if (parentPath.isEmpty()) df else df.getColumnGroup(parentPath)
-        val index = parent.getColumnIndex(it.last())
-        val col = df.getColumn(index)
+        val path = it.asList()
+
+        // Get parent of a target path
+        val effectivePath = path.dropLast(1)
+
+        // Get column name (last segment)
+        val columnName = path.last()
+
+        // Get the parent
+        val parent = if (effectivePath.isEmpty()) {
+            df
+        } else {
+            df.getColumnGroup(ColumnPath(effectivePath))
+        }
+
+        // Get the column index and the column itself
+        val index = parent.getColumnIndex(columnName)
+        val col = parent.getColumn(index)
+
         ColumnPosition(index, false, col)
     }
 
     val parentPath = targetPath.dropLast(1)
     val toInsert = removeResult.removedColumns.map {
-        val path = parentPath + it.name
-        ColumnToInsert(path, it.toColumnWithPath<C>().data, refNode)
+        val sourceCol = it.toColumnWithPath<C>()
+        val sourcePath = sourceCol.path
+        val path = if (sourcePath.size > 1) {
+            // If source is nested, preserve its structure under the target parent
+            parentPath + sourcePath.last()
+        } else {
+            parentPath + sourceCol.name()
+        }
+        ColumnToInsert(path, sourceCol.data, refNode)
     }
     return removeResult.df.insertImpl(toInsert)
 }

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

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.api
 
 import io.kotest.assertions.throwables.shouldNotThrowAny
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.junit.Test
@@ -91,4 +92,64 @@ class MoveTests {
             df.move("1").after("2") shouldBe dataFrameOf("2", "1")(2, 1)
         }
     }
+
+    @Test
+    fun `move after in nested structure`() {
+        val df = grouped.move { "a"["b"] }
+            .after { "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("d", "b")
+    }
+
+    @Test
+    fun `move after multiple columns`() {
+        val df = grouped.move { "a"["b"] and "b"["c"] }
+            .after { "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("d", "b", "c")
+        df["b"].asColumnGroup().columnNames() shouldBe listOf("d")
+    }
+
+    @Test
+    fun `move after with column selector`() {
+        val df = grouped.move { colsAtAnyDepth { it.name == "r" || it.name == "w" } }
+            .after { "a"["c"]["d"] }
+        df.columnNames() shouldBe listOf("q", "a", "b", "e")
+        df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "w", "r")
+    }
+
+    @Test
+    fun `move after between groups`() {
+        val df = grouped.move { "a"["b"] }.after { "b"["c"] }
+        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 after child`() {
+        // Simple case: direct parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").after { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a' after its own child column 'a/b'"
+
+        // Nested case: deeper parent-child relationship
+        shouldThrow<IllegalArgumentException> {
+            grouped.move("a").after { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a' after its own child column 'a/c/d'"
+
+        // Group case: moving group after its nested column
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["c"] }.after { "a"["c"]["d"] }
+        }.message shouldBe "Cannot move column 'a/c' after its own child column 'a/c/d'"
+    }
+
+    @Test
+    fun `should throw when moving column after itself`() {
+        shouldThrow<IllegalArgumentException> {
+            grouped.move { "a"["b"] }.after { "a"["b"] }
+        }.message shouldBe "Cannot move column 'a/b' after its own child column 'a/b'"
+    }
 }

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

@@ -544,6 +544,17 @@ class DataFrameTreeTests : BaseTest() {
         moved.remove { nameAndCity } shouldBe typed2.select { age and nameAndCity.name and weight }
     }
 
+    @Test
+    fun `move nested column after nested column`() {
+        val moved = typed2.move { nameAndCity.name }.after { nameAndCity.city }
+        moved.columnsCount() shouldBe 3
+        moved.nameAndCity.columnsCount() shouldBe 2
+        moved.nameAndCity.select { all() } shouldBe dataFrameOf(
+            typed2.nameAndCity.city,
+            typed2.nameAndCity.name,
+        )
+    }
+
     @Test
     fun splitFrameColumnsIntoRows() {
         val grouped = typed.groupBy { city }

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.Infer
+import org.jetbrains.kotlinx.dataframe.api.after
 import org.jetbrains.kotlinx.dataframe.api.insert
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.under
@@ -122,3 +123,14 @@ internal class Under4 : AbstractInterpreter<PluginDataFrameSchema>() {
             .toPluginDataFrameSchema()
     }
 }
+
+internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
+    val Arguments.column: SingleColumnApproximation by arg()
+    val Arguments.receiver: InsertClauseApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.df.asDataFrame()
+            .insert(receiver.column.asDataColumn()).after(column.col.path)
+            .toPluginDataFrameSchema()
+    }
+}