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 index a2e761dfee..507b788a4e 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/move.kt @@ -22,6 +22,7 @@ import org.jetbrains.kotlinx.dataframe.documentation.SelectingColumns import org.jetbrains.kotlinx.dataframe.impl.api.afterOrBefore import org.jetbrains.kotlinx.dataframe.impl.api.moveImpl import org.jetbrains.kotlinx.dataframe.impl.api.moveTo +import org.jetbrains.kotlinx.dataframe.impl.api.moveToImpl import org.jetbrains.kotlinx.dataframe.ncol import org.jetbrains.kotlinx.dataframe.util.DEPRECATED_ACCESS_API import org.jetbrains.kotlinx.dataframe.util.MOVE_TO_LEFT @@ -537,6 +538,32 @@ public fun MoveClause.under( @Interpretable("MoveTo") public fun MoveClause.to(columnIndex: Int): DataFrame = moveTo(columnIndex) +/** + * Moves columns, previously selected with [move] to a new position specified + * by [columnIndex]. If [insideGroup] is true, selected columns will be moved remaining within their [ColumnGroup], + * else they will be moved to the top level. + * + * Returns a new [DataFrame] with updated columns structure. + * + * For more information: {@include [DocumentationUrls.Move]} + * + * ### Examples: + * ```kotlin + * df.move { age and weight }.to(0, true) + * df.move("age", "weight").to(2, false) + * ``` + * + * @param [columnIndex] The index specifying the position in the [ColumnGroup] columns + * where the selected columns will be moved. + * + * @param [insideGroup] If true, selected columns will be moved remaining inside their group, + * else they will be moved to the top level. + */ +@Refine +@Interpretable("MoveTo") +public fun MoveClause.to(columnIndex: Int, insideGroup: Boolean): DataFrame = + moveToImpl(columnIndex, insideGroup) + /** * Moves columns, previously selected with [move] to the top-level within the [DataFrame]. * Moved columns name can be specified via special ColumnSelectionDsl. @@ -691,6 +718,26 @@ public fun MoveClause.toLeft(): DataFrame = to(0) @Interpretable("MoveToStart0") public fun MoveClause.toStart(): DataFrame = to(0) +/** + * If insideGroup is true, moves columns previously selected with [move] to the start of their [ColumnGroup]. + * Else, selected columns will be moved to the start of their [DataFrame] (to the top-level). + * + * Returns a new [DataFrame] with updated columns. + * + * For more information: {@include [DocumentationUrls.Move]} + * + * ### Examples: + * ```kotlin + * df.move { age and weight }.toStart(true) + * df.move { colsOf() }.toStart(true) + * df.move("age", "weight").toStart(false) + * ``` + * + * @param [insideGroup] If true, selected columns will be moved to the start remaining inside their group, + * else they will be moved to the start on top level. + */ +public fun MoveClause.toStart(insideGroup: Boolean): DataFrame = to(0, insideGroup) + @Deprecated(TO_RIGHT, ReplaceWith(TO_RIGHT_REPLACE), DeprecationLevel.ERROR) public fun MoveClause.toRight(): DataFrame = to(df.ncol) @@ -712,6 +759,28 @@ public fun MoveClause.toRight(): DataFrame = to(df.ncol) @Interpretable("MoveToEnd0") public fun MoveClause.toEnd(): DataFrame = to(df.ncol) +/** + * If insideGroup is true, moves columns previously selected with [move] to the end of their [ColumnGroup]. + * Else, selected columns will be moved to the end of their [DataFrame] (to the top-level). + * + * Returns a new [DataFrame] with updated columns. + * + * For more information: {@include [DocumentationUrls.Move]} + * + * ### Examples: + * ```kotlin + * df.move { age and weight }.toEnd(true) + * df.move { colsOf() }.toEnd(true) + * df.move("age", "weight").toEnd(false) + * ``` + * + * @param [insideGroup] If true, selected columns will be moved to the end remaining inside their group, + * else they will be moved to the end on top level. + */ +@Refine +@Interpretable("MoveToEnd0") +public fun MoveClause.toEnd(insideGroup: Boolean): DataFrame = to(df.ncol, insideGroup) + /** * An intermediate class used in the [move] operation. * 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 7d874eca3a..61365104b6 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 @@ -8,12 +8,16 @@ 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.asDataFrame 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.getColumns import org.jetbrains.kotlinx.dataframe.api.move +import org.jetbrains.kotlinx.dataframe.api.to import org.jetbrains.kotlinx.dataframe.api.toDataFrame +import org.jetbrains.kotlinx.dataframe.api.toPath import org.jetbrains.kotlinx.dataframe.columns.ColumnPath import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy @@ -23,7 +27,9 @@ 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.impl.last import org.jetbrains.kotlinx.dataframe.path +import kotlin.collections.first internal fun MoveClause.afterOrBefore(column: ColumnSelector, isAfter: Boolean): DataFrame { val removeResult = df.removeImpl(columns = columns) @@ -124,3 +130,44 @@ internal fun MoveClause.moveTo(columnIndex: Int): DataFrame { } return newColumnList.toDataFrame().cast() } + +internal fun MoveClause.moveToImpl(columnIndex: Int, insideGroup: Boolean): DataFrame { + if (!insideGroup) { + return moveTo(columnIndex) + } + + val columnsToMove = df.getColumns(columns) + + // check if columns to move have the same parent + val columnsToMoveParents = columnsToMove.map { it.path.dropLast() } + val parentOfFirst = columnsToMoveParents.first() + if (columnsToMoveParents.any { it != parentOfFirst }) { + throw IllegalArgumentException( + "Cannot move columns with different parent to an index", + ) + } + + // if columns will be moved to top level or columns to move are at top level + if (parentOfFirst.isEmpty()) { + return moveTo(columnIndex) + } + + // logic: remove columns to move and their siblings (from this point, sons), apply them moveTo, reinsert them + val parentPath = df[parentOfFirst].path + val sons = df[parentOfFirst].asColumnGroup() + // remove sons + val sonsWithFullPaths = sons.columns().map { parentPath + it.path } + val intermediateDf = df.removeImpl { sonsWithFullPaths.toColumnSet() } + // move sons and reinsert them + val columnsToMoveWithReducedPath = columnsToMove.map { it.path.last(it.path.size - parentPath.size).toPath() } + val sonsHaveBeenMoved = sons.asDataFrame().move { + columnsToMoveWithReducedPath.toColumnSet() + }.to(columnIndex).columns() + val sonsToInsert = sonsHaveBeenMoved.map { ColumnToInsert(parentPath + it.path, it) } + val secondIntermediateDf = intermediateDf.df.insertImpl(sonsToInsert) + // nested level is good but order of top level is changed -> need to fix it + val rootOfColumnsToMove = df[listOf(parentPath.first()).toPath()] + val indexOfRootOfColumnsToMove = df.columns().indexOf(rootOfColumnsToMove) + val finalDf = secondIntermediateDf.move { listOf(parentPath.first()).toPath() }.to(indexOfRootOfColumnsToMove) + return finalDf +} 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 955410e832..fafcd1bce2 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 @@ -220,4 +220,96 @@ class MoveTests { grouped.move { "a"["b"] }.before { "a"["b"] } }.message shouldBe "Cannot move column 'a/b' before its own child column 'a/b'" } + + @Test + fun `move single nested column to the start remaining inside the group`() { + val df = grouped.move { "b"["d"] }.to(0, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r") + df["b"].asColumnGroup().columnNames() shouldBe listOf("d", "c") + } + + @Test + fun `move single nested column to the end remaining inside the group`() { + val df = grouped.move { "b"["c"] }.to(2, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r") + df["b"].asColumnGroup().columnNames() shouldBe listOf("d", "c") + } + + @Test + fun `move single nested column between columns remaining inside the group`() { + // creating an appropriate df for the test + val groupedModified = grouped.move("r").before { "b"["c"] } + groupedModified["b"].asColumnGroup().columnNames() shouldBe listOf("r", "c", "d") + // test itself + val df = groupedModified.move { "b"["r"] }.to(1, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e") + df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "r", "d") + } + + @Test + fun `move single nested column to the end remaining inside the group, need to switch group's columns`() { + val df = grouped.move { "b"["c"] }.to(1, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r") + df["b"].asColumnGroup().columnNames() shouldBe listOf("d", "c") + } + + @Test + fun `move single nested column to current index of the column itself`() { + val df = grouped.move { "b"["d"] }.to(1, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e", "r") + df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "d") + } + + @Test + fun `move multiple nested columns to the start`() { + // creating an appropriate df for the test + val groupedModified = grouped.move("r").before { "b"["c"] } + groupedModified["b"].asColumnGroup().columnNames() shouldBe listOf("r", "c", "d") + // test itself + val df = groupedModified.move { "b"["c"] and "b"["d"] }.to(0, true) + df.columnNames() shouldBe listOf("q", "a", "b", "w", "e") + df["b"].asColumnGroup().columnNames() shouldBe listOf("c", "d", "r") + } + + @Test + fun `move multiple non bordering nested columns`() { + // creating an appropriate df for the test + val groupedModified = grouped.move("r", "q").before { "b"["c"] } + groupedModified["b"].asColumnGroup().columnNames() shouldBe listOf("r", "q", "c", "d") + // test itself + val df = groupedModified.move { "b"["r"] and "b"["d"] }.to(1, true) + df.columnNames() shouldBe listOf("a", "b", "w", "e") + df["b"].asColumnGroup().columnNames() shouldBe listOf("q", "r", "d", "c") + } + + @Test + fun `move single top level column to the start, insideGroup should make no difference`() { + // insideGroup is true + val dfInsideGroupIsTrue = grouped.move("e").to(0, true) + dfInsideGroupIsTrue.columnNames() shouldBe listOf("e", "q", "a", "b", "w", "r") + dfInsideGroupIsTrue["e"].asColumnGroup().columnNames() shouldBe listOf("f") + // insideGroup is false + val dfInsideGroupIsFalse = grouped.move("e").to(0, false) + dfInsideGroupIsFalse.columnNames() shouldBe listOf("e", "q", "a", "b", "w", "r") + dfInsideGroupIsFalse["e"].asColumnGroup().columnNames() shouldBe listOf("f") + } + + @Test + fun `move multiple top level columns between columns, insideGroup should make no difference`() { + // insideGroup is true + val dfInsideGroupIsTrue = grouped.move("w", "e").to(1, true) + dfInsideGroupIsTrue.columnNames() shouldBe listOf("q", "w", "e", "a", "b", "r") + dfInsideGroupIsTrue["e"].asColumnGroup().columnNames() shouldBe listOf("f") + // insideGroup is false + val dfInsideGroupIsFalse = grouped.move("w", "e").to(1, false) + dfInsideGroupIsFalse.columnNames() shouldBe listOf("q", "w", "e", "a", "b", "r") + dfInsideGroupIsFalse["e"].asColumnGroup().columnNames() shouldBe listOf("f") + } + + @Test + fun `should throw when moving columns of different groups`() { + shouldThrow { + grouped.move { "a"["b"] and "b"["c"] }.to(0, true) + }.message shouldBe "Cannot move columns with different parent to an index" + } }