Merge pull request #1088 from Kotlin/aggregate-datarow

[Compiler plugin] Support DataFrame.aggregate

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataFrame.kt

@@ -4,6 +4,8 @@ import org.jetbrains.kotlinx.dataframe.aggregation.Aggregatable
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedBody
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
 import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
 import org.jetbrains.kotlinx.dataframe.api.add
 import org.jetbrains.kotlinx.dataframe.api.cast
@@ -71,6 +73,8 @@ public interface DataFrame<out T> :
 
     // endregion
 
+    @Refine
+    @Interpretable("AggregateRow")
     public fun <R> aggregate(body: AggregateGroupedBody<T, R>): DataRow<T>
 
     // region get columns

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

@@ -72,6 +72,18 @@ class Aggregate : AbstractSchemaModificationInterpreter() {
     }
 }
 
+class AggregateRow : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.body: FirAnonymousFunctionExpression by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return aggregate(
+            GroupBy(PluginDataFrameSchema.EMPTY, receiver),
+            InterpretationErrorReporter.DEFAULT,
+            body
+        )
+    }
+}
+
 fun KotlinTypeFacade.aggregate(
     groupBy: GroupBy,
     reporter: InterpretationErrorReporter,

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

@@ -70,6 +70,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslNamedGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AggregateRow
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All2
@@ -395,6 +396,7 @@ internal inline fun <reified T> String.load(): T {
         "ToTop" -> ToTop()
         "Update0" -> Update0()
         "Aggregate" -> Aggregate()
+        "AggregateRow" -> AggregateRow()
         "DataFrameOf3" -> DataFrameOf3()
         "ValueCounts" -> ValueCounts()
         "RenameToCamelCase" -> RenameToCamelCase()

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

@@ -0,0 +1,14 @@
+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 row = dataFrameOf("a" to List(10) { it }).aggregate {
+        maxOf { a } into "max"
+        minOf { a } into "min"
+    }
+    val i: Int = row.max
+    val i1: Int = row.min
+    return "OK"
+}

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

@@ -29,6 +29,12 @@ public void testAddId() {
     runTest("testData/box/addId.kt");
   }
 
+  @Test
+  @TestMetadata("aggregateDataFrame.kt")
+  public void testAggregateDataFrame() {
+    runTest("testData/box/aggregateDataFrame.kt");
+  }
+
   @Test
   public void testAllFilesPresentInBox() {
     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("testData/box"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);