Skip to content

Commit 5efb5ff

Browse files
authored
Merge pull request #783 from Kotlin/plugin-column-selection-dsl
Add initial support for CS DSL in the compiler plugin
2 parents f5a06eb + e6b4152 commit 5efb5ff

File tree

25 files changed

+201
-53
lines changed

25 files changed

+201
-53
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
1010
import org.jetbrains.kotlinx.dataframe.DataRow
1111
import org.jetbrains.kotlinx.dataframe.Predicate
1212
import org.jetbrains.kotlinx.dataframe.RowFilter
13+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
1314
import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.BehaviorArg
1415
import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ColumnDoesNotExistArg
1516
import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ExampleArg
@@ -300,6 +301,7 @@ public interface AllColumnsSelectionDsl<out _UNUSED> {
300301
*
301302
* `df.`[select][DataFrame.select]` { `[all][ColumnsSelectionDsl.all]`() }`
302303
*/
304+
@Interpretable("All0")
303305
public fun ColumnsSelectionDsl<*>.all(): TransformableColumnSet<*> =
304306
asSingleColumn().allColumnsInternal()
305307

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter
44
import org.jetbrains.kotlinx.dataframe.DataColumn
55
import org.jetbrains.kotlinx.dataframe.DataFrame
66
import org.jetbrains.kotlinx.dataframe.DataRow
7+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
78
import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar
89
import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnGroupName
910
import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnSetName
@@ -138,6 +139,7 @@ public interface ColsAtAnyDepthColumnsSelectionDsl {
138139
*
139140
* `df.`[select][DataFrame.select]` { `[colsAtAnyDepth][ColumnsSelectionDsl.colsAtAnyDepth]` { !it.`[isColumnGroup][DataColumn.isColumnGroup]` } }`
140141
*/
142+
@Interpretable("ColsAtAnyDepth0")
141143
public fun ColumnsSelectionDsl<*>.colsAtAnyDepth(predicate: ColumnFilter<*> = { true }): ColumnSet<*> =
142144
asSingleColumn().colsAtAnyDepthInternal(predicate)
143145

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter
44
import org.jetbrains.kotlinx.dataframe.DataColumn
55
import org.jetbrains.kotlinx.dataframe.DataFrame
66
import org.jetbrains.kotlinx.dataframe.DataRow
7+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
78
import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar
89
import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnGroupName
910
import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnSetName
@@ -203,6 +204,7 @@ public fun <C> ColumnSet<*>.colsOf(
203204
* @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam]
204205
* @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return]
205206
*/
207+
@Interpretable("ColsOf1")
206208
public inline fun <reified C> ColumnSet<*>.colsOf(
207209
noinline filter: ColumnFilter<C> = { true },
208210
): TransformableColumnSet<C> = colsOf(typeOf<C>(), filter)
@@ -228,6 +230,7 @@ public fun <C> ColumnsSelectionDsl<*>.colsOf(
228230
* @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam]
229231
* @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return]
230232
*/
233+
@Interpretable("ColsOf0")
231234
public inline fun <reified C> ColumnsSelectionDsl<*>.colsOf(
232235
noinline filter: ColumnFilter<C> = { true },
233236
): TransformableColumnSet<C> = asSingleColumn().colsOf(typeOf<C>(), filter)

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
44
import org.jetbrains.kotlinx.dataframe.DataFrame
55
import org.jetbrains.kotlinx.dataframe.DataRow
66
import org.jetbrains.kotlinx.dataframe.Predicate
7+
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
78
import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnGroupName
89
import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnSetName
910
import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.PlainDslName
@@ -111,6 +112,7 @@ public interface FrameColsColumnsSelectionDsl {
111112
*
112113
* `df.`[select][DataFrame.select]` { `[frameCols][ColumnsSelectionDsl.frameCols]` { it.`[name][ColumnReference.name]`.`[startsWith][String.startsWith]`("my") } }`
113114
*/
115+
@Interpretable("FrameCols0")
114116
public fun ColumnSet<*>.frameCols(filter: Predicate<FrameColumn<*>> = { true }): TransformableColumnSet<DataFrame<*>> =
115117
frameColumnsInternal(filter)
116118

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
11
kotlin.code.style=official
2-
kotlinVersion=2.0.20-dev-5379
2+
kotlinVersion=2.0.20-Beta2-78

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@ private class Checker(val cache: FirCache<String, PluginDataFrameSchema, KotlinT
9090
val targetProjection = expression.typeArguments.getOrNull(0) as? FirTypeProjectionWithVariance ?: return
9191
val targetType = targetProjection.typeRef.coneType as? ConeClassLikeType ?: return
9292
val target = pluginDataFrameSchema(targetType)
93-
val sourceColumns = source.flatten()
94-
val targetColumns = target.flatten()
93+
val sourceColumns = source.flatten(includeFrames = true)
94+
val targetColumns = target.flatten(includeFrames = true)
9595
val sourceMap = sourceColumns.associate { it.path.path to it.column }
9696
val missingColumns = mutableListOf<String>()
9797
var valid = true

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ import org.jetbrains.kotlin.fir.extensions.FirFunctionCallRefinementExtension
3838
import org.jetbrains.kotlin.fir.moduleData
3939
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
4040
import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
41-
import org.jetbrains.kotlin.fir.resolve.calls.CallInfo
41+
import org.jetbrains.kotlin.fir.resolve.calls.candidate.CallInfo
4242
import org.jetbrains.kotlin.fir.resolve.defaultType
4343
import org.jetbrains.kotlin.fir.resolve.fqName
4444
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider

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

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,29 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

33
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
44
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
5+
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
6+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
57
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
68

7-
internal class And10 : AbstractInterpreter<List<ColumnWithPathApproximation>>() {
8-
val Arguments.other: List<ColumnWithPathApproximation> by arg()
9-
val Arguments.receiver: List<ColumnWithPathApproximation> by arg()
9+
internal class And10 : AbstractInterpreter<ColumnsResolver>() {
10+
val Arguments.other: ColumnsResolver by arg()
11+
val Arguments.receiver: ColumnsResolver by arg()
1012

11-
override fun Arguments.interpret(): List<ColumnWithPathApproximation> {
12-
return receiver + other
13+
override fun Arguments.interpret(): ColumnsResolver {
14+
return object : ColumnsResolver {
15+
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
16+
return receiver.resolve(df) + other.resolve(df)
17+
}
18+
}
1319
}
1420
}
21+
22+
class SingleColumnApproximation(val col: ColumnWithPathApproximation) : ColumnsResolver {
23+
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
24+
return listOf(col)
25+
}
26+
}
27+
28+
interface ColumnsResolver {
29+
fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation>
30+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type
2020
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
2121

2222
internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {
23-
val Arguments.columns: List<ColumnWithPathApproximation> by arg()
23+
val Arguments.columns: ColumnsResolver by arg()
2424
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
2525
override val Arguments.startingSchema get() = receiver
2626

2727
override fun Arguments.interpret(): ConvertApproximation {
28-
return ConvertApproximation(receiver, columns.map { it.path.path })
28+
return ConvertApproximation(receiver, columns.resolve(receiver).map { it.path.path })
2929
}
3030
}
3131

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
1717

1818
class DropNulls0 : AbstractSchemaModificationInterpreter() {
1919
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
20-
val Arguments.columns: List<ColumnWithPathApproximation> by arg()
20+
val Arguments.columns: ColumnsResolver by arg()
2121

2222
override fun Arguments.interpret(): PluginDataFrameSchema {
23-
return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.mapTo(mutableSetOf()) { it.path.path }, emptyList()))
23+
return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
2424
}
2525
}
2626

0 commit comments

Comments
 (0)