From f0c2ad1e44b3b6e0db62f2c5867943701c7ec012 Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 12:28:46 +0400 Subject: [PATCH 1/8] fix move for nested columns --- .../kotlinx/dataframe/impl/api/move.kt | 47 ++++++++++++++++--- .../testSets/person/DataFrameTreeTests.kt | 12 +++++ 2 files changed, 53 insertions(+), 6 deletions(-) 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 index cacbe2e028..7dbf0e92de 100644 --- 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 @@ -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 @@ -26,17 +27,51 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, 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().data, refNode) + val sourceCol = it.toColumnWithPath() + 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) } 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 index 0a757d4acc..94a68c716e 100644 --- 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 @@ -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 } From d8691ad7e4e99cce9297038b042d1097087a715c Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 12:54:12 +0400 Subject: [PATCH 2/8] move after tests --- .../jetbrains/kotlinx/dataframe/api/move.kt | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) 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 index 66a368144c..1cae51d555 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt @@ -91,4 +91,39 @@ 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") + } } From 2aaade5af481a10439d2d63d0c359f6776d721da Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 13:05:01 +0400 Subject: [PATCH 3/8] move after check --- .../kotlinx/dataframe/impl/api/move.kt | 14 ++++++++++ .../jetbrains/kotlinx/dataframe/api/move.kt | 26 +++++++++++++++++++ 2 files changed, 40 insertions(+) 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 index 7dbf0e92de..818e5c9c52 100644 --- 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 @@ -24,6 +24,20 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, val removeResult = df.removeImpl(columns = columns) val targetPath = df.getColumnWithPath(column).path + val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath().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) { 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 index 1cae51d555..541b6573b6 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/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 @@ -126,4 +127,29 @@ class MoveTests { 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 { + 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 { + 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 { + 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 { + grouped.move { "a"["b"] }.after { "a"["b"] } + }.message shouldBe "Cannot move column 'a/b' after its child column 'a/b'" + } } From eca9f8aa846cc03d6f9cf581c5b6959f1fbecf07 Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 13:14:37 +0400 Subject: [PATCH 4/8] move gen src --- .../kotlinx/dataframe/impl/api/move.kt | 63 +++++++++++++++++-- .../jetbrains/kotlinx/dataframe/api/move.kt | 61 ++++++++++++++++++ .../testSets/person/DataFrameTreeTests.kt | 12 ++++ .../kotlinx/dataframe/impl/api/move.kt | 8 ++- .../jetbrains/kotlinx/dataframe/api/move.kt | 2 +- .../testSets/person/DataFrameTreeTests.kt | 2 +- 6 files changed, 137 insertions(+), 11 deletions(-) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt index cacbe2e028..e6dc28211a 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt +++ b/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 MoveClause.afterOrBefore(column: ColumnSelector, val removeResult = df.removeImpl(columns = columns) val targetPath = df.getColumnWithPath(column).path + val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath().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().data, refNode) + val sourceCol = it.toColumnWithPath() + 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) } diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt index 66a368144c..72f581244e 100644 --- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/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 { + 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 { + 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 { + 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 { + grouped.move { "a"["b"] }.after { "a"["b"] } + }.message shouldBe "Cannot move column 'a/b' after its child column 'a/b'" + } } diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt index 0a757d4acc..6e5be2fc21 100644 --- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt +++ 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 } 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 index 818e5c9c52..e6dc28211a 100644 --- 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 @@ -31,10 +31,12 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, val sourceSegments = sourcePath.toList() val targetSegments = targetPath.toList() - if (sourceSegments.size <= targetSegments.size && - sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }) { + 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()}'") + "Cannot move column '${sourcePath.joinToString()}' after its child column '${targetPath.joinToString()}'", + ) } } 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 index 541b6573b6..72f581244e 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt @@ -116,7 +116,7 @@ class MoveTests { 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.columnNames() shouldBe listOf("q", "a", "b", "e") df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "w", "r") } 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 index 94a68c716e..6e5be2fc21 100644 --- 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 @@ -552,7 +552,7 @@ class DataFrameTreeTests : BaseTest() { moved.nameAndCity.columnsCount() shouldBe 2 moved.nameAndCity.select { all() } shouldBe dataFrameOf( typed2.nameAndCity.city, - typed2.nameAndCity.name + typed2.nameAndCity.name, ) } From 8b3e0c139a8a4aaf6526636a2e055cd254984bec Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 14:26:55 +0400 Subject: [PATCH 5/8] insert after plugin support --- .../jetbrains/kotlinx/dataframe/api/insert.kt | 2 ++ .../dataframe/plugin/impl/api/insert.kt | 12 ++++++++++++ .../dataframe/plugin/loadInterpreter.kt | 2 ++ .../testData/box/insertAfter.kt | 12 ++++++++++++ .../testData/box/insertAfterNested.kt | 13 +++++++++++++ .../testData/box/moveAfterNested.kt | 13 +++++++++++++ .../DataFrameBlackBoxCodegenTestGenerated.java | 18 ++++++++++++++++++ 7 files changed, 72 insertions(+) create mode 100644 plugins/kotlin-dataframe/testData/box/insertAfter.kt create mode 100644 plugins/kotlin-dataframe/testData/box/insertAfterNested.kt create mode 100644 plugins/kotlin-dataframe/testData/box/moveAfterNested.kt 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 index 03efde6593..74fece6d86 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt @@ -79,6 +79,8 @@ public fun InsertClause.under(column: String): DataFrame = under(pathO // region after +@Refine +@Interpretable("InsertAfter0") public fun InsertClause.after(column: ColumnSelector): DataFrame = after(df.getColumnPath(column)) public fun InsertClause.after(column: String): DataFrame = df.add(this.column).move(this.column).after(column) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/insert.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/insert.kt index e4f257e725..b34c33e181 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/insert.kt +++ b/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() { .toPluginDataFrameSchema() } } + +internal class InsertAfter0 : AbstractInterpreter() { + 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() + } +} 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 index cb4244f19e..44ec1d1f55 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt @@ -128,6 +128,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceExpression import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReduceInto import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByReducePredicate import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByXs +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.InsertAfter0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last1 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Last2 @@ -286,6 +287,7 @@ internal inline fun String.load(): T { "Under2" -> Under2() "Under3" -> Under3() "Under4" -> Under4() + "InsertAfter0" -> InsertAfter0() "Join0" -> Join0() "Match0" -> Match0() "Rename" -> Rename() diff --git a/plugins/kotlin-dataframe/testData/box/insertAfter.kt b/plugins/kotlin-dataframe/testData/box/insertAfter.kt new file mode 100644 index 0000000000..10a1877790 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/insertAfter.kt @@ -0,0 +1,12 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.api.* + +fun box(): String { + val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6) + + val dfWithD = df.insert("d") { b * c }.after { a } + + val dCol: DataColumn = dfWithD.d + + return "OK" +} diff --git a/plugins/kotlin-dataframe/testData/box/insertAfterNested.kt b/plugins/kotlin-dataframe/testData/box/insertAfterNested.kt new file mode 100644 index 0000000000..678c229779 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/insertAfterNested.kt @@ -0,0 +1,13 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.api.* + +fun box(): String { + val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6) + val grouped = df.group { a and b and c }.into("g") + + val dfWithD = grouped.insert("d") { g.b * g.c }.after { g.b } + + val dCol: DataColumn = dfWithD.g.d + + return "OK" +} diff --git a/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt b/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt new file mode 100644 index 0000000000..bb600efdaf --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt @@ -0,0 +1,13 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.api.* + +fun box(): String { + val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6) + val grouped = df.group { a and b }.into("g") + grouped.move { g.a }.after { g.b } + + val dfMovedAfterNested = grouped.move { c }.after { g.a } + val cNested: DataColumn = dfMovedAfterNested.g.c + + return "OK" +} diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java index 5d030bbf00..1710db840f 100644 --- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java +++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java @@ -286,6 +286,18 @@ public void testInsert() { runTest("testData/box/insert.kt"); } + @Test + @TestMetadata("insertAfter.kt") + public void testInsertAfter() { + runTest("testData/box/insertAfter.kt"); + } + + @Test + @TestMetadata("insertAfterNested.kt") + public void testInsertAfterNested() { + runTest("testData/box/insertAfterNested.kt"); + } + @Test @TestMetadata("inventNamesForLocalClasses.kt") public void testInventNamesForLocalClasses() { @@ -346,6 +358,12 @@ public void testMoveAfter() { runTest("testData/box/moveAfter.kt"); } + @Test + @TestMetadata("moveAfterNested.kt") + public void testMoveAfterNested() { + runTest("testData/box/moveAfterNested.kt"); + } + @Test @TestMetadata("moveInto.kt") public void testMoveInto() { From f3a71f21d46dc37a3817f14b7c01cfcd252d6348 Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Tue, 11 Mar 2025 14:31:30 +0400 Subject: [PATCH 6/8] insert gen src --- .../main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt index 03efde6593..74fece6d86 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt @@ -79,6 +79,8 @@ public fun InsertClause.under(column: String): DataFrame = under(pathO // region after +@Refine +@Interpretable("InsertAfter0") public fun InsertClause.after(column: ColumnSelector): DataFrame = after(df.getColumnPath(column)) public fun InsertClause.after(column: String): DataFrame = df.add(this.column).move(this.column).after(column) From 46392c39d9b19b58c968982b7de5ad71bf73a474 Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Mon, 17 Mar 2025 14:52:03 +0400 Subject: [PATCH 7/8] move fixes --- .../kotlin/org/jetbrains/kotlinx/dataframe/impl/api/move.kt | 5 +++-- .../kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt | 1 - plugins/kotlin-dataframe/testData/box/insertAfter.kt | 2 +- plugins/kotlin-dataframe/testData/box/moveAfterNested.kt | 1 - 4 files changed, 4 insertions(+), 5 deletions(-) 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 index e6dc28211a..027a5a161e 100644 --- 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 @@ -15,6 +15,7 @@ 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 @@ -35,7 +36,7 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] } ) { throw IllegalArgumentException( - "Cannot move column '${sourcePath.joinToString()}' after its child column '${targetPath.joinToString()}'", + "Cannot move column '${sourcePath.joinToString()}' after its own child column '${targetPath.joinToString()}'", ) } } @@ -43,7 +44,7 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, val removeRoot = removeResult.removedColumns.first().getRoot() val refNode = removeRoot.getOrPut(targetPath) { - val path = it.toList() + val path = it.asList() // Find the group reference (if any) and adjust the path val groupRefIndex = path.indexOfFirst { it.isEmpty() } 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 index 6e5be2fc21..71c4681285 100644 --- 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 @@ -527,7 +527,6 @@ 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( diff --git a/plugins/kotlin-dataframe/testData/box/insertAfter.kt b/plugins/kotlin-dataframe/testData/box/insertAfter.kt index 10a1877790..73c26a3020 100644 --- a/plugins/kotlin-dataframe/testData/box/insertAfter.kt +++ b/plugins/kotlin-dataframe/testData/box/insertAfter.kt @@ -4,7 +4,7 @@ import org.jetbrains.kotlinx.dataframe.api.* fun box(): String { val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6) - val dfWithD = df.insert("d") { b * c }.after { a } + val dfWithD = df.insert("d") { b * c }.after { a } val dCol: DataColumn = dfWithD.d diff --git a/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt b/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt index bb600efdaf..e4c924f295 100644 --- a/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt +++ b/plugins/kotlin-dataframe/testData/box/moveAfterNested.kt @@ -4,7 +4,6 @@ import org.jetbrains.kotlinx.dataframe.api.* fun box(): String { val df = dataFrameOf("a", "b", "c")(1, 2, 3, 4, 5, 6) val grouped = df.group { a and b }.into("g") - grouped.move { g.a }.after { g.b } val dfMovedAfterNested = grouped.move { c }.after { g.a } val cNested: DataColumn = dfMovedAfterNested.g.c From 45cedfd2ec2cc41b38f513e834448c7337b1e5e2 Mon Sep 17 00:00:00 2001 From: "andrei.kislitsyn" Date: Mon, 17 Mar 2025 15:10:49 +0400 Subject: [PATCH 8/8] move impl fixes --- .../kotlinx/dataframe/impl/api/move.kt | 22 +++++-------------- .../jetbrains/kotlinx/dataframe/api/move.kt | 8 +++---- 2 files changed, 9 insertions(+), 21 deletions(-) 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 index 027a5a161e..277563a67d 100644 --- 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 @@ -46,25 +46,13 @@ internal fun MoveClause.afterOrBefore(column: ColumnSelector, val refNode = removeRoot.getOrPut(targetPath) { val path = it.asList() - // Find the group reference (if any) and adjust the path - val groupRefIndex = path.indexOfFirst { it.isEmpty() } + // Get parent of a target path + val effectivePath = path.dropLast(1) - // 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 column name (last segment) + val columnName = path.last() - // Get the parent group using ColumnPath + // Get the parent val parent = if (effectivePath.isEmpty()) { df } else { 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 index 72f581244e..f78d83ead7 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt @@ -133,23 +133,23 @@ class MoveTests { // Simple case: direct parent-child relationship shouldThrow { grouped.move("a").after { "a"["b"] } - }.message shouldBe "Cannot move column 'a' after its child column 'a/b'" + }.message shouldBe "Cannot move column 'a' after its own child column 'a/b'" // Nested case: deeper parent-child relationship shouldThrow { grouped.move("a").after { "a"["c"]["d"] } - }.message shouldBe "Cannot move column 'a' after its child column '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 { grouped.move { "a"["c"] }.after { "a"["c"]["d"] } - }.message shouldBe "Cannot move column 'a/c' after its child column '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 { grouped.move { "a"["b"] }.after { "a"["b"] } - }.message shouldBe "Cannot move column 'a/b' after its child column 'a/b'" + }.message shouldBe "Cannot move column 'a/b' after its own child column 'a/b'" } }