@@ -32,6 +32,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ColumnPosition
3232import org.jetbrains.kotlinx.dataframe.impl.columns.tree.getOrPut
3333import org.jetbrains.kotlinx.dataframe.impl.last
3434import org.jetbrains.kotlinx.dataframe.path
35+ import kotlin.collections.first
3536
3637internal fun <T , C > MoveClause <T , C >.afterOrBefore (column : ColumnSelector <T , * >, isAfter : Boolean ): DataFrame <T > {
3738 val removeResult = df.removeImpl(columns = columns)
@@ -154,15 +155,22 @@ internal fun <T, C> MoveClause<T, C>.moveToImpl(columnIndex: Int, insideGroup: B
154155 return moveTo(columnIndex)
155156 }
156157
157- // idea : remove columns to move and brothers ( sons), apply moveTo to sons, insert sons
158+ // logic : remove columns to move and their siblings (from this point, sons), apply them moveTo, reinsert them
158159 val parentPath = df[parentOfFirst].path
159160 val sons = df[parentOfFirst].asColumnGroup()
160- // remove columns to move and their siblings
161- val sonsPaths = sons.columns().map { parentPath + it.path }
162- val intermediateDf = df.removeImpl { sonsPaths .toColumnSet() }
163- // move columns and insert back
161+ // remove sons
162+ val sonsWithFullPaths = sons.columns().map { parentPath + it.path }
163+ val intermediateDf = df.removeImpl { sonsWithFullPaths .toColumnSet() }
164+ // move sons and reinsert them
164165 val columnsToMoveWithReducedPath = columnsToMove.map { it.path.last(it.path.size - parentPath.size).toPath() }
165- val columnsMoved = sons.asDataFrame().move { columnsToMoveWithReducedPath.toColumnSet() }.to(columnIndex).columns()
166- val columnsMovedPaths = columnsMoved.map { ColumnToInsert (parentPath + it.path, it ) }
167- return intermediateDf.df.insertImpl(columnsMovedPaths)
166+ val sonsHaveBeenMoved = sons.asDataFrame().move {
167+ columnsToMoveWithReducedPath.toColumnSet()
168+ }.to(columnIndex).columns()
169+ val sonsToInsert = sonsHaveBeenMoved.map { ColumnToInsert (parentPath + it.path, it) }
170+ val secondIntermediateDf = intermediateDf.df.insertImpl(sonsToInsert)
171+ // nested level is good but order of top level is changed -> need to fix it
172+ val rootOfColumnsToMove = df[listOf (parentPath.first()).toPath()]
173+ val indexOfRootOfColumnsToMove = df.columns().indexOf(rootOfColumnsToMove)
174+ val finalDf = secondIntermediateDf.move { listOf (parentPath.first()).toPath() }.to(indexOfRootOfColumnsToMove)
175+ return finalDf
168176}
0 commit comments