Merge pull request #1104 from Kotlin/rename_fix

rename impl fix

Author: AndreiKingsley

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

@@ -15,8 +15,23 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.tree.map
 import org.jetbrains.kotlinx.dataframe.kind
 
 internal fun <T, C> RenameClause<T, C>.renameImpl(newNames: Array<out String>): DataFrame<T> {
-    var i = 0
-    return renameImpl { newNames[i++] }
+    val selectedColumns = df.getColumnsWithPaths(columns)
+
+    if (selectedColumns.size != newNames.size) {
+        throw IllegalArgumentException(
+            "Selected column count (${selectedColumns.size}) must match new " +
+                "column names count (${newNames.size}).",
+        )
+    }
+
+    // associate old column names with new ones
+    val oldToNew = newNames.mapIndexed { index, newName ->
+        selectedColumns[index].path to newName
+    }.toMap()
+
+    return renameImpl { column ->
+        oldToNew[column.path] ?: throw IllegalArgumentException("Unexpected column: $column")
+    }
 }
 
 internal fun <T, C> RenameClause<T, C>.renameImpl(transform: (ColumnWithPath<C>) -> String): DataFrame<T> {

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

@@ -31,6 +31,20 @@ class RenameTests : ColumnsSelectionDslTests() {
         simpleDf.rename { all() }.into("a_renamed", "b_renamed", "c_renamed") shouldBe renamedDf
     }
 
+    @Test
+    fun `test rename with String to String pairs`() {
+        val renamedDf = dataFrameOf("a_renamed", "b_renamed", "c_renamed")(
+            1, 2, 3,
+            4, 5, 6,
+        )
+
+        simpleDf.rename(
+            "c" to "c_renamed",
+            "a" to "a_renamed",
+            "b" to "b_renamed",
+        ) shouldBe renamedDf
+    }
+
     @Test
     fun `partial grouped rename`() {
         val renamedDf = dataFrameOf("a_renamed", "b", "c")(

plugins/kotlin-dataframe/testData/box/rename.kt

@@ -12,7 +12,7 @@ fun box(): String {
     val df1 = df.rename { nested.d }.into("newName")
     df1.nested.newName
 
-    val df2 = df.rename { nested.d and nested }.into("first", "second")
+    val df2 = df.rename { nested and nested.d }.into("first", "second")
     df2.first.second
     return "OK"
 }