Merge pull request #1114 from Kotlin/asGroupBy

[Compiler plugin] Support df.asGroupBy

Author: koperagen

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

@@ -10,6 +10,8 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 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.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
@@ -373,11 +375,15 @@ public fun <T, G> DataFrame<T>.asGroupBy(groupedColumn: ColumnReference<DataFram
     return asGroupBy { groups }
 }
 
+@Refine
+@Interpretable("AsGroupByDefault")
 public fun <T> DataFrame<T>.asGroupBy(): GroupBy<T, T> {
     val groupCol = columns().single { it.isFrameColumn() }.asAnyFrameColumn().castFrameColumn<T>()
     return asGroupBy { groupCol }
 }
 
+@Refine
+@Interpretable("AsGroupBy")
 public fun <T, G> DataFrame<T>.asGroupBy(selector: ColumnSelector<T, DataFrame<G>>): GroupBy<T, G> {
     val column = getColumn(selector).asFrameColumn()
     return GroupByImpl(this.move { column }.toEnd(), column) { none() }

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.types.resolvedType
 import org.jetbrains.kotlin.fir.types.typeContext
 import org.jetbrains.kotlin.fir.types.withNullability
 import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
 import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
@@ -27,12 +28,14 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.add
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
 import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
 import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
 import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 import org.jetbrains.kotlinx.dataframe.plugin.interpret
 import org.jetbrains.kotlinx.dataframe.plugin.loadInterpreter
@@ -53,6 +56,33 @@ class DataFrameGroupBy : AbstractInterpreter<GroupBy>() {
     }
 }
 
+class AsGroupBy : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): GroupBy {
+        val column = selector.resolve(receiver).singleOrNull()?.column
+        return if (column is SimpleFrameColumn) {
+            GroupBy(receiver.asDataFrame().remove { selector }.toPluginDataFrameSchema(), PluginDataFrameSchema(column.columns()))
+        } else {
+            GroupBy.EMPTY
+        }
+    }
+}
+
+class AsGroupByDefault : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+
+    override fun Arguments.interpret(): GroupBy {
+        val groups = receiver.columns().singleOrNull { it is SimpleFrameColumn } as? SimpleFrameColumn
+        return if (groups != null) {
+            GroupBy(receiver.asDataFrame().remove(groups.name).toPluginDataFrameSchema(), PluginDataFrameSchema(groups.columns()))
+        } else {
+            GroupBy.EMPTY
+        }
+    }
+}
+
 class NamedValue(val name: String, val type: ConeKotlinType)
 
 class GroupByDsl {

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

@@ -83,6 +83,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColumnsResolver
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.SingleColumnApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
 
 fun <T> KotlinTypeFacade.interpret(
     functionCall: FirFunctionCall,
@@ -400,7 +401,7 @@ private fun KotlinTypeFacade.columnWithPathApproximations(result: FirPropertyAcc
                     is ConeStarProjection -> session.builtinTypes.nullableAnyType.type
                     else -> arg as ConeClassLikeType
                 }
-                SimpleDataColumn(f(result), Marker(type))
+                simpleColumnOf(f(result), type)
             }
             Names.COLUM_GROUP_CLASS_ID -> {
                 val arg = it.typeArguments.single()

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

@@ -87,6 +87,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllFrom2
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AllUpTo2
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AsGroupBy
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AsGroupByDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ByName
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColGroups1
@@ -324,6 +326,8 @@ internal inline fun <reified T> String.load(): T {
         "Exclude1" -> Exclude1()
         "RenameInto" -> RenameInto()
         "DataFrameGroupBy" -> DataFrameGroupBy()
+        "AsGroupBy" -> AsGroupBy()
+        "AsGroupByDefault" -> AsGroupByDefault()
         "AggregateDslInto" -> AggregateDslInto()
         "GroupByToDataFrame" -> GroupByToDataFrame()
         "GroupByInto" -> GroupByInto()

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

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("i", "group")(1, dataFrameOf("a", "b")(111, 222))
+    val aggregated1 = df.asGroupBy { group }.aggregate { maxOf { a } into "max" }
+    val aggregated2 = df.asGroupBy().aggregate { maxOf { a } into "max" }
+
+    val i: Int = aggregated1.max[0]
+
+    compareSchemas(aggregated1, aggregated2)
+    return "OK"
+}

plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

@@ -40,6 +40,12 @@ public void testAllFilesPresentInBox() {
     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("testData/box"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
   }
 
+  @Test
+  @TestMetadata("asGroupBy.kt")
+  public void testAsGroupBy() {
+    runTest("testData/box/asGroupBy.kt");
+  }
+
   @Test
   @TestMetadata("castTo.kt")
   public void testCastTo() {