[Compiler plugin] Make interpreters access columns through ColumnsResolver to further support CS DSL

koperagen · koperagen · commit 2b5fe8bf2cd1 · 2024-07-09T14:03:47.000+03:00
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
@@ -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<List<ColumnWithPathApproximation>>() {
-    val Arguments.other: List<ColumnWithPathApproximation> by arg()
-    val Arguments.receiver: List<ColumnWithPathApproximation> by arg()
+internal class And10 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.other: ColumnsResolver by arg()
+    val Arguments.receiver: ColumnsResolver by arg()
 
-    override fun Arguments.interpret(): List<ColumnWithPathApproximation> {
-        return receiver + other
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return receiver.resolve(df) + other.resolve(df)
+            }
+        }
     }
 }
+
+class SingleColumnApproximation(val col: ColumnWithPathApproximation) : ColumnsResolver {
+    override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+        return listOf(col)
+    }
+}
+
+interface ColumnsResolver {
+    fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation>
+}
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
@@ -20,12 +20,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
 
 internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {
-    val Arguments.columns: List<ColumnWithPathApproximation> 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
@@ -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<ColumnWithPathApproximation> 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
@@ -16,12 +16,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))
     val Arguments.receiver: PluginDataFrameSchema by dataFrame()
-    val Arguments.selector: List<ColumnWithPathApproximation>? 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
@@ -12,21 +12,21 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 
 class Group0 : AbstractInterpreter<GroupClauseApproximation>() {
     val Arguments.receiver: PluginDataFrameSchema by dataFrame()
-    val Arguments.columns: List<ColumnWithPathApproximation> 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<ColumnWithPathApproximation>)
+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
@@ -32,10 +32,10 @@ class GroupBy(val df: PluginDataFrameSchema, val keys: List<ColumnWithPathApprox
 class DataFrameGroupBy : AbstractInterpreter<GroupBy>() {
     val Arguments.receiver: PluginDataFrameSchema by dataFrame()
     val Arguments.moveToTop: Boolean by arg(defaultValue = Present(true))
-    val Arguments.cols: List<ColumnWithPathApproximation> 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
@@ -18,7 +18,7 @@ internal class Join0 : AbstractInterpreter<PluginDataFrameSchema>() {
     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<SimpleColumnGroup>().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
@@ -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<ColumnMatchApproximation>() {
-    val Arguments.receiver: List<ColumnWithPathApproximation> by arg()
-    val Arguments.other: List<ColumnWithPathApproximation> 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
@@ -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<ColumnWithPathApproximation> 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
@@ -14,22 +14,23 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
 
 class Rename : AbstractInterpreter<RenameClauseApproximation>() {
     private val Arguments.receiver by dataFrame()
-    private val Arguments.columns: List<ColumnWithPathApproximation> 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<ColumnWithPathApproximation>)
+class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver)
 
 class RenameInto : AbstractSchemaModificationInterpreter() {
     val Arguments.receiver: RenameClauseApproximation by arg()
     val Arguments.newNames: List<String> 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
@@ -14,28 +14,32 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 
 internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.receiver: PluginDataFrameSchema by dataFrame()
-    val Arguments.columns: List<ColumnWithPathApproximation> 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<List<ColumnWithPathApproximation>>() {
+internal class Expr0 : AbstractInterpreter<ColumnsResolver>() {
     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<ColumnWithPathApproximation> {
-        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<List<ColumnWithPathApproximation>>() {
-    val Arguments.receiver: List<ColumnWithPathApproximation> by arg()
-    val Arguments.other: List<ColumnWithPathApproximation> by arg()
+internal class And0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.other: ColumnsResolver by arg()
 
-    override fun Arguments.interpret(): List<ColumnWithPathApproximation> {
-        return receiver + other
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                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
@@ -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<ColumnWithPathApproximation> 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
@@ -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 <T> KotlinTypeFacade.interpret(
     functionCall: FirFunctionCall,
@@ -331,7 +333,7 @@ fun KotlinTypeFacade.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType)
     return PluginDataFrameSchema(columns)
 }
 
-private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAccessExpression): List<ColumnWithPathApproximation> {
+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<ColumnWithPathApproximation> {
+                    return emptyList()
+                }
+            }
         }
-        listOf(
+        SingleColumnApproximation(
             ColumnWithPathApproximation(
                 path = ColumnPathApproximation(path(result)),
                 column

Original file line number	Diff line number	Diff line change
`@@ -20,12 +20,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.type`
`20`	`20`	`import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString`
`21`	`21`
`22`	`22`	`internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {`
`23`		`- val Arguments.columns: List<ColumnWithPathApproximation> by arg()`
	`23`	`+ val Arguments.columns: ColumnsResolver by arg()`
`24`	`24`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`25`	`25`	`override val Arguments.startingSchema get() = receiver`
`26`	`26`
`27`	`27`	`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 })`
`29`	`29`	`}`
`30`	`30`	`}`
`31`	`31`
Original file line number	Diff line number	Diff line change
`@@ -17,10 +17,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame`
`17`	`17`
`18`	`18`	`class DropNulls0 : AbstractSchemaModificationInterpreter() {`
`19`	`19`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`20`		`- val Arguments.columns: List<ColumnWithPathApproximation> by arg()`
	`20`	`+ val Arguments.columns: ColumnsResolver by arg()`
`21`	`21`
`22`	`22`	`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()))`
`24`	`24`	`}`
`25`	`25`	`}`
`26`	`26`
Original file line number	Diff line number	Diff line change
`@@ -16,12 +16,12 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame`
`16`	`16`	`internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {`
`17`	`17`	`val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))`
`18`	`18`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`19`		`- val Arguments.selector: List<ColumnWithPathApproximation>? by arg(defaultValue = Present(null))`
	`19`	`+ val Arguments.selector: ColumnsResolver? by arg(defaultValue = Present(null))`
`20`	`20`	`override val Arguments.startingSchema get() = receiver`
`21`	`21`
`22`	`22`	`override fun Arguments.interpret(): PluginDataFrameSchema {`
`23`	`23`	`val columns = selector ?: TODO()`
`24`		`- return receiver.explodeImpl(dropEmpty, columns.map { ColumnPathApproximation(it.path.path) })`
	`24`	`+ return receiver.explodeImpl(dropEmpty, columns.resolve(receiver).map { ColumnPathApproximation(it.path.path) })`
`25`	`25`	`}`
`26`	`26`	`}`
`27`	`27`
Original file line number	Diff line number	Diff line change
`@@ -12,21 +12,21 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame`
`12`	`12`
`13`	`13`	`class Group0 : AbstractInterpreter<GroupClauseApproximation>() {`
`14`	`14`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`15`		`- val Arguments.columns: List<ColumnWithPathApproximation> by arg()`
	`15`	`+ val Arguments.columns: ColumnsResolver by arg()`
`16`	`16`
`17`	`17`	`override fun Arguments.interpret(): GroupClauseApproximation {`
`18`	`18`	`return GroupClauseApproximation(receiver, columns)`
`19`	`19`	`}`
`20`	`20`	`}`
`21`	`21`
`22`		`-class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: List<ColumnWithPathApproximation>)`
	`22`	`+class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)`
`23`	`23`
`24`	`24`	`class Into0 : AbstractSchemaModificationInterpreter() {`
`25`	`25`	`val Arguments.receiver: GroupClauseApproximation by arg()`
`26`	`26`	`val Arguments.column: String by arg()`
`27`	`27`
`28`	`28`	`override fun Arguments.interpret(): PluginDataFrameSchema {`
`29`		`- val grouped = groupImpl(receiver.df.columns(), receiver.columns.mapTo(mutableSetOf()) { it.path.path }, column)`
	`29`	`+ val grouped = groupImpl(receiver.df.columns(), receiver.columns.resolve(receiver.df).mapTo(mutableSetOf()) { it.path.path }, column)`
`30`	`30`	`return grouped`
`31`	`31`	`}`
`32`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -32,10 +32,10 @@ class GroupBy(val df: PluginDataFrameSchema, val keys: List<ColumnWithPathApprox`
`32`	`32`	`class DataFrameGroupBy : AbstractInterpreter<GroupBy>() {`
`33`	`33`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`34`	`34`	`val Arguments.moveToTop: Boolean by arg(defaultValue = Present(true))`
`35`		`- val Arguments.cols: List<ColumnWithPathApproximation> by arg()`
	`35`	`+ val Arguments.cols: ColumnsResolver by arg()`
`36`	`36`
`37`	`37`	`override fun Arguments.interpret(): GroupBy {`
`38`		`- return GroupBy(receiver, cols, moveToTop)`
	`38`	`+ return GroupBy(receiver, cols.resolve(receiver), moveToTop)`
`39`	`39`	`}`
`40`	`40`	`}`
`41`	`41`
Original file line number	Diff line number	Diff line change
`@@ -11,10 +11,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame`
`11`	`11`
`12`	`12`	`class Remove0 : AbstractSchemaModificationInterpreter() {`
`13`	`13`	`val Arguments.receiver: PluginDataFrameSchema by dataFrame()`
`14`		`- val Arguments.columns: List<ColumnWithPathApproximation> by arg()`
	`14`	`+ val Arguments.columns: ColumnsResolver by arg()`
`15`	`15`
`16`	`16`	`override fun Arguments.interpret(): PluginDataFrameSchema {`
`17`		`- val removeResult = removeImpl(receiver.columns(), columns.mapTo(mutableSetOf()) { it.path.path })`
	`17`	`+ val removeResult = removeImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path })`
`18`	`18`	`return PluginDataFrameSchema(removeResult.updatedColumns)`
`19`	`19`	`}`
`20`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,22 +14,23 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString`
`14`	`14`
`15`	`15`	`class Rename : AbstractInterpreter<RenameClauseApproximation>() {`
`16`	`16`	`private val Arguments.receiver by dataFrame()`
`17`		`- private val Arguments.columns: List<ColumnWithPathApproximation> by arg()`
	`17`	`+ private val Arguments.columns: ColumnsResolver by arg()`
`18`	`18`	`override fun Arguments.interpret(): RenameClauseApproximation {`
`19`	`19`	`return RenameClauseApproximation(receiver, columns)`
`20`	`20`	`}`
`21`	`21`	`}`
`22`	`22`
`23`		`-class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: List<ColumnWithPathApproximation>)`
	`23`	`+class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver)`
`24`	`24`
`25`	`25`	`class RenameInto : AbstractSchemaModificationInterpreter() {`
`26`	`26`	`val Arguments.receiver: RenameClauseApproximation by arg()`
`27`	`27`	`val Arguments.newNames: List<String> by varargString()`
`28`	`28`
`29`	`29`	`override fun Arguments.interpret(): PluginDataFrameSchema {`
`30`		`- require(receiver.columns.size == newNames.size)`
	`30`	`+ val columns = receiver.columns.resolve(receiver.schema)`
	`31`	`+ require(columns.size == newNames.size)`
`31`	`32`	`var i = 0`
`32`		`- return receiver.schema.map(receiver.columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } })`
	`33`	`+ return receiver.schema.map(columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } })`
`33`	`34`	`}`
`34`	`35`	`}`
`35`	`36`