Skip to content

Commit d447799

Browse files
JolanrensenJolanrensen
committed
added map function for TreeNode. improved renameImpl according to reviews and using new map function
1 parent 147ee90 commit d447799

File tree

6 files changed

+86
-38
lines changed
  • core
    • generated-sources/src
      • main/kotlin/org/jetbrains/kotlinx/dataframe/impl
      • test/kotlin/org/jetbrains/kotlinx/dataframe/api
    • src
      • main/kotlin/org/jetbrains/kotlinx/dataframe/impl
      • test/kotlin/org/jetbrains/kotlinx/dataframe/api

6 files changed

+86
-38
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/rename.kt

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@ package org.jetbrains.kotlinx.dataframe.impl.api
22

33
import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.api.RenameClause
5+
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
56
import org.jetbrains.kotlinx.dataframe.api.cast
67
import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
7-
import org.jetbrains.kotlinx.dataframe.api.insert
8-
import org.jetbrains.kotlinx.dataframe.api.under
8+
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
9+
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
910
import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
1011
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
1112
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.allChildrenNotNull
1213
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.collectTree
14+
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.map
1315
import org.jetbrains.kotlinx.dataframe.kind
1416

1517
internal fun <T, C> RenameClause<T, C>.renameImpl(newNames: Array<out String>): DataFrame<T> {
@@ -18,29 +20,40 @@ internal fun <T, C> RenameClause<T, C>.renameImpl(newNames: Array<out String>):
1820
}
1921

2022
internal fun <T, C> RenameClause<T, C>.renameImpl(transform: (ColumnWithPath<C>) -> String): DataFrame<T> {
21-
val selectedColumns = df.getColumnsWithPaths(columns)
23+
// get all selected columns and their paths
24+
val selectedColumnsWithPath = df.getColumnsWithPaths(columns)
25+
.associateBy { it.data }
26+
// gather a tree of all columns where the nodes will be renamed
2227
val tree = df.getColumnsWithPaths { all().rec() }.collectTree()
2328

2429
// perform rename in nodes
2530
tree.allChildrenNotNull().forEach { node ->
26-
val column = selectedColumns.find { it.data == node.data } ?: return@forEach
27-
val newName = transform(column)
28-
node.name = newName
31+
// Check if the current node/column is a selected column and, if so, get its ColumnWithPath
32+
val column = selectedColumnsWithPath[node.data] ?: return@forEach
33+
// Use the found selected ColumnWithPath to query for the new name
34+
val newColumnName = transform(column)
35+
node.name = newColumnName
2936
}
3037

31-
// build up a new DataFrame using the modified names
32-
var newDf = DataFrame.empty(df.rowsCount()).cast<T>()
33-
tree.allChildrenNotNull().forEach { node ->
34-
val path = node.pathFromRoot().dropLast(1)
35-
val col = node.data.rename(node.name)
36-
37-
when (col.kind) {
38+
// use the mapping function to convert the tree to a ColumnGroup/ValueColumn structure
39+
// The result will be a ColumnGroup, since the root node's data is null
40+
val renamedDfAsColumnGroup = tree.map { node, children ->
41+
val col = node.data
42+
when (col?.kind) {
43+
// if the column is a value column or a frame column, rename it using the node's (new) name
3844
ColumnKind.Value, ColumnKind.Frame ->
39-
newDf = newDf.insert(col).under(path)
40-
41-
ColumnKind.Group -> Unit
45+
col.rename(node.name)
46+
47+
// if the column is a group column, create a new column group using the node's (new) name and children
48+
// if the column is null, node is the root, so we'll create a column group as well
49+
ColumnKind.Group, null ->
50+
children
51+
.toDataFrame()
52+
.asColumnGroup(node.name)
4253
}
43-
}
54+
} as ColumnGroup<*>
4455

45-
return newDf
56+
// convert the created ColumnGroup to a DataFrame
57+
val renamedDf = renamedDfAsColumnGroup.columns().toDataFrame()
58+
return renamedDf.cast()
4659
}

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/tree/Utils.kt

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.jetbrains.kotlinx.dataframe.impl.columns.tree
22

33
import org.jetbrains.kotlinx.dataframe.AnyCol
4+
import org.jetbrains.kotlinx.dataframe.DataFrame
45
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
56
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
67
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
@@ -61,6 +62,16 @@ internal fun <T> TreeNode<T>.topmostChildrenExcluding(excludeRoot: TreeNode<*>):
6162
return result
6263
}
6364

65+
/**
66+
* Mapping function for [ReadonlyTreeNodes][ReadonlyTreeNode] (like [TreeNode])
67+
* which can convert the tree-structure (depth-first) to any other tree-type structure (e.g. [DataFrame]).
68+
*/
69+
@Suppress("UNCHECKED_CAST")
70+
internal fun <T : ReadonlyTreeNode<*>, R> T.map(operation: (node: T, children: List<R>) -> R): R {
71+
val children = children.map { (it as T).map(operation) }
72+
return operation(this, children)
73+
}
74+
6475
internal fun <T> TreeNode<T?>.allChildrenNotNull(): List<TreeNode<T>> =
6576
allChildren { it.data != null } as List<TreeNode<T>>
6677

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/rename.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ class RenameToCamelCaseTests {
109109
).first()
110110
)
111111

112-
val df = doublyNestedColumnGroup.renameToCamelCase()//.alsoDebug()
112+
val df = doublyNestedColumnGroup.renameToCamelCase()
113113
df.columnNames() shouldBe listOf("testName")
114114
df["testName"].asColumnGroup().columnNames() shouldBe listOf("anotherName")
115115
df["testName"]["anotherName"].asColumnGroup().columnNames() shouldBe listOf("thirdName")

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

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@ package org.jetbrains.kotlinx.dataframe.impl.api
22

33
import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.api.RenameClause
5+
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
56
import org.jetbrains.kotlinx.dataframe.api.cast
67
import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
7-
import org.jetbrains.kotlinx.dataframe.api.insert
8-
import org.jetbrains.kotlinx.dataframe.api.under
8+
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
9+
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
910
import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
1011
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
1112
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.allChildrenNotNull
1213
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.collectTree
14+
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.map
1315
import org.jetbrains.kotlinx.dataframe.kind
1416

1517
internal fun <T, C> RenameClause<T, C>.renameImpl(newNames: Array<out String>): DataFrame<T> {
@@ -18,29 +20,40 @@ internal fun <T, C> RenameClause<T, C>.renameImpl(newNames: Array<out String>):
1820
}
1921

2022
internal fun <T, C> RenameClause<T, C>.renameImpl(transform: (ColumnWithPath<C>) -> String): DataFrame<T> {
21-
val selectedColumns = df.getColumnsWithPaths(columns)
23+
// get all selected columns and their paths
24+
val selectedColumnsWithPath = df.getColumnsWithPaths(columns)
25+
.associateBy { it.data }
26+
// gather a tree of all columns where the nodes will be renamed
2227
val tree = df.getColumnsWithPaths { all().rec() }.collectTree()
2328

2429
// perform rename in nodes
2530
tree.allChildrenNotNull().forEach { node ->
26-
val column = selectedColumns.find { it.data == node.data } ?: return@forEach
27-
val newName = transform(column)
28-
node.name = newName
31+
// Check if the current node/column is a selected column and, if so, get its ColumnWithPath
32+
val column = selectedColumnsWithPath[node.data] ?: return@forEach
33+
// Use the found selected ColumnWithPath to query for the new name
34+
val newColumnName = transform(column)
35+
node.name = newColumnName
2936
}
3037

31-
// build up a new DataFrame using the modified names
32-
var newDf = DataFrame.empty(df.rowsCount()).cast<T>()
33-
tree.allChildrenNotNull().forEach { node ->
34-
val path = node.pathFromRoot().dropLast(1)
35-
val col = node.data.rename(node.name)
36-
37-
when (col.kind) {
38+
// use the mapping function to convert the tree to a ColumnGroup/ValueColumn structure
39+
// The result will be a ColumnGroup, since the root node's data is null
40+
val renamedDfAsColumnGroup = tree.map { node, children ->
41+
val col = node.data
42+
when (col?.kind) {
43+
// if the column is a value column or a frame column, rename it using the node's (new) name
3844
ColumnKind.Value, ColumnKind.Frame ->
39-
newDf = newDf.insert(col).under(path)
40-
41-
ColumnKind.Group -> Unit
45+
col.rename(node.name)
46+
47+
// if the column is a group column, create a new column group using the node's (new) name and children
48+
// if the column is null, node is the root, so we'll create a column group as well
49+
ColumnKind.Group, null ->
50+
children
51+
.toDataFrame()
52+
.asColumnGroup(node.name)
4253
}
43-
}
54+
} as ColumnGroup<*>
4455

45-
return newDf
56+
// convert the created ColumnGroup to a DataFrame
57+
val renamedDf = renamedDfAsColumnGroup.columns().toDataFrame()
58+
return renamedDf.cast()
4659
}

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/tree/Utils.kt

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.jetbrains.kotlinx.dataframe.impl.columns.tree
22

33
import org.jetbrains.kotlinx.dataframe.AnyCol
4+
import org.jetbrains.kotlinx.dataframe.DataFrame
45
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
56
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
67
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
@@ -61,6 +62,16 @@ internal fun <T> TreeNode<T>.topmostChildrenExcluding(excludeRoot: TreeNode<*>):
6162
return result
6263
}
6364

65+
/**
66+
* Mapping function for [ReadonlyTreeNodes][ReadonlyTreeNode] (like [TreeNode])
67+
* which can convert the tree-structure (depth-first) to any other tree-type structure (e.g. [DataFrame]).
68+
*/
69+
@Suppress("UNCHECKED_CAST")
70+
internal fun <T : ReadonlyTreeNode<*>, R> T.map(operation: (node: T, children: List<R>) -> R): R {
71+
val children = children.map { (it as T).map(operation) }
72+
return operation(this, children)
73+
}
74+
6475
internal fun <T> TreeNode<T?>.allChildrenNotNull(): List<TreeNode<T>> =
6576
allChildren { it.data != null } as List<TreeNode<T>>
6677

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ class RenameToCamelCaseTests {
109109
).first()
110110
)
111111

112-
val df = doublyNestedColumnGroup.renameToCamelCase()//.alsoDebug()
112+
val df = doublyNestedColumnGroup.renameToCamelCase()
113113
df.columnNames() shouldBe listOf("testName")
114114
df["testName"].asColumnGroup().columnNames() shouldBe listOf("anotherName")
115115
df["testName"]["anotherName"].asColumnGroup().columnNames() shouldBe listOf("thirdName")

0 commit comments

Comments
 (0)