From 2b5fe8bf2cd11eb8a46e559d6aa1456df5bfa905 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Thu, 27 Jun 2024 20:22:19 +0300 Subject: [PATCH 1/3] [Compiler plugin] Make interpreters access columns through ColumnsResolver to further support CS DSL --- .../plugin/impl/api/columnsSelectionDsl.kt | 26 +++++++++++++++---- .../dataframe/plugin/impl/api/convert.kt | 4 +-- .../dataframe/plugin/impl/api/dropNulls.kt | 4 +-- .../dataframe/plugin/impl/api/explode.kt | 4 +-- .../dataframe/plugin/impl/api/group.kt | 6 ++--- .../dataframe/plugin/impl/api/groupBy.kt | 4 +-- .../kotlinx/dataframe/plugin/impl/api/join.kt | 2 +- .../dataframe/plugin/impl/api/joinDsl.kt | 9 +++---- .../dataframe/plugin/impl/api/remove.kt | 4 +-- .../dataframe/plugin/impl/api/rename.kt | 9 ++++--- .../dataframe/plugin/impl/api/select.kt | 24 ++++++++++------- .../dataframe/plugin/impl/api/ungroup.kt | 4 +-- .../kotlinx/dataframe/plugin/interpret.kt | 12 ++++++--- 13 files changed, 69 insertions(+), 43 deletions(-) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/columnsSelectionDsl.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/columnsSelectionDsl.kt index 84c38a855..e76b73f74 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/columnsSelectionDsl.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/columnsSelectionDsl.kt @@ -2,13 +2,29 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments +import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema +import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation -internal class And10 : AbstractInterpreter>() { - val Arguments.other: List by arg() - val Arguments.receiver: List by arg() +internal class And10 : AbstractInterpreter() { + val Arguments.other: ColumnsResolver by arg() + val Arguments.receiver: ColumnsResolver by arg() - override fun Arguments.interpret(): List { - return receiver + other + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return receiver.resolve(df) + other.resolve(df) + } + } } } + +class SingleColumnApproximation(val col: ColumnWithPathApproximation) : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return listOf(col) + } +} + +interface ColumnsResolver { + fun resolve(df: PluginDataFrameSchema): List +} diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt index 4eee297fe..ed6bf251d 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt @@ -20,12 +20,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString internal class Convert0 : AbstractInterpreter() { - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() val Arguments.receiver: PluginDataFrameSchema by dataFrame() override val Arguments.startingSchema get() = receiver override fun Arguments.interpret(): ConvertApproximation { - return ConvertApproximation(receiver, columns.map { it.path.path }) + return ConvertApproximation(receiver, columns.resolve(receiver).map { it.path.path }) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dropNulls.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dropNulls.kt index 5840b6ac9..0977904f7 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dropNulls.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dropNulls.kt @@ -17,10 +17,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame class DropNulls0 : AbstractSchemaModificationInterpreter() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): PluginDataFrameSchema { - return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.mapTo(mutableSetOf()) { it.path.path }, emptyList())) + return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList())) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/explode.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/explode.kt index 58ae1301d..48869e648 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/explode.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/explode.kt @@ -16,12 +16,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame internal class Explode0 : AbstractInterpreter() { val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true)) val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.selector: List? by arg(defaultValue = Present(null)) + val Arguments.selector: ColumnsResolver? by arg(defaultValue = Present(null)) override val Arguments.startingSchema get() = receiver override fun Arguments.interpret(): PluginDataFrameSchema { val columns = selector ?: TODO() - return receiver.explodeImpl(dropEmpty, columns.map { ColumnPathApproximation(it.path.path) }) + return receiver.explodeImpl(dropEmpty, columns.resolve(receiver).map { ColumnPathApproximation(it.path.path) }) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/group.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/group.kt index 565a742c6..df8fb4961 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/group.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/group.kt @@ -12,21 +12,21 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame class Group0 : AbstractInterpreter() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): GroupClauseApproximation { return GroupClauseApproximation(receiver, columns) } } -class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: List) +class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver) class Into0 : AbstractSchemaModificationInterpreter() { val Arguments.receiver: GroupClauseApproximation by arg() val Arguments.column: String by arg() override fun Arguments.interpret(): PluginDataFrameSchema { - val grouped = groupImpl(receiver.df.columns(), receiver.columns.mapTo(mutableSetOf()) { it.path.path }, column) + val grouped = groupImpl(receiver.df.columns(), receiver.columns.resolve(receiver.df).mapTo(mutableSetOf()) { it.path.path }, column) return grouped } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt index 6b61ad373..28b47a037 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt @@ -32,10 +32,10 @@ class GroupBy(val df: PluginDataFrameSchema, val keys: List() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() val Arguments.moveToTop: Boolean by arg(defaultValue = Present(true)) - val Arguments.cols: List by arg() + val Arguments.cols: ColumnsResolver by arg() override fun Arguments.interpret(): GroupBy { - return GroupBy(receiver, cols, moveToTop) + return GroupBy(receiver, cols.resolve(receiver), moveToTop) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/join.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/join.kt index b55e9435e..4f92e4e7d 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/join.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/join.kt @@ -18,7 +18,7 @@ internal class Join0 : AbstractInterpreter() { override fun Arguments.interpret(): PluginDataFrameSchema { val nameGenerator = ColumnNameGenerator() val left = receiver.columns() - val right = removeImpl(other.columns(), setOf(selector.right.path.path)).updatedColumns + val right = removeImpl(other.columns(), setOf(selector.right.resolve(receiver).single().path.path)).updatedColumns val rightColumnGroups = right.filterIsInstance().associateBy { it.name } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinDsl.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinDsl.kt index d18af79c7..2c1bb2a78 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinDsl.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/joinDsl.kt @@ -2,15 +2,14 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments -import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation -internal data class ColumnMatchApproximation(val left: ColumnWithPathApproximation, val right: ColumnWithPathApproximation) +internal data class ColumnMatchApproximation(val left: ColumnsResolver, val right: ColumnsResolver) internal class Match0 : AbstractInterpreter() { - val Arguments.receiver: List by arg() - val Arguments.other: List by arg() + val Arguments.receiver: ColumnsResolver by arg() + val Arguments.other: ColumnsResolver by arg() override fun Arguments.interpret(): ColumnMatchApproximation { - return ColumnMatchApproximation(receiver.single(), other.single()) + return ColumnMatchApproximation(receiver, other) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/remove.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/remove.kt index 7c3486d72..a9a782213 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/remove.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/remove.kt @@ -11,10 +11,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame class Remove0 : AbstractSchemaModificationInterpreter() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): PluginDataFrameSchema { - val removeResult = removeImpl(receiver.columns(), columns.mapTo(mutableSetOf()) { it.path.path }) + val removeResult = removeImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }) return PluginDataFrameSchema(removeResult.updatedColumns) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt index 4c55f8e3b..04aae32e1 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt @@ -14,22 +14,23 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString class Rename : AbstractInterpreter() { private val Arguments.receiver by dataFrame() - private val Arguments.columns: List by arg() + private val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): RenameClauseApproximation { return RenameClauseApproximation(receiver, columns) } } -class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: List) +class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) class RenameInto : AbstractSchemaModificationInterpreter() { val Arguments.receiver: RenameClauseApproximation by arg() val Arguments.newNames: List by varargString() override fun Arguments.interpret(): PluginDataFrameSchema { - require(receiver.columns.size == newNames.size) + val columns = receiver.columns.resolve(receiver.schema) + require(columns.size == newNames.size) var i = 0 - return receiver.schema.map(receiver.columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } }) + return receiver.schema.map(columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } }) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt index ec98ff1ae..38cdfbf2c 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt @@ -14,28 +14,32 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type internal class Select0 : AbstractInterpreter() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): PluginDataFrameSchema { - return PluginDataFrameSchema(columns.map { it.column }) + return PluginDataFrameSchema(columns.resolve(receiver).map { it.column }) } } -internal class Expr0 : AbstractInterpreter>() { +internal class Expr0 : AbstractInterpreter() { val Arguments.name: String by arg(defaultValue = Present("untitled")) val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls)) val Arguments.expression: TypeApproximation by type() - override fun Arguments.interpret(): List { - return listOf(ColumnWithPathApproximation(ColumnPathApproximation(listOf(name)), SimpleDataColumn(name, expression))) + override fun Arguments.interpret(): ColumnsResolver { + return SingleColumnApproximation(ColumnWithPathApproximation(ColumnPathApproximation(listOf(name)), SimpleDataColumn(name, expression))) } } -internal class And0 : AbstractInterpreter>() { - val Arguments.receiver: List by arg() - val Arguments.other: List by arg() +internal class And0 : AbstractInterpreter() { + val Arguments.receiver: ColumnsResolver by arg() + val Arguments.other: ColumnsResolver by arg() - override fun Arguments.interpret(): List { - return receiver + other + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return receiver.resolve(df) + other.resolve(df) + } + } } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ungroup.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ungroup.kt index 3d44d3e3a..b14c31a66 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ungroup.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/ungroup.kt @@ -11,10 +11,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame class Ungroup0 : AbstractSchemaModificationInterpreter() { val Arguments.receiver: PluginDataFrameSchema by dataFrame() - val Arguments.columns: List by arg() + val Arguments.columns: ColumnsResolver by arg() override fun Arguments.interpret(): PluginDataFrameSchema { - val ungrouped = ungroupImpl(receiver.columns(), columns.mapTo(mutableSetOf()) { it.path.path }, emptyList()) + val ungrouped = ungroupImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()) return PluginDataFrameSchema(ungrouped) } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt index 15181f1ac..de830972f 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt @@ -76,6 +76,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnsResolver +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.SingleColumnApproximation fun KotlinTypeFacade.interpret( functionCall: FirFunctionCall, @@ -331,7 +333,7 @@ fun KotlinTypeFacade.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType) return PluginDataFrameSchema(columns) } -private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAccessExpression): List { +private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAccessExpression): ColumnsResolver { return result.resolvedType.let { val column = when (it.classId) { Names.DATA_COLUMN_CLASS_ID -> { @@ -346,9 +348,13 @@ private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAcc val path = f(result) SimpleColumnGroup(path, pluginDataFrameSchema(arg).columns()) } - else -> return emptyList() + else -> return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return emptyList() + } + } } - listOf( + SingleColumnApproximation( ColumnWithPathApproximation( path = ColumnPathApproximation(path(result)), column From 48247e11c654bdc9a771f0bbd377e57619dd2bc2 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Tue, 9 Jul 2024 14:03:35 +0300 Subject: [PATCH 2/3] [Compiler plugin] Support colsOf, colsAtAnyDepth, frameCols --- .../jetbrains/kotlinx/dataframe/api/all.kt | 2 + .../kotlinx/dataframe/api/colsAtAnyDepth.kt | 2 + .../jetbrains/kotlinx/dataframe/api/colsOf.kt | 3 + .../kotlinx/dataframe/api/frameCols.kt | 2 + .../ExpressionAnalysisAdditionalChecker.kt | 4 +- .../plugin/impl/api/flattenPluginSchema.kt | 10 +-- .../dataframe/plugin/impl/api/select.kt | 82 ++++++++++++++++++- .../dataframe/plugin/loadInterpreter.kt | 10 +++ .../testData/box/selectionDsl.kt | 17 ++++ ...DataFrameBlackBoxCodegenTestGenerated.java | 6 ++ 10 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 plugins/kotlin-dataframe/testData/box/selectionDsl.kt diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt index 55e5d0a57..46f1f718c 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt @@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow import org.jetbrains.kotlinx.dataframe.Predicate import org.jetbrains.kotlinx.dataframe.RowFilter +import org.jetbrains.kotlinx.dataframe.annotations.Interpretable import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.BehaviorArg import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ColumnDoesNotExistArg import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ExampleArg @@ -300,6 +301,7 @@ public interface AllColumnsSelectionDsl { * * `df.`[select][DataFrame.select]` { `[all][ColumnsSelectionDsl.all]`() }` */ + @Interpretable("All0") public fun ColumnsSelectionDsl<*>.all(): TransformableColumnSet<*> = asSingleColumn().allColumnsInternal() diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt index c80ca0d85..f977cb5e9 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter import org.jetbrains.kotlinx.dataframe.DataColumn import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow +import org.jetbrains.kotlinx.dataframe.annotations.Interpretable import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnGroupName import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnSetName @@ -138,6 +139,7 @@ public interface ColsAtAnyDepthColumnsSelectionDsl { * * `df.`[select][DataFrame.select]` { `[colsAtAnyDepth][ColumnsSelectionDsl.colsAtAnyDepth]` { !it.`[isColumnGroup][DataColumn.isColumnGroup]` } }` */ + @Interpretable("ColsAtAnyDepth0") public fun ColumnsSelectionDsl<*>.colsAtAnyDepth(predicate: ColumnFilter<*> = { true }): ColumnSet<*> = asSingleColumn().colsAtAnyDepthInternal(predicate) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt index 8d1ad745b..90122cea0 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter import org.jetbrains.kotlinx.dataframe.DataColumn import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow +import org.jetbrains.kotlinx.dataframe.annotations.Interpretable import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnGroupName import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnSetName @@ -203,6 +204,7 @@ public fun ColumnSet<*>.colsOf( * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam] * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return] */ +@Interpretable("ColsOf1") public inline fun ColumnSet<*>.colsOf( noinline filter: ColumnFilter = { true }, ): TransformableColumnSet = colsOf(typeOf(), filter) @@ -228,6 +230,7 @@ public fun ColumnsSelectionDsl<*>.colsOf( * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam] * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return] */ +@Interpretable("ColsOf0") public inline fun ColumnsSelectionDsl<*>.colsOf( noinline filter: ColumnFilter = { true }, ): TransformableColumnSet = asSingleColumn().colsOf(typeOf(), filter) diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt index faa74e5d3..935bb3d68 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt @@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow import org.jetbrains.kotlinx.dataframe.Predicate +import org.jetbrains.kotlinx.dataframe.annotations.Interpretable import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnGroupName import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnSetName import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.PlainDslName @@ -111,6 +112,7 @@ public interface FrameColsColumnsSelectionDsl { * * `df.`[select][DataFrame.select]` { `[frameCols][ColumnsSelectionDsl.frameCols]` { it.`[name][ColumnReference.name]`.`[startsWith][String.startsWith]`("my") } }` */ + @Interpretable("FrameCols0") public fun ColumnSet<*>.frameCols(filter: Predicate> = { true }): TransformableColumnSet> = frameColumnsInternal(filter) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt index 7b4707061..5d985eacf 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt @@ -90,8 +90,8 @@ private class Checker(val cache: FirCache() var valid = true diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt index 8290b5b7a..7d9943bef 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt @@ -8,24 +8,24 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation -fun PluginDataFrameSchema.flatten(): List { +fun PluginDataFrameSchema.flatten(includeFrames: Boolean): List { if (columns().isEmpty()) return emptyList() val columns = mutableListOf() - flattenImpl(columns(), emptyList(), columns) + flattenImpl(columns(), emptyList(), columns, includeFrames) return columns } -fun flattenImpl(columns: List, path: List, flatList: MutableList) { +fun flattenImpl(columns: List, path: List, flatList: MutableList, includeFrames: Boolean) { columns.forEach { column -> val fullPath = path + listOf(column.name) when (column) { is SimpleColumnGroup -> { flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column)) - flattenImpl(column.columns(), fullPath, flatList) + flattenImpl(column.columns(), fullPath, flatList, includeFrames) } is SimpleFrameColumn -> { flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column)) - flattenImpl(column.columns(), fullPath, flatList) + flattenImpl(column.columns(), fullPath, flatList, includeFrames) } is SimpleDataColumn -> { flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column)) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt index 38cdfbf2c..c21f79e12 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt @@ -1,11 +1,14 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api +import org.jetbrains.kotlin.fir.types.ConeKotlinType +import org.jetbrains.kotlin.fir.types.isSubtypeOf import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments import org.jetbrains.kotlinx.dataframe.plugin.impl.Present import org.jetbrains.kotlinx.dataframe.api.Infer import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn +import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame @@ -27,7 +30,12 @@ internal class Expr0 : AbstractInterpreter() { val Arguments.expression: TypeApproximation by type() override fun Arguments.interpret(): ColumnsResolver { - return SingleColumnApproximation(ColumnWithPathApproximation(ColumnPathApproximation(listOf(name)), SimpleDataColumn(name, expression))) + return SingleColumnApproximation( + ColumnWithPathApproximation( + ColumnPathApproximation(listOf(name)), + SimpleDataColumn(name, expression) + ) + ) } } @@ -43,3 +51,75 @@ internal class And0 : AbstractInterpreter() { } } } + +internal class All0 : AbstractInterpreter() { + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return df.columns().map { + val path = ColumnPathApproximation(listOf(it.name)) + ColumnWithPathApproximation(path, it) + } + } + } + } +} + +internal class ColsOf0 : AbstractInterpreter() { + val Arguments.typeArg0: TypeApproximation by arg() + + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + val cols = df.columns().map { + val path = ColumnPathApproximation(listOf(it.name)) + ColumnWithPathApproximation(path, it) + } + return colsOf(cols, typeArg0.type) + } + } + } + +} + +private fun Arguments.colsOf(cols: List, type: ConeKotlinType) = + cols + .filter { + val column = it.column + column is SimpleDataColumn && column.type.type.isSubtypeOf(type, session) + } + +internal class ColsAtAnyDepth0 : AbstractInterpreter() { + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return df.flatten(includeFrames = false) + } + } + } +} + +internal class ColsOf1 : AbstractInterpreter() { + val Arguments.receiver: ColumnsResolver by arg() + val Arguments.typeArg0: TypeApproximation by arg() + + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return colsOf(receiver.resolve(df), typeArg0.type) + } + } + } +} + +internal class FrameCols0 : AbstractInterpreter() { + val Arguments.receiver: ColumnsResolver by arg() + + override fun Arguments.interpret(): ColumnsResolver { + return object : ColumnsResolver { + override fun resolve(df: PluginDataFrameSchema): List { + return receiver.resolve(df).filter { it.column is SimpleFrameColumn } + } + } + } +} diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt index 9ab8ed1f1..1d49a67f9 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt @@ -67,6 +67,11 @@ import org.jetbrains.kotlin.fir.types.classId import org.jetbrains.kotlin.fir.types.coneType import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.name.Name +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? { @@ -161,6 +166,11 @@ internal inline fun String.load(): T { "ReadDelimStr" -> ReadDelimStr() "GroupByToDataFrame" -> GroupByToDataFrame() "ToDataFrameFrom0" -> ToDataFrameFrom() + "All0" -> All0() + "ColsOf0" -> ColsOf0() + "ColsOf1" -> ColsOf1() + "ColsAtAnyDepth0" -> ColsAtAnyDepth0() + "FrameCols0" -> FrameCols0() else -> error("$this") } as T } diff --git a/plugins/kotlin-dataframe/testData/box/selectionDsl.kt b/plugins/kotlin-dataframe/testData/box/selectionDsl.kt new file mode 100644 index 000000000..778aa4b1b --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/selectionDsl.kt @@ -0,0 +1,17 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.annotations.* +import org.jetbrains.kotlinx.dataframe.api.* +import org.jetbrains.kotlinx.dataframe.io.* + +data class Nested(val d: List) + +data class Record(val a: String, val b: Int, val nested: List) + +fun box(): String { + val df = listOf(Record("112", 42, listOf(Nested(listOf(3.0))))).toDataFrame(maxDepth = 1) + + df.group { nested }.into("group").convert { colsAtAnyDepth().frameCols() }.with { 1 }.compareSchemas() + df.group { b }.into("group").convert { colsAtAnyDepth().colsOf() }.with { "" }.compareSchemas() + df.convert { all().frameCols() }.with { 1 }.compareSchemas() + return "OK" +} diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java index dadc68e54..299796e33 100644 --- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java +++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java @@ -310,6 +310,12 @@ public void testSelectThis() { runTest("testData/box/selectThis.kt"); } + @Test + @TestMetadata("selectionDsl.kt") + public void testSelectionDsl() { + runTest("testData/box/selectionDsl.kt"); + } + @Test @TestMetadata("toDataFrame.kt") public void testToDataFrame() { From e6b41528055715e09dafb4ec3d2b1de8bf524811 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Tue, 16 Jul 2024 15:09:30 +0300 Subject: [PATCH 3/3] Update compiler plugin to 2.0.20-Beta2-78 --- plugins/kotlin-dataframe/gradle.properties | 2 +- .../dataframe/plugin/extensions/FunctionCallTransformer.kt | 2 +- .../testData/diagnostics/selectDuringTyping.fir.txt | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/plugins/kotlin-dataframe/gradle.properties b/plugins/kotlin-dataframe/gradle.properties index cfe50e0df..c97a3b48c 100644 --- a/plugins/kotlin-dataframe/gradle.properties +++ b/plugins/kotlin-dataframe/gradle.properties @@ -1,2 +1,2 @@ kotlin.code.style=official -kotlinVersion=2.0.20-dev-5379 +kotlinVersion=2.0.20-Beta2-78 diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt index d3e077858..8f778ed12 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt @@ -38,7 +38,7 @@ import org.jetbrains.kotlin.fir.extensions.FirFunctionCallRefinementExtension import org.jetbrains.kotlin.fir.moduleData import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference -import org.jetbrains.kotlin.fir.resolve.calls.CallInfo +import org.jetbrains.kotlin.fir.resolve.calls.candidate.CallInfo import org.jetbrains.kotlin.fir.resolve.defaultType import org.jetbrains.kotlin.fir.resolve.fqName import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider diff --git a/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt b/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt index c296fc526..4479d8415 100644 --- a/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt +++ b/plugins/kotlin-dataframe/testData/diagnostics/selectDuringTyping.fir.txt @@ -35,7 +35,7 @@ FILE: selectDuringTyping.kt } - ^ R|/it|.R|org/jetbrains/kotlinx/dataframe/api/explode|( = explode@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl|.(it: R|@R|kotlin/ParameterName|(name = String(it)) org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl|): R|org/jetbrains/kotlinx/dataframe/columns/ColumnsResolver<*>| { + ^ R|/it|.R|org/jetbrains/kotlinx/dataframe/api/explode|( = explode@fun R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl|.(it: R|org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl|): R|org/jetbrains/kotlinx/dataframe/columns/ColumnsResolver<*>| { ^ # } )