[Compiler plugin] Support DF operations in AggregateGroupedBody scope

Author: koperagen

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

@@ -1,3 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.aggregation
 
+import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
+
+@HasSchema(schemaArg = 0)
 public abstract class AggregateGroupedDsl<out T> : AggregateDsl<T>()

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

@@ -0,0 +1,26 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.io.*
+
+data class Name(val firstName: String, val lastName: String)
+
+data class Score(val subject: String, val value: Int)
+
+data class Student(val name: Name, val age: Int, val scores: List<Score>)
+
+fun box(): String {
+    val students = listOf(
+        Student(Name("Alice", "Cooper"), 15, listOf(Score("math", 4), Score("biology", 3))),
+        Student(Name("Bob", "Marley"), 20, listOf(Score("music", 5))),
+    )
+
+    val df = students.toDataFrame().groupBy { expr { name.firstName} }
+        .aggregate {
+            remove { age } into "a"
+        }
+
+    df.compareSchemas(strict = true)
+    return "OK"
+}

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

@@ -292,6 +292,12 @@ public void testMapToFrame() {
     runTest("testData/box/mapToFrame.kt");
   }
 
+  @Test
+  @TestMetadata("modifySchemaInAggregate.kt")
+  public void testModifySchemaInAggregate() {
+    runTest("testData/box/modifySchemaInAggregate.kt");
+  }
+
   @Test
   @TestMetadata("moveAfter.kt")
   public void testMoveAfter() {