Skip to content

Commit dee7e29

Browse files
koperagenkoperagen
committed
Revert insert implementation to only work with DataFrame structure
For compiler plugin, it's easier to convert PluginDataFrameSchema to DataFrame as an adapter for tree operations
1 parent 785dafa commit dee7e29

File tree

8 files changed

+102
-201
lines changed

8 files changed

+102
-201
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
66
import org.jetbrains.kotlinx.dataframe.DataFrame
77
import org.jetbrains.kotlinx.dataframe.RowExpression
88
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
9+
import org.jetbrains.kotlinx.dataframe.annotations.Refine
910
import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
1011
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1112
import org.jetbrains.kotlinx.dataframe.impl.api.insertImpl
@@ -49,19 +50,24 @@ public class InsertClause<T>(internal val df: DataFrame<T>, internal val column:
4950

5051
// region under
5152

53+
@Refine
5254
@Interpretable("Under0")
5355
public fun <T> InsertClause<T>.under(column: ColumnSelector<T, *>): DataFrame<T> = under(df.getColumnPath(column))
5456

57+
@Refine
5558
@Interpretable("Under1")
5659
public fun <T> InsertClause<T>.under(columnPath: ColumnPath): DataFrame<T> =
5760
df.insertImpl(columnPath + column.name, column)
5861

62+
@Refine
5963
@Interpretable("Under2")
6064
public fun <T> InsertClause<T>.under(column: ColumnAccessor<*>): DataFrame<T> = under(column.path())
6165

66+
@Refine
6267
@Interpretable("Under3")
6368
public fun <T> InsertClause<T>.under(column: KProperty<*>): DataFrame<T> = under(column.columnName)
6469

70+
@Refine
6571
@Interpretable("Under4")
6672
public fun <T> InsertClause<T>.under(column: String): DataFrame<T> = under(pathOf(column))
6773

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import org.jetbrains.kotlinx.dataframe.AnyBaseCol
44
import org.jetbrains.kotlinx.dataframe.AnyCol
55
import org.jetbrains.kotlinx.dataframe.AnyRow
66
import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
7-
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumn
87
import org.jetbrains.kotlinx.dataframe.impl.asList
98
import org.jetbrains.kotlinx.dataframe.impl.columnName
109
import org.jetbrains.kotlinx.dataframe.impl.columns.DataColumnInternal
@@ -19,9 +18,7 @@ import kotlin.reflect.KType
1918
*
2019
* @param T type of values contained in column.
2120
*/
22-
public interface BaseColumn<out T> :
23-
ColumnReference<T>,
24-
GenericColumn {
21+
public interface BaseColumn<out T> : ColumnReference<T> {
2522

2623
// region info
2724

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnGroup.kt

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
99
import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
1010
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
1111
import org.jetbrains.kotlinx.dataframe.api.columnGroup
12-
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumnGroup
1312
import kotlin.reflect.KProperty
1413

1514
/**
@@ -28,8 +27,7 @@ import kotlin.reflect.KProperty
2827
@HasSchema(schemaArg = 0)
2928
public interface ColumnGroup<out T> :
3029
BaseColumn<DataRow<T>>,
31-
DataFrame<T>,
32-
GenericColumnGroup<BaseColumn<*>> {
30+
DataFrame<T> {
3331

3432
/**
3533
* Gets the rows at given indices.

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

Lines changed: 23 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,12 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
66
import org.jetbrains.kotlinx.dataframe.DataFrame
77
import org.jetbrains.kotlinx.dataframe.api.cast
88
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
9-
import org.jetbrains.kotlinx.dataframe.columns.BaseColumn
9+
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
1010
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1111
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ReadonlyTreeNode
1212
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.ReferenceData
1313
import org.jetbrains.kotlinx.dataframe.impl.columns.tree.getAncestor
14+
import org.jetbrains.kotlinx.dataframe.impl.columns.withDf
1415

1516
internal 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
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/DataFrameAdapter.kt

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
@file:Suppress("INVISIBLE_REFERENCE")
2+
13
package org.jetbrains.kotlinx.dataframe.plugin.impl
24

35
import org.jetbrains.kotlinx.dataframe.AnyCol
@@ -7,21 +9,34 @@ import org.jetbrains.kotlinx.dataframe.api.cast
79
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
810
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
911
import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
12+
import org.jetbrains.kotlinx.dataframe.impl.columns.ColumnGroupImpl
1013
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
1114

12-
fun PluginDataFrameSchema.asDataFrame(): DataFrame<ConeTypesAdapter> = columns().map()
15+
fun PluginDataFrameSchema.asDataFrame(): DataFrame<ConeTypesAdapter> {
16+
val df = columns().map()
17+
return df
18+
}
1319

1420
fun DataFrame<ConeTypesAdapter>.toPluginDataFrameSchema() = PluginDataFrameSchema(columns().mapBack())
1521

1622
interface ConeTypesAdapter
1723

18-
private fun List<SimpleCol>.map(): DataFrame<ConeTypesAdapter> = map {
19-
when (it) {
20-
is SimpleDataColumn -> DataColumn.create(it.name, listOf(it.type))
21-
is SimpleColumnGroup -> DataColumn.createColumnGroup(it.name, it.columns().map())
22-
is SimpleFrameColumn -> DataColumn.createFrameColumn(it.name, listOf(it.columns().map()))
24+
private fun List<SimpleCol>.map(): DataFrame<ConeTypesAdapter> {
25+
val columns = map {
26+
it.asDataColumn()
2327
}
24-
}.let { dataFrameOf(it).cast() }
28+
return dataFrameOf(columns).cast()
29+
}
30+
31+
@Suppress("INVISIBLE_REFERENCE")
32+
fun SimpleCol.asDataColumn(): DataColumn<*> {
33+
val column = when (this) {
34+
is SimpleDataColumn -> DataColumn.create(this.name, listOf(this.type))
35+
is SimpleColumnGroup -> DataColumn.createColumnGroup(this.name, this.columns().map()) as ColumnGroupImpl<*>
36+
is SimpleFrameColumn -> DataColumn.createFrameColumn(this.name, listOf(this.columns().map()))
37+
}
38+
return column
39+
}
2540

2641
private fun List<AnyCol>.mapBack(): List<SimpleCol> = map {
2742
when (it) {

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,6 @@ import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
66
import org.jetbrains.kotlin.fir.types.ConeKotlinType
77
import org.jetbrains.kotlin.fir.types.ConeNullability
88
import org.jetbrains.kotlin.fir.types.isNullable
9-
import org.jetbrains.kotlinx.dataframe.impl.api.DataFrameLikeContainer
10-
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumn
11-
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumnGroup
129
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
1310
import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
1411
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
@@ -17,11 +14,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
1714

1815
data class PluginDataFrameSchema(
1916
private val columns: List<SimpleCol>
20-
) : DataFrameLikeContainer<SimpleCol> {
17+
) {
2118
companion object {
2219
val EMPTY = PluginDataFrameSchema(emptyList())
2320
}
24-
override fun columns(): List<SimpleCol> {
21+
fun columns(): List<SimpleCol> {
2522
return columns
2623
}
2724

@@ -47,10 +44,10 @@ private fun List<SimpleCol>.asString(indent: String = ""): String {
4744
}
4845
}
4946

50-
sealed interface SimpleCol : GenericColumn {
47+
sealed interface SimpleCol {
5148
val name: String
5249

53-
override fun name(): String {
50+
fun name(): String {
5451
return name
5552
}
5653

@@ -60,7 +57,7 @@ sealed interface SimpleCol : GenericColumn {
6057
data class SimpleDataColumn(
6158
override val name: String,
6259
val type: TypeApproximation
63-
) : GenericColumn, SimpleCol {
60+
) : SimpleCol {
6461

6562
override fun name(): String {
6663
return name
@@ -79,8 +76,8 @@ data class SimpleDataColumn(
7976
data class SimpleFrameColumn(
8077
override val name: String,
8178
private val columns: List<SimpleCol>
82-
) : GenericColumnGroup<SimpleCol>, SimpleCol {
83-
override fun columns(): List<SimpleCol> {
79+
) : SimpleCol {
80+
fun columns(): List<SimpleCol> {
8481
return columns
8582
}
8683

@@ -92,9 +89,9 @@ data class SimpleFrameColumn(
9289
data class SimpleColumnGroup(
9390
override val name: String,
9491
private val columns: List<SimpleCol>
95-
) : GenericColumnGroup<SimpleCol>, SimpleCol {
92+
) : SimpleCol {
9693

97-
override fun columns(): List<SimpleCol> {
94+
fun columns(): List<SimpleCol> {
9895
return columns
9996
}
10097

0 commit comments

Comments
 (0)