Skip to content

Commit d857715

Browse files
koperagenkoperagen
committed
[Compiler plugin] Use simpleColumnOf where the library uses DataColumn.create
1 parent 38a376c commit d857715

File tree

11 files changed

+3690
-49
lines changed

11 files changed

+3690
-49
lines changed

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22

33
package org.jetbrains.kotlinx.dataframe.plugin.impl
44

5+
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
56
import org.jetbrains.kotlin.fir.types.ConeKotlinType
67
import org.jetbrains.kotlin.fir.types.ConeNullability
7-
import org.jetbrains.kotlin.fir.types.classId
88
import org.jetbrains.kotlin.fir.types.isNullable
99
import org.jetbrains.kotlinx.dataframe.impl.api.DataFrameLikeContainer
1010
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumn
@@ -101,17 +101,17 @@ data class SimpleColumnGroup(
101101
}
102102

103103
fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
104-
return if (type.classId == Names.DATA_ROW_CLASS_ID) {
105-
val schema = pluginDataFrameSchema(type)
104+
return if (type.fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID) {
105+
val schema = pluginDataFrameSchema(type.typeArguments[0])
106106
val group = SimpleColumnGroup(name, schema.columns())
107107
val column = if (type.isNullable) {
108108
makeNullable(group)
109109
} else {
110110
group
111111
}
112112
column
113-
} else if (type.classId == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
114-
val schema = pluginDataFrameSchema(type)
113+
} else if (type.fullyExpandedClassId(session) == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
114+
val schema = pluginDataFrameSchema(type.typeArguments[0])
115115
SimpleFrameColumn(name, schema.columns())
116116
} else {
117117
SimpleDataColumn(name, type.wrap())
@@ -121,7 +121,7 @@ fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleC
121121
private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
122122
return when (column) {
123123
is SimpleColumnGroup -> {
124-
SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) })
124+
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
125125
}
126126
is SimpleFrameColumn -> column
127127
is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
77
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
88
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
99
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
10-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1110
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
11+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
1212
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
1313
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
1414

@@ -20,7 +20,7 @@ class Add : AbstractSchemaModificationInterpreter() {
2020
val Arguments.type: TypeApproximation by type(name("expression"))
2121

2222
override fun Arguments.interpret(): PluginDataFrameSchema {
23-
return PluginDataFrameSchema(receiver.columns() + SimpleDataColumn(name, type))
23+
return PluginDataFrameSchema(receiver.columns() + simpleColumnOf(name, type.type))
2424
}
2525
}
2626

@@ -30,7 +30,7 @@ class From : AbstractInterpreter<Unit>() {
3030
val Arguments.type: TypeApproximation by type(name("expression"))
3131

3232
override fun Arguments.interpret() {
33-
dsl.columns += SimpleDataColumn(receiver, type)
33+
dsl.columns += simpleColumnOf(receiver, type.type)
3434
}
3535
}
3636

@@ -40,7 +40,7 @@ class Into : AbstractInterpreter<Unit>() {
4040
val Arguments.name: String by string()
4141

4242
override fun Arguments.interpret() {
43-
dsl.columns += SimpleDataColumn(name, receiver)
43+
dsl.columns += simpleColumnOf(name, receiver.type)
4444
}
4545
}
4646

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

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,20 @@
11
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

3-
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
4-
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
5-
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
6-
import org.jetbrains.kotlin.fir.types.classId
7-
import org.jetbrains.kotlin.fir.types.isNullable
3+
import org.jetbrains.kotlinx.dataframe.api.Infer
84
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
95
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
106
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
117
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
12-
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
13-
import org.jetbrains.kotlinx.dataframe.api.Infer
148
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
9+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
1510
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
16-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1711
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
12+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1813
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
1914
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
2015
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
2116
import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
17+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
2218
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
2319
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
2420
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
@@ -73,17 +69,8 @@ internal fun KotlinTypeFacade.convertImpl(
7369
type: TypeApproximation
7470
): PluginDataFrameSchema {
7571
return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
76-
require(column is SimpleDataColumn) {
77-
"$path must be ${SimpleDataColumn::class}, but was ${column::class}"
78-
}
7972
val unwrappedType = type.type
80-
// TODO: AnyRow
81-
if (unwrappedType is ConeClassLikeType && unwrappedType.classId == Names.DF_CLASS_ID && !unwrappedType.isNullable) {
82-
val f = unwrappedType.typeArguments.single()
83-
SimpleFrameColumn(column.name, pluginDataFrameSchema(f).columns())
84-
} else {
85-
column.changeType(type)
86-
}
73+
simpleColumnOf(column.name, unwrappedType)
8774
}
8875
}
8976

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

Lines changed: 2 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
2121
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
2222
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
2323
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
24-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
2524
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
2625
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
2726
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
2827
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
28+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
2929

3030
class GroupBy(val df: PluginDataFrameSchema, val keys: List<ColumnWithPathApproximation>, val moveToTop: Boolean)
3131

@@ -82,27 +82,8 @@ fun KotlinTypeFacade.aggregate(
8282
)
8383
}
8484

85-
// important to create FrameColumns, nullable DataRows?
8685
val cols = createPluginDataFrameSchema(groupBy.keys, groupBy.moveToTop).columns() + dsl.columns.map {
87-
when (it.type.classId) {
88-
Names.DATA_ROW_CLASS_ID -> {
89-
when (it.type.nullability) {
90-
ConeNullability.NULLABLE -> SimpleDataColumn(
91-
it.name,
92-
TypeApproximation(it.type)
93-
)
94-
ConeNullability.UNKNOWN -> SimpleDataColumn(
95-
it.name,
96-
TypeApproximation(it.type)
97-
)
98-
ConeNullability.NOT_NULL -> {
99-
val typeProjection = it.type.typeArguments[0]
100-
SimpleColumnGroup(it.name, pluginDataFrameSchema(typeProjection).columns())
101-
}
102-
}
103-
}
104-
else -> SimpleDataColumn(it.name, TypeApproximation(it.type))
105-
}
86+
simpleColumnOf(it.name, it.type)
10687
}
10788
PluginDataFrameSchema(cols)
10889
} else {

0 commit comments

Comments
 (0)