11package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22
3- import org.jetbrains.kotlinx.dataframe.api.rename
43import org.jetbrains.kotlinx.dataframe.api.renameToCamelCase
54import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
65import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
@@ -18,9 +17,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
1817class Rename : AbstractInterpreter <RenameClauseApproximation >() {
1918 private val Arguments .receiver by dataFrame()
2019 private val Arguments .columns: ColumnsResolver by arg()
21- override fun Arguments.interpret (): RenameClauseApproximation {
22- return RenameClauseApproximation (receiver, columns)
23- }
20+
21+ override fun Arguments.interpret (): RenameClauseApproximation = RenameClauseApproximation (receiver, columns)
2422}
2523
2624class RenameClauseApproximation (val schema : PluginDataFrameSchema , val columns : ColumnsResolver )
@@ -33,58 +31,72 @@ class RenameInto : AbstractSchemaModificationInterpreter() {
3331 val columns = receiver.columns.resolve(receiver.schema)
3432 require(columns.size == newNames.size)
3533 var i = 0
36- return receiver.schema.map(columns.mapTo(mutableSetOf ()) { it.path.path }, nextName = { newNames[i].also { i + = 1 } })
34+ return receiver.schema.map(
35+ selected = columns.mapTo(mutableSetOf ()) { it.path.path },
36+ nextName = { newNames[i].also { i + = 1 } },
37+ )
3738 }
3839}
3940
40- internal fun PluginDataFrameSchema.map (selected : ColumnsSet , nextName : () -> String ): PluginDataFrameSchema {
41- return PluginDataFrameSchema (
42- f(columns(), nextName, selected, emptyList())
41+ internal fun PluginDataFrameSchema.map (selected : ColumnsSet , nextName : () -> String ): PluginDataFrameSchema =
42+ PluginDataFrameSchema (
43+ f(columns(), nextName, selected, emptyList()),
4344 )
44- }
4545
46- internal fun f (columns : List <SimpleCol >, nextName : () -> String , selected : ColumnsSet , path : List <String >): List <SimpleCol > {
47- return columns.map {
46+ internal fun f (
47+ columns : List <SimpleCol >,
48+ nextName : () -> String ,
49+ selected : ColumnsSet ,
50+ path : List <String >,
51+ ): List <SimpleCol > =
52+ columns.map {
4853 val fullPath = path + listOf (it.name)
4954 when (it) {
5055 is SimpleColumnGroup -> {
5156 val group = if (fullPath in selected) {
5257 it.rename(nextName())
53- } else {
58+ } else {
5459 it
5560 }
5661 group.map(selected, fullPath, nextName)
5762 }
63+
5864 is SimpleFrameColumn -> {
5965 val frame = if (fullPath in selected) {
6066 it.rename(nextName())
61- } else {
67+ } else {
6268 it
6369 }
6470 frame.map(selected, fullPath, nextName)
6571 }
72+
6673 is SimpleDataColumn -> if (fullPath in selected) {
6774 it.rename(nextName())
6875 } else {
6976 it
7077 }
7178 }
7279 }
73- }
7480
75- internal fun SimpleColumnGroup.map (selected : ColumnsSet , path : List <String >, nextName : () -> String ): SimpleColumnGroup {
76- return SimpleColumnGroup (
77- name,
78- f(columns(), nextName, selected, path)
81+ internal fun SimpleColumnGroup.map (
82+ selected : ColumnsSet ,
83+ path : List <String >,
84+ nextName : () -> String ,
85+ ): SimpleColumnGroup =
86+ SimpleColumnGroup (
87+ name = name,
88+ columns = f(columns(), nextName, selected, path),
7989 )
80- }
8190
82- internal fun SimpleFrameColumn.map (selected : ColumnsSet , path : List <String >, nextName : () -> String ): SimpleFrameColumn {
83- return SimpleFrameColumn (
84- name,
85- f(columns(), nextName, selected, path)
91+ internal fun SimpleFrameColumn.map (
92+ selected : ColumnsSet ,
93+ path : List <String >,
94+ nextName : () -> String ,
95+ ): SimpleFrameColumn =
96+ SimpleFrameColumn (
97+ name = name,
98+ columns = f(columns(), nextName, selected, path),
8699 )
87- }
88100
89101class RenameToCamelCase : AbstractSchemaModificationInterpreter () {
90102 val Arguments .receiver by dataFrame()
0 commit comments