Skip to content

Commit 40e1db1

Browse files
AndreiKingsleyAndreiKingsley
authored
Merge pull request #1092 from Kotlin/move_fix
Move/Insert After fix
2 parents 76dd759 + 45cedfd commit 40e1db1

File tree

14 files changed

+321
-12
lines changed

14 files changed

+321
-12
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathO
7979

8080
// region after
8181

82+
@Refine
83+
@Interpretable("InsertAfter0")
8284
public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T> = after(df.getColumnPath(column))
8385

8486
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‎

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import org.jetbrains.kotlinx.dataframe.api.getColumn
1111
import org.jetbrains.kotlinx.dataframe.api.getColumnGroup
1212
import org.jetbrains.kotlinx.dataframe.api.getColumnWithPath
1313
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
14+
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1415
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
1516
import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
1617
import org.jetbrains.kotlinx.dataframe.impl.DataFrameReceiver
@@ -23,20 +24,70 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
2324
val removeResult = df.removeImpl(columns = columns)
2425

2526
val targetPath = df.getColumnWithPath(column).path
27+
val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath<C>().path }
28+
29+
// Check if any source path is a prefix of the target path
30+
sourcePaths.forEach { sourcePath ->
31+
val sourceSegments = sourcePath.toList()
32+
val targetSegments = targetPath.toList()
33+
34+
if (sourceSegments.size <= targetSegments.size &&
35+
sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }
36+
) {
37+
throw IllegalArgumentException(
38+
"Cannot move column '${sourcePath.joinToString()}' after its child column '${targetPath.joinToString()}'",
39+
)
40+
}
41+
}
42+
2643
val removeRoot = removeResult.removedColumns.first().getRoot()
2744

2845
val refNode = removeRoot.getOrPut(targetPath) {
29-
val parentPath = it.dropLast(1)
30-
val parent = if (parentPath.isEmpty()) df else df.getColumnGroup(parentPath)
31-
val index = parent.getColumnIndex(it.last())
32-
val col = df.getColumn(index)
46+
val path = it.toList()
47+
48+
// Find the group reference (if any) and adjust the path
49+
val groupRefIndex = path.indexOfFirst { it.isEmpty() }
50+
51+
// Calculate effective path and column name based on group reference
52+
val effectivePath = if (groupRefIndex >= 0) {
53+
// Nested column reference
54+
path.take(groupRefIndex)
55+
} else {
56+
path.dropLast(1)
57+
}
58+
59+
val columnName = if (groupRefIndex >= 0) {
60+
// Use the next segment after group reference, or previous if no next segment
61+
path.getOrNull(groupRefIndex + 1) ?: path[groupRefIndex - 1]
62+
} else {
63+
path.last()
64+
}
65+
66+
// Get the parent group using ColumnPath
67+
val parent = if (effectivePath.isEmpty()) {
68+
df
69+
} else {
70+
df.getColumnGroup(ColumnPath(effectivePath))
71+
}
72+
73+
// Get the column index and the column itself
74+
val index = parent.getColumnIndex(columnName)
75+
val col = parent.getColumn(index)
76+
3377
ColumnPosition(index, false, col)
3478
}
3579

3680
val parentPath = targetPath.dropLast(1)
3781
val toInsert = removeResult.removedColumns.map {
38-
val path = parentPath + it.name
39-
ColumnToInsert(path, it.toColumnWithPath<C>().data, refNode)
82+
val sourceCol = it.toColumnWithPath<C>()
83+
val sourcePath = sourceCol.path
84+
val path = if (sourcePath.size > 1) {
85+
// If source is nested, preserve its structure under the target parent
86+
parentPath + sourcePath.last()
87+
} else {
88+
parentPath + sourceCol.name()
89+
}
90+
ColumnToInsert(path, sourceCol.data, refNode)
4091
}
4192
return removeResult.df.insertImpl(toInsert)
4293
}

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.assertions.throwables.shouldNotThrowAny
4+
import io.kotest.assertions.throwables.shouldThrow
45
import io.kotest.matchers.shouldBe
56
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
67
import org.junit.Test
@@ -91,4 +92,64 @@ class MoveTests {
9192
df.move("1").after("2") shouldBe dataFrameOf("2", "1")(2, 1)
9293
}
9394
}
95+
96+
@Test
97+
fun `move after in nested structure`() {
98+
val df = grouped.move { "a"["b"] }
99+
.after { "a"["c"]["d"] }
100+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
101+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
102+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "b")
103+
}
104+
105+
@Test
106+
fun `move after multiple columns`() {
107+
val df = grouped.move { "a"["b"] and "b"["c"] }
108+
.after { "a"["c"]["d"] }
109+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
110+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
111+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "b", "c")
112+
df["b"].asColumnGroup().columnNames() shouldBe listOf("d")
113+
}
114+
115+
@Test
116+
fun `move after with column selector`() {
117+
val df = grouped.move { colsAtAnyDepth { it.name == "r" || it.name == "w" } }
118+
.after { "a"["c"]["d"] }
119+
df.columnNames() shouldBe listOf("q", "a", "b", "e")
120+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "w", "r")
121+
}
122+
123+
@Test
124+
fun `move after between groups`() {
125+
val df = grouped.move { "a"["b"] }.after { "b"["c"] }
126+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
127+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
128+
df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "b", "d")
129+
}
130+
131+
@Test
132+
fun `should throw when moving parent after child`() {
133+
// Simple case: direct parent-child relationship
134+
shouldThrow<IllegalArgumentException> {
135+
grouped.move("a").after { "a"["b"] }
136+
}.message shouldBe "Cannot move column 'a' after its child column 'a/b'"
137+
138+
// Nested case: deeper parent-child relationship
139+
shouldThrow<IllegalArgumentException> {
140+
grouped.move("a").after { "a"["c"]["d"] }
141+
}.message shouldBe "Cannot move column 'a' after its child column 'a/c/d'"
142+
143+
// Group case: moving group after its nested column
144+
shouldThrow<IllegalArgumentException> {
145+
grouped.move { "a"["c"] }.after { "a"["c"]["d"] }
146+
}.message shouldBe "Cannot move column 'a/c' after its child column 'a/c/d'"
147+
}
148+
149+
@Test
150+
fun `should throw when moving column after itself`() {
151+
shouldThrow<IllegalArgumentException> {
152+
grouped.move { "a"["b"] }.after { "a"["b"] }
153+
}.message shouldBe "Cannot move column 'a/b' after its child column 'a/b'"
154+
}
94155
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,7 @@ class DataFrameTreeTests : BaseTest() {
527527
@Test
528528
fun moveAfter() {
529529
val moved = typed2.move { age }.after { nameAndCity.name }
530+
println(moved)
530531
moved.columnsCount() shouldBe 2
531532
moved.nameAndCity.columnsCount() shouldBe 3
532533
moved.nameAndCity.select { all() } shouldBe dataFrameOf(
@@ -544,6 +545,17 @@ class DataFrameTreeTests : BaseTest() {
544545
moved.remove { nameAndCity } shouldBe typed2.select { age and nameAndCity.name and weight }
545546
}
546547

548+
@Test
549+
fun `move nested column after nested column`() {
550+
val moved = typed2.move { nameAndCity.name }.after { nameAndCity.city }
551+
moved.columnsCount() shouldBe 3
552+
moved.nameAndCity.columnsCount() shouldBe 2
553+
moved.nameAndCity.select { all() } shouldBe dataFrameOf(
554+
typed2.nameAndCity.city,
555+
typed2.nameAndCity.name,
556+
)
557+
}
558+
547559
@Test
548560
fun splitFrameColumnsIntoRows() {
549561
val grouped = typed.groupBy { city }

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathO
7979

8080
// region after
8181

82+
@Refine
83+
@Interpretable("InsertAfter0")
8284
public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T> = after(df.getColumnPath(column))
8385

8486
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‎

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ import org.jetbrains.kotlinx.dataframe.api.getColumn
1111
import org.jetbrains.kotlinx.dataframe.api.getColumnGroup
1212
import org.jetbrains.kotlinx.dataframe.api.getColumnWithPath
1313
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
14+
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1415
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
1516
import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
1617
import org.jetbrains.kotlinx.dataframe.impl.DataFrameReceiver
18+
import org.jetbrains.kotlinx.dataframe.impl.asList
1719
import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnWithPath
1820
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ColumnPosition
1921
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.getOrPut
@@ -23,20 +25,58 @@ internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>,
2325
val removeResult = df.removeImpl(columns = columns)
2426

2527
val targetPath = df.getColumnWithPath(column).path
28+
val sourcePaths = removeResult.removedColumns.map { it.toColumnWithPath<C>().path }
29+
30+
// Check if any source path is a prefix of the target path
31+
sourcePaths.forEach { sourcePath ->
32+
val sourceSegments = sourcePath.toList()
33+
val targetSegments = targetPath.toList()
34+
35+
if (sourceSegments.size <= targetSegments.size &&
36+
sourceSegments.indices.all { targetSegments[it] == sourceSegments[it] }
37+
) {
38+
throw IllegalArgumentException(
39+
"Cannot move column '${sourcePath.joinToString()}' after its own child column '${targetPath.joinToString()}'",
40+
)
41+
}
42+
}
43+
2644
val removeRoot = removeResult.removedColumns.first().getRoot()
2745

2846
val refNode = removeRoot.getOrPut(targetPath) {
29-
val parentPath = it.dropLast(1)
30-
val parent = if (parentPath.isEmpty()) df else df.getColumnGroup(parentPath)
31-
val index = parent.getColumnIndex(it.last())
32-
val col = df.getColumn(index)
47+
val path = it.asList()
48+
49+
// Get parent of a target path
50+
val effectivePath = path.dropLast(1)
51+
52+
// Get column name (last segment)
53+
val columnName = path.last()
54+
55+
// Get the parent
56+
val parent = if (effectivePath.isEmpty()) {
57+
df
58+
} else {
59+
df.getColumnGroup(ColumnPath(effectivePath))
60+
}
61+
62+
// Get the column index and the column itself
63+
val index = parent.getColumnIndex(columnName)
64+
val col = parent.getColumn(index)
65+
3366
ColumnPosition(index, false, col)
3467
}
3568

3669
val parentPath = targetPath.dropLast(1)
3770
val toInsert = removeResult.removedColumns.map {
38-
val path = parentPath + it.name
39-
ColumnToInsert(path, it.toColumnWithPath<C>().data, refNode)
71+
val sourceCol = it.toColumnWithPath<C>()
72+
val sourcePath = sourceCol.path
73+
val path = if (sourcePath.size > 1) {
74+
// If source is nested, preserve its structure under the target parent
75+
parentPath + sourcePath.last()
76+
} else {
77+
parentPath + sourceCol.name()
78+
}
79+
ColumnToInsert(path, sourceCol.data, refNode)
4080
}
4181
return removeResult.df.insertImpl(toInsert)
4282
}

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.assertions.throwables.shouldNotThrowAny
4+
import io.kotest.assertions.throwables.shouldThrow
45
import io.kotest.matchers.shouldBe
56
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
67
import org.junit.Test
@@ -91,4 +92,64 @@ class MoveTests {
9192
df.move("1").after("2") shouldBe dataFrameOf("2", "1")(2, 1)
9293
}
9394
}
95+
96+
@Test
97+
fun `move after in nested structure`() {
98+
val df = grouped.move { "a"["b"] }
99+
.after { "a"["c"]["d"] }
100+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
101+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
102+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "b")
103+
}
104+
105+
@Test
106+
fun `move after multiple columns`() {
107+
val df = grouped.move { "a"["b"] and "b"["c"] }
108+
.after { "a"["c"]["d"] }
109+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
110+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
111+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "b", "c")
112+
df["b"].asColumnGroup().columnNames() shouldBe listOf("d")
113+
}
114+
115+
@Test
116+
fun `move after with column selector`() {
117+
val df = grouped.move { colsAtAnyDepth { it.name == "r" || it.name == "w" } }
118+
.after { "a"["c"]["d"] }
119+
df.columnNames() shouldBe listOf("q", "a", "b", "e")
120+
df["a"]["c"].asColumnGroup().columnNames() shouldBe listOf("d", "w", "r")
121+
}
122+
123+
@Test
124+
fun `move after between groups`() {
125+
val df = grouped.move { "a"["b"] }.after { "b"["c"] }
126+
df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r")
127+
df["a"].asColumnGroup().columnNames() shouldBe listOf("c")
128+
df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "b", "d")
129+
}
130+
131+
@Test
132+
fun `should throw when moving parent after child`() {
133+
// Simple case: direct parent-child relationship
134+
shouldThrow<IllegalArgumentException> {
135+
grouped.move("a").after { "a"["b"] }
136+
}.message shouldBe "Cannot move column 'a' after its own child column 'a/b'"
137+
138+
// Nested case: deeper parent-child relationship
139+
shouldThrow<IllegalArgumentException> {
140+
grouped.move("a").after { "a"["c"]["d"] }
141+
}.message shouldBe "Cannot move column 'a' after its own child column 'a/c/d'"
142+
143+
// Group case: moving group after its nested column
144+
shouldThrow<IllegalArgumentException> {
145+
grouped.move { "a"["c"] }.after { "a"["c"]["d"] }
146+
}.message shouldBe "Cannot move column 'a/c' after its own child column 'a/c/d'"
147+
}
148+
149+
@Test
150+
fun `should throw when moving column after itself`() {
151+
shouldThrow<IllegalArgumentException> {
152+
grouped.move { "a"["b"] }.after { "a"["b"] }
153+
}.message shouldBe "Cannot move column 'a/b' after its own child column 'a/b'"
154+
}
94155
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,17 @@ class DataFrameTreeTests : BaseTest() {
544544
moved.remove { nameAndCity } shouldBe typed2.select { age and nameAndCity.name and weight }
545545
}
546546

547+
@Test
548+
fun `move nested column after nested column`() {
549+
val moved = typed2.move { nameAndCity.name }.after { nameAndCity.city }
550+
moved.columnsCount() shouldBe 3
551+
moved.nameAndCity.columnsCount() shouldBe 2
552+
moved.nameAndCity.select { all() } shouldBe dataFrameOf(
553+
typed2.nameAndCity.city,
554+
typed2.nameAndCity.name,
555+
)
556+
}
557+
547558
@Test
548559
fun splitFrameColumnsIntoRows() {
549560
val grouped = typed.groupBy { city }

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

33
import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.api.Infer
5+
import org.jetbrains.kotlinx.dataframe.api.after
56
import org.jetbrains.kotlinx.dataframe.api.insert
67
import org.jetbrains.kotlinx.dataframe.api.pathOf
78
import org.jetbrains.kotlinx.dataframe.api.under
@@ -122,3 +123,14 @@ internal class Under4 : AbstractInterpreter<PluginDataFrameSchema>() {
122123
.toPluginDataFrameSchema()
123124
}
124125
}
126+
127+
internal class InsertAfter0 : AbstractInterpreter<PluginDataFrameSchema>() {
128+
val Arguments.column: SingleColumnApproximation by arg()
129+
val Arguments.receiver: InsertClauseApproximation by arg()
130+
131+
override fun Arguments.interpret(): PluginDataFrameSchema {
132+
return receiver.df.asDataFrame()
133+
.insert(receiver.column.asDataColumn()).after(column.col.path)
134+
.toPluginDataFrameSchema()
135+
}
136+
}

0 commit comments

Comments
 (0)