From 046cd0c445dab06eb2591fb09f3647b6c19ca028 Mon Sep 17 00:00:00 2001 From: "Anatoly.Nikitin" Date: Fri, 8 Dec 2023 00:47:55 +0300 Subject: [PATCH] Fix performance problem in `rename` implementation. Use column path instead of column content as a key for `associateBy`. --- .../org/jetbrains/kotlinx/dataframe/impl/api/rename.kt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/rename.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/rename.kt index 2e7e8ece2c..b6130abf36 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/rename.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/rename.kt @@ -22,14 +22,14 @@ internal fun RenameClause.renameImpl(newNames: Array): internal fun RenameClause.renameImpl(transform: (ColumnWithPath) -> String): DataFrame { // get all selected columns and their paths val selectedColumnsWithPath = df.getColumnsWithPaths(columns) - .associateBy { it.data } + .associateBy { it.path } // gather a tree of all columns where the nodes will be renamed val tree = df.getColumnsWithPaths { all().rec() }.collectTree() // perform rename in nodes - tree.allChildrenNotNull().forEach { node -> + tree.allChildrenNotNull().map { it to it.pathFromRoot() }.forEach { (node, originalPath) -> // Check if the current node/column is a selected column and, if so, get its ColumnWithPath - val column = selectedColumnsWithPath[node.data] ?: return@forEach + val column = selectedColumnsWithPath[originalPath] ?: return@forEach // Use the found selected ColumnWithPath to query for the new name val newColumnName = transform(column) node.name = newColumnName