wip making Aggregators more modular

Author: Jolanrensen

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

@@ -34,7 +34,7 @@ import kotlin.reflect.typeOf
 
 // region DataColumn
 
-public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Double = Aggregators.mean(skipNA).aggregate(this)
+public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Double = Aggregators.mean(skipNA).aggregateSingleColumn(this)
 
 public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
     skipNA: Boolean = skipNA_default,

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

@@ -13,6 +13,7 @@ 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
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.aggregateCalculatingValueType
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
@@ -27,7 +28,8 @@ import kotlin.reflect.KProperty
 
 public fun <T : Comparable<T>> DataColumn<T?>.median(): T = medianOrNull().suggestIfNull("median")
 
-public fun <T : Comparable<T>> DataColumn<T?>.medianOrNull(): T? = Aggregators.median.cast<T>().aggregate(this)
+public fun <T : Comparable<T>> DataColumn<T?>.medianOrNull(): T? =
+    Aggregators.median.cast<T>().aggregateSingleColumn(this)
 
 public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOfOrNull(noinline expression: (T) -> R?): R? =
     Aggregators.median.cast<R?>().aggregateOf(this, expression)
@@ -40,7 +42,7 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOf(noinline
 // region DataRow
 
 public fun AnyRow.rowMedianOrNull(): Any? =
-    Aggregators.median.aggregateCalculatingType(
+    Aggregators.median.aggregateCalculatingValueType(
         values = values().filterIsInstance<Comparable<Any?>>().asIterable(),
         valueTypes = df().columns().filter { it.valuesAreComparable() }.map { it.type() }.toSet(),
     )

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

@@ -28,7 +28,7 @@ import kotlin.reflect.KProperty
 
 public fun <T : Comparable<T>> DataColumn<T?>.min(): T = minOrNull().suggestIfNull("min")
 
-public fun <T : Comparable<T>> DataColumn<T?>.minOrNull(): T? = Aggregators.min<T>().aggregate(this)
+public fun <T : Comparable<T>> DataColumn<T?>.minOrNull(): T? = Aggregators.min<T>().aggregateSingleColumn(this)
 
 public inline fun <T, reified R : Comparable<R>> DataColumn<T>.minBy(noinline selector: (T) -> R?): T =
     minByOrNull(selector).suggestIfNull("minBy")

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

@@ -27,7 +27,7 @@ public fun <T : Comparable<T>> DataColumn<T?>.percentile(percentile: Double): T
     percentileOrNull(percentile).suggestIfNull("percentile")
 
 public fun <T : Comparable<T>> DataColumn<T?>.percentileOrNull(percentile: Double): T? =
-    Aggregators.percentile(percentile).cast<T>().aggregate(this)
+    Aggregators.percentile(percentile).cast<T>().aggregateSingleColumn(this)
 
 public inline fun <T, reified R : Comparable<R>> DataColumn<T>.percentileOfOrNull(
     percentile: Double,
@@ -44,7 +44,7 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.percentileOf(
 // region DataRow
 
 public fun AnyRow.rowPercentileOrNull(percentile: Double): Any? =
-    Aggregators.percentile(percentile).aggregate(
+    Aggregators.percentile(percentile).aggregateSingleColumn(
         values().filterIsInstance<Comparable<Any?>>().toValueColumn(),
     )
 

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

@@ -26,7 +26,7 @@ import kotlin.reflect.typeOf
 // region DataColumn
 
 public fun <T : Number> DataColumn<T?>.std(skipNA: Boolean = skipNA_default, ddof: Int = ddof_default): Double =
-    Aggregators.std(skipNA, ddof).aggregate(this) ?: .0
+    Aggregators.std(skipNA, ddof).aggregateSingleColumn(this) ?: .0
 
 public inline fun <T, reified R : Number> DataColumn<T>.stdOf(
     skipNA: Boolean = skipNA_default,

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

@@ -42,14 +42,14 @@ import kotlin.reflect.typeOf
 // region DataColumn
 
 @JvmName("sumShort")
-public fun DataColumn<Short?>.sum(): Int = Aggregators.sum.aggregate(this) as Int
+public fun DataColumn<Short?>.sum(): Int = Aggregators.sum.aggregateSingleColumn(this) as Int
 
 @JvmName("sumByte")
-public fun DataColumn<Byte?>.sum(): Int = Aggregators.sum.aggregate(this) as Int
+public fun DataColumn<Byte?>.sum(): Int = Aggregators.sum.aggregateSingleColumn(this) as Int
 
 @Suppress("UNCHECKED_CAST")
 @JvmName("sumNumber")
-public fun <T : Number> DataColumn<T?>.sum(): T = Aggregators.sum.aggregate(this) as T
+public fun <T : Number> DataColumn<T?>.sum(): T = Aggregators.sum.aggregateSingleColumn(this) as T
 
 @JvmName("sumOfShort")
 @OverloadResolutionByLambdaReturnType

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt

@@ -25,63 +25,83 @@ internal interface Aggregator<in Value, out Return> {
     /**
      * Base function of [Aggregator].
      *
-     * Aggregates the given values, taking [type] into account,
+     * Aggregates the given values, taking [valueType] into account,
      * filtering nulls (only if [type.isMarkedNullable][KType.isMarkedNullable]),
      * and computes a single resulting value.
      *
-     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregator] argument.
+     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregateSingle] argument.
      *
-     * When the exact [type] is unknown, use [aggregateCalculatingType].
+     * When the exact [valueType] is unknown, use [aggregateCalculatingValueType].
      */
-    fun aggregate(values: Iterable<Value?>, type: KType): Return
+    fun aggregateSingleIterable(values: Iterable<Value?>, valueType: KType): Return
+
+    fun calculateValueType(valueTypes: Set<KType>): KType
+
+    fun calculateValueType(values: Iterable<Value?>): KType
 
     /**
      * Aggregates the data in the given column and computes a single resulting value.
-     * Calls [aggregate] (with [Iterable] and [KType]).
+     * Calls [aggregateSingleColumn] (with [Iterable] and [KType]).
      *
-     * See [AggregatorBase.aggregate].
+     * See [AggregatorBase.aggregateSingleIterable].
      */
-    fun aggregate(column: DataColumn<Value?>): Return
+    fun aggregateSingleColumn(column: DataColumn<Value?>): Return
 
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
      */
-    fun aggregate(columns: Iterable<DataColumn<Value?>>): Return
+    fun aggregateMultipleColumns(columns: Iterable<DataColumn<Value?>>): Return
 
     /**
-     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
-     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
-     *
-     * @param values The values to be aggregated.
-     * @param valueTypes The types of the values.
-     *   If provided, this can be used to avoid calculating the types of [values] at runtime with reflection.
-     *   It should contain all types of [values].
-     *   If `null` or empty, the types of [values] will be calculated at runtime (heavy!).
-     */
-    fun aggregateCalculatingType(values: Iterable<Value?>, valueTypes: Set<KType>? = null): Return
-
-    /**
-     * Function that can give the return type of [aggregate] as [KType], given the type of the input.
+     * Function that can give the return type of [aggregateSingleIterable] as [KType], given the type of the input.
      * This allows aggregators to avoid runtime type calculations.
      *
      * @param type The type of the input values.
      * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
-     * @return The return type of [aggregate] as [KType].
+     * @return The return type of [aggregateSingleIterable] as [KType].
      */
     fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
 
     /**
-     * Function that can give the return type of [aggregate] with columns as [KType],
+     * Function that can give the return type of [aggregateSingleIterable] with columns as [KType],
      * given the multiple types of the input.
      * This allows aggregators to avoid runtime type calculations.
      *
      * @param colTypes The types of the input columns.
      * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
-     * @return The return type of [aggregate] as [KType].
+     * @return The return type of [aggregateSingleIterable] as [KType].
      */
-    fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
+    fun calculateReturnTypeMultipleColumnsOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
+
+    val ref: Aggregator<Value, Return> get() = this
 }
 
+internal fun <Value, Return> Aggregator<Value, Return>.aggregate(values: Iterable<Value?>, valueType: KType) =
+    aggregateSingleIterable(values, valueType)
+
+internal fun <Value, Return> Aggregator<Value, Return>.calculateValueType(
+    values: Iterable<Value?>,
+    valueTypes: Set<KType>? = null,
+) = if (valueTypes != null && valueTypes.isNotEmpty()) {
+    calculateValueType(valueTypes)
+} else {
+    calculateValueType(values)
+}
+
+internal fun <Value, Return> Aggregator<Value, Return>.aggregateCalculatingValueType(
+    values: Iterable<Value?>,
+    valueTypes: Set<KType>? = null,
+) = aggregateSingleIterable(
+    values = values,
+    valueType = calculateValueType(values, valueTypes),
+)
+
+internal fun <Value, Return> Aggregator<Value, Return>.aggregate(column: DataColumn<Value?>) =
+    aggregateSingleColumn(column)
+
+internal fun <Value, Return> Aggregator<Value, Return>.aggregate(columns: Iterable<DataColumn<Value?>>) =
+    aggregateMultipleColumns(columns)
+
 @Suppress("UNCHECKED_CAST")
 @PublishedApi
 internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
@@ -90,15 +110,18 @@ internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Ag
 @PublishedApi
 internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>
 
-/** Type alias for [Aggregator.calculateReturnTypeOrNull] */
+/** Type alias for [Aggregator.calculateReturnTypeMultipleColumnsOrNull] */
 internal typealias CalculateReturnTypeOrNull = (type: KType, emptyInput: Boolean) -> KType?
 
 /**
- * Type alias for the argument for [Aggregator.aggregate].
+ * Type alias for the argument for [Aggregator.aggregateSingleIterable].
  * Nulls have already been filtered out when this argument is called.
  */
 internal typealias Aggregate<Value, Return> = Iterable<Value & Any>.(type: KType) -> Return
 
+internal typealias AggregateBy<Source, Value, Return> =
+    Iterable<Source>.(sourceType: KType, valueType: KType, selector: (Source) -> Value) -> Return
+
 /** Common case for [CalculateReturnTypeOrNull], preserves return type, but makes it nullable for empty inputs. */
 internal val preserveReturnTypeNullIfEmpty: CalculateReturnTypeOrNull = { type, emptyInput ->
     type.withNullability(emptyInput)

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt

@@ -2,9 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.asIterable
-import org.jetbrains.kotlinx.dataframe.api.asSequence
-import org.jetbrains.kotlinx.dataframe.impl.commonType
-import org.jetbrains.kotlinx.dataframe.impl.nothingType
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
@@ -15,100 +12,49 @@ import kotlin.reflect.full.withNullability
  * or multiple [DataColumns][DataColumn].
  *
  * @param name The name of this aggregator.
- * @param aggregator Functional argument for the [aggregate] function. Nulls are filtered out before this is called.
  */
 internal abstract class AggregatorBase<in Value, out Return>(
     override val name: String,
-    protected val getReturnTypeOrNull: CalculateReturnTypeOrNull,
-    protected val aggregator: Aggregate<Value, Return>,
+    private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    private val aggregateSingle: Aggregate<Value, Return>,
 ) : Aggregator<Value, Return> {
 
     /**
      * Base function of [Aggregator].
      *
-     * Aggregates the given values, taking [type] into account,
+     * Aggregates the given values, taking [valueType] into account,
      * filtering nulls (only if [type.isMarkedNullable][KType.isMarkedNullable]),
      * and computes a single resulting value.
      *
-     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregator] argument.
+     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregateSingle] argument.
      *
-     * When the exact [type] is unknown, use [aggregateCalculatingType].
+     * When the exact [valueType] is unknown, use [Aggregator.aggregateCalculatingValueType].
      */
     @Suppress("UNCHECKED_CAST")
-    override fun aggregate(values: Iterable<Value?>, type: KType): Return =
-        aggregator(
+    override fun aggregateSingleIterable(values: Iterable<Value?>, valueType: KType): Return =
+        aggregateSingle(
             // values =
-            if (type.isMarkedNullable) {
+            if (valueType.isMarkedNullable) {
                 values.asSequence().filterNotNull().asIterable()
             } else {
                 values as Iterable<Value & Any>
             },
             // type =
-            type.withNullability(false),
+            valueType.withNullability(false),
         )
 
-    /**
-     * Function that can give the return type of [aggregate] as [KType], given the type of the input.
-     * This allows aggregators to avoid runtime type calculations.
-     *
-     * Uses [getReturnTypeOrNull] to calculate the return type.
-     *
-     * @param type The type of the input values.
-     * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
-     * @return The return type of [aggregate] as [KType].
-     */
-    override fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType? =
-        getReturnTypeOrNull(type.withNullability(false), emptyInput)
-
     /**
      * Aggregates the data in the given column and computes a single resulting value.
      *
-     * Nulls are filtered out by default, then [aggregate] (with [Iterable] and [KType]) is called.
+     * Nulls are filtered out by default, then [aggregateSingleColumn] (with [Iterable] and [KType]) is called.
      */
     @Suppress("UNCHECKED_CAST")
-    override fun aggregate(column: DataColumn<Value?>): Return =
-        aggregate(
+    final override fun aggregateSingleColumn(column: DataColumn<Value?>): Return =
+        aggregateSingleIterable(
             values = column.asIterable(),
-            type = column.type(),
+            valueType = column.type(),
         )
 
-    /** @include [Aggregator.aggregateCalculatingType] */
-    override fun aggregateCalculatingType(values: Iterable<Value?>, valueTypes: Set<KType>?): Return {
-        val commonType = if (valueTypes != null && valueTypes.isNotEmpty()) {
-            valueTypes.commonType(false)
-        } else {
-            var hasNulls = false
-            val classes = values.mapNotNull {
-                if (it == null) {
-                    hasNulls = true
-                    null
-                } else {
-                    it.javaClass.kotlin
-                }
-            }
-            if (classes.isEmpty()) {
-                nothingType(hasNulls)
-            } else {
-                classes.commonType(hasNulls)
-            }
-        }
-        return aggregate(values, commonType)
-    }
-
-    /**
-     * Aggregates the data in the multiple given columns and computes a single resulting value.
-     * Must be overridden to use.
-     */
-    abstract override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return
-
-    /**
-     * Function that can give the return type of [aggregate] with columns as [KType],
-     * given the multiple types of the input.
-     * This allows aggregators to avoid runtime type calculations.
-     *
-     * @param colTypes The types of the input columns.
-     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
-     * @return The return type of [aggregate] as [KType].
-     */
-    abstract override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
+    final override fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType? =
+        getReturnTypeOrNull(type, emptyInput)
 }