[Compiler plugin] Support transformation of functions with GroupBy type

Author: koperagen

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

@@ -29,6 +29,7 @@ import kotlin.reflect.KProperty
  *
  * `df.add("columnName") { "someColumn"<Int>() + 15 }.groupBy("columnName")`
  */
+@Refine
 @Interpretable("DataFrameGroupBy")
 public fun <T> DataFrame<T>.groupBy(moveToTop: Boolean = true, cols: ColumnsSelector<T, *>): GroupBy<T, T> =
     groupByImpl(moveToTop, cols)

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

@@ -16,8 +16,11 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.aggregate
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.createPluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.DF_CLASS_ID
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names.GROUP_BY_CLASS_ID
 
 fun KotlinTypeFacade.analyzeRefinedCallShape(call: FirFunctionCall, reporter: InterpretationErrorReporter): CallResult? {
     val callReturnType = call.resolvedType
@@ -61,6 +64,53 @@ fun KotlinTypeFacade.analyzeRefinedCallShape(call: FirFunctionCall, reporter: In
     return CallResult(rootMarker, newSchema)
 }
 
+fun KotlinTypeFacade.analyzeRefinedGroupByCallShape(call: FirFunctionCall, reporter: InterpretationErrorReporter): GroupByCallResult? {
+    val callReturnType = call.resolvedType
+    if (callReturnType.classId != GROUP_BY_CLASS_ID) return null
+    val keyMarker = callReturnType.typeArguments[0]
+    val groupMarker = callReturnType.typeArguments[1]
+    // rootMarker is expected to be a token generated by the plugin.
+    // it's implied by "refined call"
+    // thus ConeClassLikeType
+    if (keyMarker !is ConeClassLikeType || groupMarker !is ConeClassLikeType) {
+        return null
+    }
+
+    val newSchema = call.interpreterName(session)?.let { name ->
+        name.load<Interpreter<*>>().let { processor ->
+            val dataFrameSchema = interpret(call, processor, reporter = reporter)
+                .let {
+                    val value = it?.value
+                    if (value !is GroupBy) {
+                        if (!reporter.errorReported) {
+                            reporter.reportInterpretationError(call, "${processor::class} must return ${PluginDataFrameSchema::class}, but was ${value}")
+                        }
+                        return null
+                    }
+                    value
+                }
+
+            val keySchema = createPluginDataFrameSchema(dataFrameSchema.keys, dataFrameSchema.moveToTop)
+            val groupSchema = PluginDataFrameSchema(dataFrameSchema.df.columns())
+            GroupBySchema(keySchema, groupSchema)
+        }
+    } ?: return null
+
+    return GroupByCallResult(keyMarker, newSchema.keySchema, groupMarker, newSchema.groupSchema)
+}
+
+data class GroupByCallResult(
+    val keyMarker: ConeClassLikeType,
+    val keySchema: PluginDataFrameSchema,
+    val groupMarker: ConeClassLikeType,
+    val groupSchema: PluginDataFrameSchema,
+)
+
+data class GroupBySchema(
+    val keySchema: PluginDataFrameSchema,
+    val groupSchema: PluginDataFrameSchema,
+)
+
 data class CallResult(val rootMarker: ConeClassLikeType, val newSchema: PluginDataFrameSchema)
 
 class RefinedArguments(val refinedArguments: List<RefinedArgument>) : List<RefinedArgument> by refinedArguments