WIP rework of aggregator implementation

Author: Jolanrensen

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

@@ -39,8 +39,9 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOf(noinline
 // region DataRow
 
 public fun AnyRow.rowMedianOrNull(): Any? =
-    Aggregators.median.aggregateMixed(
-        values().filterIsInstance<Comparable<Any?>>().asIterable(),
+    Aggregators.median.aggregateCalculatingType(
+        values = values().filterIsInstance<Comparable<Any?>>().asIterable(),
+        valueTypes = df().columns().filter { it.valuesAreComparable() }.map { it.type() }.toSet(),
     )
 
 public fun AnyRow.rowMedian(): Any = rowMedianOrNull().suggestIfNull("rowMedian")

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

@@ -44,9 +44,9 @@ public inline fun <T, reified R : Number> DataColumn<T>.sumOf(crossinline expres
 // region DataRow
 
 public fun AnyRow.rowSum(): Number =
-    Aggregators.sum.aggregateMixed(
+    Aggregators.sum.aggregateCalculatingType(
         values = values().filterIsInstance<Number>(),
-        types = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
+        valueTypes = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
     ) ?: 0
 
 public inline fun <reified T : Number> AnyRow.rowSumOf(): T = values().filterIsInstance<T>().sum(typeOf<T>())

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt

@@ -137,3 +137,15 @@ internal fun Iterable<Number>.convertToUnifiedNumberType(
         converter(it) ?: error("Can not convert $it to $commonNumberType")
     }
 }
+
+/** @include [Iterable.convertToUnifiedNumberType] */
+@JvmName("convertToUnifiedNumberTypeSequence")
+@Suppress("UNCHECKED_CAST")
+internal fun Sequence<Number>.convertToUnifiedNumberType(
+    commonNumberType: KType = asIterable().types().unifiedNumberType(),
+): Sequence<Number> {
+    val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
+    return map {
+        converter(it) ?: error("Can not convert $it to $commonNumberType")
+    }
+}

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

@@ -3,22 +3,63 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import kotlin.reflect.KType
 
+/**
+ * Base interface for all aggregators.
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * The [AggregatorBase] class is a base implementation of this interface.
+ *
+ * @param Value The type of the values to be aggregated.
+ *   This can be nullable for [Iterables][Iterable] or not, depending on the use case.
+ *   For columns, [Value] will always be considered nullable; nulls are filtered out from columns anyway.
+ * @param Return The type of the resulting value. It doesn't matter if this is nullable or not, as the aggregator
+ *   will always return a [Return]`?`.
+ */
 @PublishedApi
-internal interface Aggregator<C, R> {
+internal interface Aggregator<Value, Return> {
 
+    /** The name of this aggregator. */
     val name: String
 
-    fun aggregate(column: DataColumn<C?>): R?
-
+    /** If `true`, [Value][Value]`  ==  ` [Return][Return]. */
     val preservesType: Boolean
 
-    fun aggregate(columns: Iterable<DataColumn<C?>>): R?
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregator] argument.
+     */
+    fun aggregate(values: Iterable<Value>, type: KType): Return?
+
+    /**
+     * Aggregates the data in the given column and computes a single resulting value.
+     * Nulls are filtered out by default, then the aggregation function (with [Iterable] and [KType]) is called.
+     *
+     * See [AggregatorBase.aggregate].
+     */
+    fun aggregate(column: DataColumn<Value?>): Return?
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     *
+     * Must be overridden when using [AggregatorBase].
+     */
+    fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
 
-    fun aggregate(values: Iterable<C>, type: KType): R?
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     */
+    fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>? = null): Return?
 }
 
 @PublishedApi
-internal fun <T> Aggregator<*, *>.cast(): Aggregator<T, T> = this as Aggregator<T, T>
+internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
 
 @PublishedApi
-internal fun <T, P> Aggregator<*, *>.cast2(): Aggregator<T, P> = this as Aggregator<T, P>
+internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>

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

@@ -3,19 +3,69 @@ 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 kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
-internal abstract class AggregatorBase<C, R>(
+/**
+ * Base class for [aggregators][Aggregator].
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * @param name The name of this aggregator.
+ * @param aggregator Functional argument for the [aggregate] function.
+ */
+internal abstract class AggregatorBase<Value, Return>(
     override val name: String,
-    protected val aggregator: (Iterable<C>, KType) -> R?,
-) : Aggregator<C, R> {
+    protected val aggregator: (values: Iterable<Value>, type: KType) -> Return?,
+) : Aggregator<Value, Return> {
 
-    override fun aggregate(column: DataColumn<C?>): R? =
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     * Uses [aggregator] to compute the result.
+     */
+    override fun aggregate(values: Iterable<Value>, type: KType): Return? = aggregator(values, type)
+
+    /**
+     * Aggregates the data in the given column and computes a single resulting value.
+     * Nulls are filtered out before calling the aggregation function with [Iterable] and [KType].
+     */
+    override fun aggregate(column: DataColumn<Value?>): Return? =
         if (column.hasNulls()) {
-            aggregate(column.asSequence().filterNotNull().asIterable(), column.type())
+            aggregate(column.asSequence().filterNotNull().asIterable(), column.type().withNullability(false))
+        } else {
+            aggregate(column.asIterable() as Iterable<Value>, column.type().withNullability(false))
+        }
+
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     */
+    override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return? {
+        val commonType = if (valueTypes != null) {
+            valueTypes.commonType(false)
         } else {
-            aggregate(column.asIterable() as Iterable<C>, column.type())
+            var hasNulls = false
+            val classes = values.mapNotNull {
+                if (it == null) {
+                    hasNulls = true
+                    null
+                } else {
+                    it.javaClass.kotlin
+                }
+            }
+            classes.commonType(hasNulls)
         }
+        return aggregate(values, commonType)
+    }
 
-    override fun aggregate(values: Iterable<C>, type: KType): R? = aggregator(values, type)
+    /**
+     * 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?
 }

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

@@ -1,33 +1,72 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
-import kotlin.reflect.KProperty
-
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require a single parameter.
+ *
+ * Aggregators are cached by their parameter value.
+ * @see AggregatorOptionSwitch2
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch<P, C, R>(val name: String, val getAggregator: (P) -> AggregatorProvider<C, R>) {
+internal class AggregatorOptionSwitch1<Param1, AggregatorType : Aggregator<*, *>>(
+    val name: String,
+    val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
+) {
 
-    private val cache = mutableMapOf<P, Aggregator<C, R>>()
+    private val cache: MutableMap<Param1, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option: P) = cache.getOrPut(option) { getAggregator(option).create(name) }
+    operator fun invoke(param1: Param1): AggregatorType =
+        cache.getOrPut(param1) {
+            getAggregator(param1).create(name)
+        }
 
-    class Factory<P, C, R>(val getAggregator: (P) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch1].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch1.Factory { param1: Param1 ->
+     *   MyAggregator.Factory(param1)
+     * }
+     */
+    class Factory<Param1, AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch1(name, getAggregator)
+        })
 }
 
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require two parameters.
+ *
+ * Aggregators are cached by their parameter values.
+ * @see AggregatorOptionSwitch1
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch2<P1, P2, C, R>(
+internal class AggregatorOptionSwitch2<Param1, Param2, AggregatorType : Aggregator<*, *>>(
     val name: String,
-    val getAggregator: (P1, P2) -> AggregatorProvider<C, R>,
+    val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
 ) {
 
-    private val cache = mutableMapOf<Pair<P1, P2>, Aggregator<C, R>>()
+    private val cache: MutableMap<Pair<Param1, Param2>, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option1: P1, option2: P2) =
-        cache.getOrPut(option1 to option2) {
-            getAggregator(option1, option2).create(name)
+    operator fun invoke(param1: Param1, param2: Param2): AggregatorType =
+        cache.getOrPut(param1 to param2) {
+            getAggregator(param1, param2).create(name)
         }
 
-    class Factory<P1, P2, C, R>(val getAggregator: (P1, P2) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch2(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch2].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch2.Factory { param1: Param1, param2: Param2 ->
+     *   MyAggregator.Factory(param1, param2)
+     * }
+     * ```
+     */
+    class Factory<Param1, Param2, AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch2(name, getAggregator)
+        })
 }

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

@@ -2,9 +2,27 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import kotlin.reflect.KProperty
 
-internal interface AggregatorProvider<C, R> {
+/**
+ * Common interface for providers or "factory" objects that create anything of type [T].
+ *
+ * When implemented, this allows the object to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myNamedValue by MyFactory
+ * ```
+ */
+internal fun interface Provider<T> {
 
-    operator fun getValue(obj: Any?, property: KProperty<*>): Aggregator<C, R> = create(property.name)
-
-    fun create(name: String): Aggregator<C, R>
+    fun create(name: String): T
 }
+
+internal operator fun <T> Provider<T>.getValue(obj: Any?, property: KProperty<*>): T = create(property.name)
+
+/**
+ * Common interface for providers of [Aggregators][Aggregator] or "factory" objects that create aggregators.
+ *
+ * When implemented, this allows an aggregator to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myAggregator by MyAggregator.Factory
+ * ```
+ */
+internal fun interface AggregatorProvider<AggregatorType : Aggregator<*, *>> : Provider<AggregatorType>