[Compiler plugin] GroupBy.into(column)

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/aggregation/AggregateDsl.kt

@@ -17,7 +17,7 @@ public abstract class AggregateDsl<out T> :
     DataFrame<T>,
     ColumnSelectionDsl<T> {
 
-    @Interpretable("GroupByInto")
+    @Interpretable("AggregateDslInto")
     public inline infix fun <reified R> R.into(name: String): NamedValue =
         internal().yield(pathOf(name), this, typeOf<R>())
 

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

@@ -5,6 +5,8 @@ import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.internal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.withExpr
@@ -14,6 +16,8 @@ import kotlin.reflect.typeOf
 
 // region GroupBy
 
+@Refine
+@Interpretable("GroupByInto")
 public fun <T, G> GroupBy<T, G>.into(column: String): DataFrame<T> = toDataFrame(column)
 
 @AccessApiOverload

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

@@ -44,7 +44,7 @@ class GroupByDsl {
     val columns = mutableListOf<NamedValue>()
 }
 
-class GroupByInto : AbstractInterpreter<Unit>() {
+class AggregateDslInto : AbstractInterpreter<Unit>() {
     val Arguments.dsl: GroupByDsl by arg()
     val Arguments.receiver: FirExpression by arg(lens = Interpreter.Id)
     val Arguments.name: String by arg()
@@ -146,6 +146,18 @@ fun KotlinTypeFacade.createPluginDataFrameSchema(keys: List<ColumnWithPathApprox
     return PluginDataFrameSchema(rootColumns)
 }
 
+class GroupByInto : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.column: String by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val grouped = listOf(SimpleFrameColumn(column, receiver.groups.columns()))
+        return PluginDataFrameSchema(
+            receiver.keys.columns() + grouped
+        )
+    }
+}
+
 class GroupByToDataFrame : AbstractSchemaModificationInterpreter() {
     val Arguments.receiver: GroupBy by groupBy()
     val Arguments.groupedColumnName: String? by arg(defaultValue = Present(null))

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

@@ -24,7 +24,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Explode0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Expr0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.From
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Group0
-import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByInto
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateDslInto
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Insert1
@@ -89,6 +89,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByAdd
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupByInto
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Merge0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MergeId
@@ -242,11 +243,12 @@ internal inline fun <reified T> String.load(): T {
         "Exclude1" -> Exclude1()
         "RenameInto" -> RenameInto()
         "DataFrameGroupBy" -> DataFrameGroupBy()
-        "GroupByInto" -> GroupByInto()
+        "AggregateDslInto" -> AggregateDslInto()
         "ReadJsonStr" -> ReadJsonStr()
         "DataRowReadJsonStr" -> DataRowReadJsonStr()
         "ReadDelimStr" -> ReadDelimStr()
         "GroupByToDataFrame" -> GroupByToDataFrame()
+        "GroupByInto" -> GroupByInto()
         "ToDataFrameFrom0" -> ToDataFrameFrom()
         "All0" -> All0()
         "ColsOf0" -> ColsOf0()

plugins/kotlin-dataframe/testData/box/groupBy_toDataFrame.kt

@@ -11,5 +11,9 @@ fun box(): String {
     val df1 = df.groupBy { b }.toDataFrame()
     df1.group[0].a
     df1.group[0].b
+
+    val df2 = df.groupBy { b }.into("gr")
+    df2.gr[0].a
+    df2.gr[0].b
     return "OK"
 }