Add Compiler Plugin support for statistics on GroupBy (#1077)

* Add GroupBySumOf functionality to groupBy operations

Introduces the `GroupBySumOf` interpreter for aggregation, enabling the calculation of column sums with customizable expressions and result names in grouped DataFrames. Adds tests and updates APIs to support and validate this feature.

* Add commented GroupBySum0 support with updated scenarios and tests

* Refactor GroupBy statistics functionality and tests.

Updated statistical aggregation functions for GroupBy with comments addressing open questions. Added comprehensive tests to verify behavior across various statistics (sum, mean, median, std, min, and max), replacing older test cases for cleaner coverage.

* Add support for GroupBy mean and median operations.

* Add support for min and max functions

* Added support for std function

* Updated support for sum/sumFor

* Added support for all statistics but faced with limitation of FiR

* Fixed comparable types for median

* Fixed for max/min

* Fixed for std/mean

* added missed casts to median/percentile. Could result in Comparable<Any?> columns

* Refactor groupBy for enhanced type safety and comparability

Replaced direct subtype checks with `isIntraComparable` to improve type safety when resolving columns. Updated documentation syntax for better consistency and clarity. Added schema comparison in test to validate grouping behavior.

* Refactor `GroupBy` aggregation classes and test handling.

Revised `GroupBy` aggregation logic by restructuring classes, improving naming consistency, and refining comments/documentation. Updated test cases to address initializer type mismatches and better handle scenarios involving multiple columns. Added relevant TODOs for unresolved cases linked to issue #1090.

---------

Co-authored-by: Jolan Rensen <jolan.rensen@jetbrains.com>

Author: zaleslaw

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

@@ -133,9 +133,12 @@ public fun <T, C : Comparable<C>> DataFrame<T>.maxByOrNull(column: KProperty<C?>
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupByMax1")
 public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(interComparableColumns())
 
+@Refine
+@Interpretable("GroupByMax0")
 public fun <T, C : Comparable<C>> Grouped<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.max.aggregateFor(this, columns)
 
@@ -149,6 +152,8 @@ public fun <T, C : Comparable<C>> Grouped<T>.maxFor(vararg columns: ColumnRefere
 public fun <T, C : Comparable<C>> Grouped<T>.maxFor(vararg columns: KProperty<C?>): DataFrame<T> =
     maxFor { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMax0")
 public fun <T, C : Comparable<C>> Grouped<T>.max(name: String? = null, columns: ColumnsSelector<T, C?>): DataFrame<T> =
     Aggregators.max.aggregateAll(this, name, columns)
 

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

@@ -8,6 +8,8 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 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.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.toColumnsSetOf
@@ -98,9 +100,12 @@ public inline fun <T, reified D : Number> DataFrame<T>.meanOf(
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupByMean1")
 public fun <T> Grouped<T>.mean(skipNA: Boolean = skipNA_default): DataFrame<T> = meanFor(skipNA, numberColumns())
 
+@Refine
+@Interpretable("GroupByMean0")
 public fun <T, C : Number> Grouped<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsForAggregateSelector<T, C?>,
@@ -121,6 +126,8 @@ public fun <T, C : Number> Grouped<T>.meanFor(
     skipNA: Boolean = skipNA_default,
 ): DataFrame<T> = meanFor(skipNA) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMean0")
 public fun <T, C : Number> Grouped<T>.mean(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
@@ -147,6 +154,8 @@ public fun <T, C : Number> Grouped<T>.mean(
     skipNA: Boolean = skipNA_default,
 ): DataFrame<T> = mean(name, skipNA) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMeanOf")
 public inline fun <T, reified R : Number> Grouped<T>.meanOf(
     name: String? = null,
     skipNA: Boolean = skipNA_default,

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

@@ -8,6 +8,8 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 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.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
@@ -16,6 +18,7 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.columns.toComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull
@@ -103,9 +106,12 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.medianOf(
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupByMedian1")
 public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(interComparableColumns())
 
+@Refine
+@Interpretable("GroupByMedian0")
 public fun <T, C : Comparable<C>> Grouped<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.median.aggregateFor(this, columns)
 
@@ -119,6 +125,8 @@ public fun <T, C : Comparable<C>> Grouped<T>.medianFor(vararg columns: ColumnRef
 public fun <T, C : Comparable<C>> Grouped<T>.medianFor(vararg columns: KProperty<C?>): DataFrame<T> =
     medianFor { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMedian0")
 public fun <T, C : Comparable<C>> Grouped<T>.median(
     name: String? = null,
     columns: ColumnsSelector<T, C?>,
@@ -137,10 +145,12 @@ public fun <T, C : Comparable<C>> Grouped<T>.median(
 public fun <T, C : Comparable<C>> Grouped<T>.median(vararg columns: KProperty<C?>, name: String? = null): DataFrame<T> =
     median(name) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMedianOf")
 public inline fun <T, reified R : Comparable<R>> Grouped<T>.medianOf(
     name: String? = null,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.median.aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.median.cast<R?>().aggregateOf(this, name, expression)
 
 // endregion
 
@@ -227,6 +237,6 @@ public fun <T, C : Comparable<C>> PivotGroupBy<T>.median(vararg columns: KProper
 
 public inline fun <T, reified R : Comparable<R>> PivotGroupBy<T>.medianOf(
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.median.aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.median.cast<R?>().aggregateOf(this, expression)
 
 // endregion

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

@@ -133,9 +133,12 @@ public fun <T, C : Comparable<C>> DataFrame<T>.minByOrNull(column: KProperty<C?>
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupByMin1")
 public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(interComparableColumns())
 
+@Refine
+@Interpretable("GroupByMin0")
 public fun <T, C : Comparable<C>> Grouped<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.min.aggregateFor(this, columns)
 
@@ -149,6 +152,8 @@ public fun <T, C : Comparable<C>> Grouped<T>.minFor(vararg columns: ColumnRefere
 public fun <T, C : Comparable<C>> Grouped<T>.minFor(vararg columns: KProperty<C?>): DataFrame<T> =
     minFor { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByMin0")
 public fun <T, C : Comparable<C>> Grouped<T>.min(name: String? = null, columns: ColumnsSelector<T, C?>): DataFrame<T> =
     Aggregators.min.aggregateAll(this, name, columns)
 

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

@@ -177,7 +177,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.percentileOf(
     percentile: Double,
     name: String? = null,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.percentile(percentile).aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.percentile(percentile).cast<R?>().aggregateOf(this, name, expression)
 
 // endregion
 
@@ -289,6 +289,6 @@ public fun <T, C : Comparable<C>> PivotGroupBy<T>.percentile(
 public inline fun <T, reified R : Comparable<R>> PivotGroupBy<T>.percentileOf(
     percentile: Double,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.percentile(percentile).aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.percentile(percentile).cast<R?>().aggregateOf(this, expression)
 
 // endregion

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

@@ -8,6 +8,8 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 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.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.toColumnsSetOf
@@ -102,10 +104,13 @@ public inline fun <T, reified R : Number> DataFrame<T>.stdOf(
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupByStd1")
 public fun <T> Grouped<T>.std(skipNA: Boolean = skipNA_default, ddof: Int = ddof_default): DataFrame<T> =
     stdFor(skipNA, ddof, numberColumns())
 
+@Refine
+@Interpretable("GroupByStd0")
 public fun <T> Grouped<T>.stdFor(
     skipNA: Boolean = skipNA_default,
     ddof: Int = ddof_default,
@@ -118,6 +123,7 @@ public fun <T> Grouped<T>.stdFor(
     ddof: Int = ddof_default,
 ): DataFrame<T> = stdFor(skipNA, ddof) { columns.toColumnsSetOf() }
 
+@AccessApiOverload
 public fun <T, C : Number> Grouped<T>.stdFor(
     vararg columns: ColumnReference<C?>,
     skipNA: Boolean = skipNA_default,
@@ -131,13 +137,16 @@ public fun <T, C : Number> Grouped<T>.stdFor(
     ddof: Int = ddof_default,
 ): DataFrame<T> = stdFor(skipNA, ddof) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByStd0")
 public fun <T> Grouped<T>.std(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     ddof: Int = ddof_default,
     columns: ColumnsSelector<T, Number?>,
 ): DataFrame<T> = Aggregators.std(skipNA, ddof).aggregateAll(this, name, columns)
 
+@AccessApiOverload
 public fun <T> Grouped<T>.std(
     vararg columns: ColumnReference<Number?>,
     name: String? = null,
@@ -160,6 +169,8 @@ public fun <T> Grouped<T>.std(
     ddof: Int = ddof_default,
 ): DataFrame<T> = std(name, skipNA, ddof) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupByStdOf")
 public inline fun <T, reified R : Number> Grouped<T>.stdOf(
     name: String? = null,
     skipNA: Boolean = skipNA_default,

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

@@ -8,6 +8,8 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 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.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.toColumnsSetOf
@@ -89,9 +91,12 @@ public inline fun <T, reified C : Number?> DataFrame<T>.sumOf(crossinline expres
 // endregion
 
 // region GroupBy
-
+@Refine
+@Interpretable("GroupBySum1")
 public fun <T> Grouped<T>.sum(): DataFrame<T> = sumFor(numberColumns())
 
+@Refine
+@Interpretable("GroupBySum0")
 public fun <T, C : Number> Grouped<T>.sumFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.sum.aggregateFor(this, columns)
 
@@ -105,6 +110,8 @@ public fun <T, C : Number> Grouped<T>.sumFor(vararg columns: ColumnReference<C?>
 public fun <T, C : Number> Grouped<T>.sumFor(vararg columns: KProperty<C?>): DataFrame<T> =
     sumFor { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupBySum0")
 public fun <T, C : Number> Grouped<T>.sum(name: String? = null, columns: ColumnsSelector<T, C?>): DataFrame<T> =
     Aggregators.sum.aggregateAll(this, name, columns)
 
@@ -119,6 +126,8 @@ public fun <T, C : Number> Grouped<T>.sum(vararg columns: ColumnReference<C?>, n
 public fun <T, C : Number> Grouped<T>.sum(vararg columns: KProperty<C?>, name: String? = null): DataFrame<T> =
     sum(name) { columns.toColumnSet() }
 
+@Refine
+@Interpretable("GroupBySumOf")
 public inline fun <T, reified R : Number> Grouped<T>.sumOf(
     resultName: String? = null,
     crossinline expression: RowExpression<T, R?>,