Move to(Start) keeping inside group

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
@@ -219,6 +220,53 @@ public fun <T> DataFrame<T>.moveTo(newColumnIndex: Int, vararg columns: AnyColum
 public fun <T> DataFrame<T>.moveTo(newColumnIndex: Int, vararg columns: KProperty<*>): DataFrame<T> =
     moveTo(newColumnIndex) { columns.toColumnSet() }
 
+/**
+ * Moves the specified [columns\] 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.
+ *
+ * @include [CommonMoveToDocs]
+ * @include [SelectingColumns.Dsl] {@include [SetMoveToOperationArg]}
+ * ### Examples:
+ * ```kotlin
+ * df.moveTo(0, true) { length and age }
+ * df.moveTo(2, false) { cols(1..5) }
+ * ```
+ * @param [newColumnIndex] The index specifying the position in the [DataFrame] 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.
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ */
+public fun <T> DataFrame<T>.moveTo(
+    newColumnIndex: Int,
+    insideGroup: Boolean,
+    columns: ColumnsSelector<T, *>,
+): DataFrame<T> = move(columns).to(newColumnIndex, insideGroup)
+
+/**
+ * Moves the specified [columns\] 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.
+ *
+ * @include [CommonMoveToDocs]
+ * @include [SelectingColumns.ColumnNames] {@include [SetMoveToOperationArg]}
+ * ### Examples:
+ * ```kotlin
+ * df.moveTo(0, true) { length and age }
+ * df.moveTo(2, false) { cols(1..5) }
+ * ```
+ * @param [newColumnIndex] The index specifying the position in the [DataFrame] 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.
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ */
+public fun <T> DataFrame<T>.moveTo(newColumnIndex: Int, insideGroup: Boolean, vararg columns: String): DataFrame<T> =
+    moveTo(newColumnIndex, insideGroup) { columns.toColumnSet() }
+
 // endregion
 
 // region moveToStart
@@ -264,6 +312,18 @@ public fun <T> DataFrame<T>.moveToLeft(columns: ColumnsSelector<T, *>): DataFram
 @Interpretable("MoveToStart1")
 public fun <T> DataFrame<T>.moveToStart(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toStart()
 
+/**
+ * @include [CommonMoveToStartDocs]
+ * @include [SelectingColumns.Dsl.WithExample] {@include [SetMoveToStartOperationArg]}
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ * @param [insideGroup] If true, selected columns will be moved to the start remaining inside their group,
+ * else they will be moved to the start of the top level.
+ */
+@Refine
+@Interpretable("MoveToStart1")
+public fun <T> DataFrame<T>.moveToStart(insideGroup: Boolean, columns: ColumnsSelector<T, *>): DataFrame<T> =
+    move(columns).toStart(insideGroup)
+
 @Deprecated(MOVE_TO_LEFT, ReplaceWith(MOVE_TO_LEFT_REPLACE), DeprecationLevel.ERROR)
 public fun <T> DataFrame<T>.moveToLeft(vararg columns: String): DataFrame<T> = moveToStart { columns.toColumnSet() }
 
@@ -274,6 +334,16 @@ public fun <T> DataFrame<T>.moveToLeft(vararg columns: String): DataFrame<T> = m
  */
 public fun <T> DataFrame<T>.moveToStart(vararg columns: String): DataFrame<T> = moveToStart { columns.toColumnSet() }
 
+/**
+ * @include [CommonMoveToStartDocs]
+ * @include [SelectingColumns.ColumnNames.WithExample] {@include [SetMoveToStartOperationArg]}
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ * @param [insideGroup] If true, selected columns will be moved to the start remaining inside their group,
+ * else they will be moved to the start of the top level.
+ */
+public fun <T> DataFrame<T>.moveToStart(insideGroup: Boolean, vararg columns: String): DataFrame<T> =
+    moveToStart(insideGroup) { columns.toColumnSet() }
+
 @Deprecated(MOVE_TO_LEFT, ReplaceWith(MOVE_TO_LEFT_REPLACE), DeprecationLevel.ERROR)
 @AccessApiOverload
 public fun <T> DataFrame<T>.moveToLeft(vararg columns: AnyColumnReference): DataFrame<T> =
@@ -339,6 +409,18 @@ public fun <T> DataFrame<T>.moveToRight(columns: ColumnsSelector<T, *>): DataFra
 @Interpretable("MoveToEnd1")
 public fun <T> DataFrame<T>.moveToEnd(columns: ColumnsSelector<T, *>): DataFrame<T> = move(columns).toEnd()
 
+/**
+ * @include [CommonMoveToEndDocs]
+ * @include [SelectingColumns.Dsl.WithExample] {@include [SetMoveToEndOperationArg]}
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ * @param [insideGroup] If true, selected columns will be moved to the end remaining inside their group,
+ * else they will be moved to the end of the top level.
+ */
+@Refine
+@Interpretable("MoveToEnd1")
+public fun <T> DataFrame<T>.moveToEnd(insideGroup: Boolean, columns: ColumnsSelector<T, *>): DataFrame<T> =
+    move(columns).toEnd(insideGroup)
+
 @Deprecated(MOVE_TO_RIGHT, ReplaceWith(MOVE_TO_RIGHT_REPLACE), DeprecationLevel.ERROR)
 public fun <T> DataFrame<T>.moveToRight(vararg columns: String): DataFrame<T> = moveToEnd { columns.toColumnSet() }
 
@@ -349,6 +431,16 @@ public fun <T> DataFrame<T>.moveToRight(vararg columns: String): DataFrame<T> =
  */
 public fun <T> DataFrame<T>.moveToEnd(vararg columns: String): DataFrame<T> = moveToEnd { columns.toColumnSet() }
 
+/**
+ * @include [CommonMoveToEndDocs]
+ * @include [SelectingColumns.ColumnNames.WithExample] {@include [SetMoveToEndOperationArg]}
+ * @param [columns\] The [Columns Selector][ColumnsSelector] used to select the columns of this [DataFrame] to move.
+ * @param [insideGroup] If true, selected columns will be moved to the end remaining inside their group,
+ * else they will be moved to the end of the top level.
+ */
+public fun <T> DataFrame<T>.moveToEnd(insideGroup: Boolean, vararg columns: String): DataFrame<T> =
+    moveToEnd(insideGroup) { columns.toColumnSet() }
+
 @Deprecated(MOVE_TO_RIGHT, ReplaceWith(MOVE_TO_RIGHT_REPLACE), DeprecationLevel.ERROR)
 @AccessApiOverload
 public fun <T> DataFrame<T>.moveToRight(vararg columns: AnyColumnReference): DataFrame<T> =
@@ -537,6 +629,32 @@ public fun <T, C> MoveClause<T, C>.under(
 @Interpretable("MoveTo")
 public fun <T, C> MoveClause<T, C>.to(columnIndex: Int): DataFrame<T> = 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 <T, C> MoveClause<T, C>.to(columnIndex: Int, insideGroup: Boolean): DataFrame<T> =
+    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 +809,28 @@ public fun <T, C> MoveClause<T, C>.toLeft(): DataFrame<T> = to(0)
 @Interpretable("MoveToStart0")
 public fun <T, C> MoveClause<T, C>.toStart(): DataFrame<T> = 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<String>() }.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.
+ */
+@Refine
+@Interpretable("MoveToStart0")
+public fun <T, C> MoveClause<T, C>.toStart(insideGroup: Boolean): DataFrame<T> = to(0, insideGroup)
+
 @Deprecated(TO_RIGHT, ReplaceWith(TO_RIGHT_REPLACE), DeprecationLevel.ERROR)
 public fun <T, C> MoveClause<T, C>.toRight(): DataFrame<T> = to(df.ncol)
 
@@ -712,6 +852,28 @@ public fun <T, C> MoveClause<T, C>.toRight(): DataFrame<T> = to(df.ncol)
 @Interpretable("MoveToEnd0")
 public fun <T, C> MoveClause<T, C>.toEnd(): DataFrame<T> = 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<String>() }.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 <T, C> MoveClause<T, C>.toEnd(insideGroup: Boolean): DataFrame<T> = to(df.ncol, insideGroup)
+
 /**
  * An intermediate class used in the [move] operation.
  *

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/exceptions/ColumnsWithDifferentParentException.kt

@@ -0,0 +1,8 @@
+package org.jetbrains.kotlinx.dataframe.exceptions
+
+public class ColumnsWithDifferentParentException(message: String) :
+    IllegalArgumentException(),
+    DataFrameError {
+
+    override val message: String = message
+}

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

@@ -8,22 +8,29 @@ 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.replace
+import org.jetbrains.kotlinx.dataframe.api.to
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.api.with
 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.exceptions.ColumnsWithDifferentParentException
 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
+import kotlin.collections.first
 
 internal fun <T, C> MoveClause<T, C>.afterOrBefore(column: ColumnSelector<T, *>, isAfter: Boolean): DataFrame<T> {
     val removeResult = df.removeImpl(columns = columns)
@@ -124,3 +131,33 @@ internal fun <T, C> MoveClause<T, C>.moveTo(columnIndex: Int): DataFrame<T> {
             }
     return newColumnList.toDataFrame().cast()
 }
+
+internal fun <T, C> MoveClause<T, C>.moveToImpl(columnIndex: Int, insideGroup: Boolean): DataFrame<T> {
+    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 ColumnsWithDifferentParentException(
+            "Cannot move columns to an index remaining inside group if they have different parent",
+        )
+    }
+
+    // if columns will be moved to top level or columns to move are at top level
+    if (parentOfFirst.isEmpty()) {
+        return moveTo(columnIndex)
+    }
+
+    // replace the level where columns to move are with a new one where columns are moved
+    val columnsToMoveNames = columnsToMove.map { it.name() }
+    return df.replace { parentOfFirst.asColumnGroup() }.with {
+        it.asDataFrame()
+            .move { columnsToMoveNames.toColumnSet() }.to(columnIndex)
+            .asColumnGroup(it.name())
+    }
+}

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<IllegalArgumentException> {
+            grouped.move { "a"["b"] and "b"["c"] }.to(0, true)
+        }.message shouldBe "Cannot move columns to an index remaining inside group if they have different parent"
+    }
 }