@@ -6,11 +6,12 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
66import org.jetbrains.kotlinx.dataframe.DataFrame
77import org.jetbrains.kotlinx.dataframe.api.cast
88import org.jetbrains.kotlinx.dataframe.api.toDataFrame
9- import org.jetbrains.kotlinx.dataframe.columns.BaseColumn
9+ import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
1010import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1111import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ReadonlyTreeNode
1212import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ReferenceData
1313import org.jetbrains.kotlinx.dataframe.impl.columns.tree.getAncestor
14+ import org.jetbrains.kotlinx.dataframe.impl.columns.withDf
1415
1516internal data class ColumnToInsert (
1617 val insertionPath : ColumnPath ,
@@ -36,125 +37,31 @@ internal fun <T> insertImpl(
3637 columns : List <ColumnToInsert >,
3738 treeNode : ReadonlyTreeNode <ReferenceData >? ,
3839 depth : Int ,
39- ): DataFrame <T > =
40- if (columns.isEmpty()) {
41- df ? : DataFrame .empty().cast()
42- } else {
43- insertImplDataFrame(df, columns, treeNode, depth)
44- }
45-
46- internal fun <T > insertImplDataFrame (
47- df : DataFrame <T >? ,
48- columns : List <ColumnToInsert >,
49- treeNode : ReadonlyTreeNode <ReferenceData >? ,
50- depth : Int ,
5140): DataFrame <T > {
52- class DfAdapter <T >(val df : DataFrame <T >) : DataFrameLikeContainer<BaseColumn<*>> {
53- override fun columns (): List <DataColumn <* >> = this .df.columns()
54- }
55-
56- return if (columns.isEmpty()) {
57- df ? : DataFrame .empty().cast()
58- } else {
59- insertImplGenericContainer(
60- df?.let { DfAdapter (it) },
61- columns.map { GenericColumnsToInsert (it.insertionPath, it.column, it.referenceNode) },
62- treeNode,
63- depth,
64- factory = { DfAdapter (it.toDataFrame().cast()) },
65- empty = DfAdapter (DataFrame .Empty .cast()),
66- rename = { rename(it) },
67- createColumnGroup = { name, columns ->
68- DataColumn .createColumnGroup(name, columns.toDataFrame())
69- },
70- ).df
71- }
72- }
73-
74- internal interface DataFrameLikeContainer <T : GenericColumn > {
75- fun columns (): List <T >
76- }
77-
78- internal fun <T : DataFrameLikeContainer <Column >, Column : GenericColumn , ColumnGroup > insertImplGenericContainer (
79- df : T ? ,
80- columns : List <GenericColumnsToInsert <Column >>,
81- treeNode : ReadonlyTreeNode <ReferenceData >? ,
82- depth : Int ,
83- factory : (List <Column >) -> T ,
84- empty : T ,
85- rename : Column .(String ) -> Column ,
86- createColumnGroup : (String , List <Column >) -> Column ,
87- ): T where ColumnGroup : GenericColumnGroup <Column > {
88- if (columns.isEmpty()) return df ? : empty
89-
90- val res: List <Column > = insertImplGenericTree(
91- columns = columns,
92- treeNode = treeNode,
93- depth = depth,
94- existingColumns = df?.columns(),
95- rename = rename,
96- createColumnGroup = createColumnGroup,
97- )
98- return factory(res)
99- }
41+ if (columns.isEmpty()) return df ? : DataFrame .empty().cast()
10042
101- public interface GenericColumn {
102- public fun name (): String
103- }
104-
105- public interface GenericColumnGroup <Column : GenericColumn > : GenericColumn {
106- public fun columns (): List <Column >
107- }
108-
109- internal data class GenericColumnsToInsert <Column : GenericColumn >(
110- val insertionPath : ColumnPath ,
111- val column : Column ,
112- val referenceNode : ReadonlyTreeNode <ReferenceData >? = null ,
113- )
114-
115- internal fun <Column : GenericColumn , ColumnGroup : GenericColumnGroup <Column >> insertImplGenericTree (
116- columns : List <GenericColumnsToInsert <Column >>,
117- treeNode : ReadonlyTreeNode <ReferenceData >? ,
118- depth : Int ,
119- existingColumns : List <Column >? ,
120- rename : Column .(String ) -> Column ,
121- createColumnGroup : (String , List <Column >) -> Column ,
122- ): List <Column > {
12343 val childDepth = depth + 1
12444
12545 val columnsMap = columns.groupBy { it.insertionPath[depth] }.toMutableMap() // map: columnName -> columnsToAdd
12646
127- val newColumns = mutableListOf<Column >()
47+ val newColumns = mutableListOf<AnyBaseCol >()
12848
12949 // insert new columns under existing
130- existingColumns ?.forEach {
50+ df?.columns() ?.forEach {
13151 val subTree = columnsMap[it.name()]
13252 if (subTree != null ) {
13353 // assert that new columns go directly under current column so they have longer paths
13454 val invalidPath = subTree.firstOrNull { it.insertionPath.size == childDepth }
13555 check(invalidPath == null ) {
136- " Can not insert column '${
137- invalidPath!! .insertionPath.joinToString(" ." )
138- } ' because column with this path already exists in DataFrame"
56+ val text = invalidPath!! .insertionPath.joinToString(" ." )
57+ " Can not insert column `$text ` because column with this path already exists in DataFrame"
13958 }
140- val group = it as ? ColumnGroup
141- check(
142- group != null ,
143- ) { " Can not insert columns under a column '${it.name()} ', because it is not a column group" }
144- val column = if (subTree.isEmpty()) {
145- group as Column
146- } else {
147- val res = insertImplGenericTree(
148- columns = subTree,
149- treeNode = treeNode?.get(it.name()),
150- depth = childDepth,
151- existingColumns = group.columns(),
152- rename = rename,
153- createColumnGroup = createColumnGroup,
154- )
155- createColumnGroup(group.name(), res)
59+ val group = it as ? ColumnGroup <* >
60+ check(group != null ) {
61+ " Can not insert columns under a column '${it.name()} ', because it is not a column group"
15662 }
157- val newCol = column
63+ val newDf = insertImpl(group, subTree, treeNode?.get(it.name()), childDepth)
64+ val newCol = group.withDf(newDf)
15865 newColumns.add(newCol)
15966 columnsMap.remove(it.name())
16067 } else {
@@ -215,39 +122,21 @@ internal fun <Column : GenericColumn, ColumnGroup : GenericColumnGroup<Column>>
215122 check(columns.count { it.insertionPath.size == childDepth } == 1 ) {
216123 " Can not insert more than one column into the path ${nodeToInsert.insertionPath} "
217124 }
218- column as ColumnGroup
219- val columns1 = columns.filter { it.insertionPath.size > childDepth }
220- val newDf = if (columns1.isEmpty()) {
221- listOf (column)
222- } else {
223- insertImplGenericTree(
224- columns = columns1,
225- treeNode = treeNode?.get(name),
226- depth = childDepth,
227- existingColumns = column.columns(),
228- rename = rename,
229- createColumnGroup = createColumnGroup,
230- )
231- }
232- createColumnGroup(name, newDf)
125+ check(column is ColumnGroup <* >)
126+ val newDf = insertImpl(
127+ column,
128+ columns.filter { it.insertionPath.size > childDepth },
129+ treeNode?.get(name),
130+ childDepth,
131+ )
132+ column.withDf(newDf)
233133 } else {
234134 column.rename(name)
235135 }
236136 } else {
237137 val newDf =
238- if (columns.isEmpty()) {
239- emptyList()
240- } else {
241- insertImplGenericTree(
242- columns = columns,
243- treeNode = treeNode?.get(name),
244- depth = childDepth,
245- existingColumns = emptyList(),
246- rename = rename,
247- createColumnGroup = createColumnGroup,
248- )
249- }
250- createColumnGroup(name, newDf) // new node needs to be created
138+ insertImpl<Unit >(null , columns, treeNode?.get(name), childDepth)
139+ DataColumn .createColumnGroup(name, newDf) // new node needs to be created
251140 }
252141 if (insertionIndex == Int .MAX_VALUE ) {
253142 newColumns.add(newCol)
@@ -257,5 +146,5 @@ internal fun <Column : GenericColumn, ColumnGroup : GenericColumnGroup<Column>>
257146 }
258147 }
259148
260- return newColumns
149+ return newColumns.toDataFrame().cast()
261150}
0 commit comments