Merge branch 'mean' into sum

Author: Jolanrensen

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.primitiveNumberColumns
 import org.jetbrains.kotlinx.dataframe.impl.columns.toNumberColumns
 import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
 import kotlin.reflect.KProperty
+import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
 /*
@@ -38,7 +39,7 @@ import kotlin.reflect.typeOf
 // region DataColumn
 
 public fun <T : Number> DataColumn<T?>.mean(skipNA: Boolean = skipNA_default): Double =
-    Aggregators.mean(skipNA).aggregate(this)!!
+    Aggregators.mean(skipNA).aggregate(this)
 
 public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
     skipNA: Boolean = skipNA_default,
@@ -47,7 +48,6 @@ public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
     Aggregators.mean(skipNA)
         .cast2<R?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
 // endregion
 
@@ -56,15 +56,14 @@ public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
 public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Double =
     Aggregators.mean(skipNA).aggregateOfRow(this) {
         colsOf<Number?> { it.isPrimitiveNumber() }
-    } ?: Double.NaN
+    }
 
-public inline fun <reified T : Number> AnyRow.rowMeanOf(skipNA: Boolean = skipNA_default): Double {
-    require(typeOf<T>() in primitiveNumberTypes) {
+public inline fun <reified T : Number?> AnyRow.rowMeanOf(skipNA: Boolean = skipNA_default): Double {
+    require(typeOf<T>().withNullability(false) in primitiveNumberTypes) {
         "Type ${T::class.simpleName} is not a primitive number type. Mean only supports primitive number types."
     }
     return Aggregators.mean(skipNA)
         .aggregateOfRow(this) { colsOf<T>() }
-        ?: Double.NaN
 }
 
 // endregion
@@ -97,7 +96,7 @@ public fun <T, C : Number> DataFrame<T>.meanFor(
 public fun <T, C : Number> DataFrame<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): Double = Aggregators.mean(skipNA).aggregateAll(this, columns) ?: Double.NaN
+): Double = Aggregators.mean(skipNA).aggregateAll(this, columns)
 
 public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Double =
     mean(skipNA) { columns.toNumberColumns() }
@@ -115,7 +114,7 @@ public fun <T, C : Number> DataFrame<T>.mean(vararg columns: KProperty<C?>, skip
 public inline fun <T, reified D : Number> DataFrame<T>.meanOf(
     skipNA: Boolean = skipNA_default,
     noinline expression: RowExpression<T, D?>,
-): Double = Aggregators.mean(skipNA).of(this, expression) ?: Double.NaN
+): Double = Aggregators.mean(skipNA).of(this, expression)
 
 // endregion
 

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

@@ -15,8 +15,7 @@ import kotlin.reflect.full.withNullability
  * @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]`?`.
+ * @param Return The type of the resulting value. Can optionally be nullable.
  */
 @PublishedApi
 internal interface Aggregator<in Value, out Return> {
@@ -33,20 +32,20 @@ internal interface Aggregator<in Value, out Return> {
      *
      * When the exact [type] is unknown, use [aggregateCalculatingType].
      */
-    fun aggregate(values: Iterable<Value>, type: KType): Return?
+    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 [aggregate] (with [Iterable] and [KType]) is called.
      *
      * See [AggregatorBase.aggregate].
      */
-    fun aggregate(column: DataColumn<Value?>): Return?
+    fun aggregate(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 aggregate(columns: Iterable<DataColumn<Value?>>): Return
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
@@ -58,7 +57,7 @@ internal interface Aggregator<in Value, out Return> {
      *   It should contain all types of [values].
      *   If `null`, the types of [values] will be calculated at runtime (heavy!).
      */
-    fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>? = null): Return?
+    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.
@@ -82,17 +81,19 @@ internal interface Aggregator<in Value, out Return> {
     fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
 }
 
+@Suppress("UNCHECKED_CAST")
 @PublishedApi
 internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
 
+@Suppress("UNCHECKED_CAST")
 @PublishedApi
 internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>
 
 /** Type alias for [Aggregator.calculateReturnTypeOrNull] */
 internal typealias CalculateReturnTypeOrNull = (type: KType, emptyInput: Boolean) -> KType?
 
 /** Type alias for [Aggregator.aggregate]. */
-internal typealias Aggregate<Value, Return> = Iterable<Value>.(type: KType) -> Return?
+internal typealias Aggregate<Value, Return> = Iterable<Value>.(type: KType) -> Return
 
 /** Common case for [CalculateReturnTypeOrNull], preserves return type, but makes it nullable for empty inputs. */
 internal val preserveReturnTypeNullIfEmpty: CalculateReturnTypeOrNull = { type, emptyInput ->

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

@@ -31,7 +31,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
      *
      * When the exact [type] is unknown, use [aggregateCalculatingType].
      */
-    override fun aggregate(values: Iterable<Value>, type: KType): Return? = aggregator(values, type)
+    override fun aggregate(values: Iterable<Value>, type: KType): Return = aggregator(values, type)
 
     /**
      * Function that can give the return type of [aggregate] as [KType], given the type of the input.
@@ -52,7 +52,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * Nulls are filtered out by default, then [aggregate] (with [Iterable] and [KType]) is called.
      */
     @Suppress("UNCHECKED_CAST")
-    override fun aggregate(column: DataColumn<Value?>): Return? =
+    override fun aggregate(column: DataColumn<Value?>): Return =
         aggregate(
             values =
                 if (column.hasNulls()) {
@@ -71,7 +71,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
      *   If provided, this can be used to avoid calculating the types of [values][org.jetbrains.kotlinx.dataframe.values] at runtime with reflection.
      *   It should contain all types of [values][org.jetbrains.kotlinx.dataframe.values].
      * If `null`, the types of [values][org.jetbrains.kotlinx.dataframe.values] will be calculated at runtime (heavy!). */
-    override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return? {
+    override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return {
         val commonType = if (valueTypes != null) {
             valueTypes.commonType(false)
         } else {
@@ -93,7 +93,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * 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?
+    abstract override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return
 
     /**
      * Function that can give the return type of [aggregate] with columns as [KType],

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt

@@ -30,7 +30,7 @@ internal object Aggregators {
      *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
      *     -> Iterable<Return> // nulls filtered out
      *     -> stepTwoAggregator(Iterable<Return>, common valueType)
-     *     -> Return?
+     *     -> Return
      * ```
      *
      * It can also be used as a "simple" aggregator by providing the same function for both steps.
@@ -43,7 +43,7 @@ internal object Aggregators {
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
      */
-    private fun <Type> twoStepPreservingType(aggregator: Aggregate<Type, Type>) =
+    private fun <Type> twoStepPreservingType(aggregator: Aggregate<Type, Type?>) =
         TwoStepAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             stepOneAggregator = aggregator,
@@ -68,7 +68,7 @@ internal object Aggregators {
      *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
      *     -> Iterable<Return> // nulls filtered out
      *     -> stepTwoAggregator(Iterable<Return>, common valueType)
-     *     -> Return?
+     *     -> Return
      * ```
      *
      * It can also be used as a "simple" aggregator by providing the same function for both steps.
@@ -106,7 +106,7 @@ internal object Aggregators {
      * Iterable<Column<Value?>>
      *     -> Iterable<Value> // flattened without nulls
      *     -> aggregator(Iterable<Value>, common colType)
-     *     -> Return?
+     *     -> Return
      * ```
      *
      * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
@@ -119,7 +119,7 @@ internal object Aggregators {
      * @param aggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator.aggregate] function.
      *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator.aggregate].
      */
-    private fun <Type> flatteningPreservingTypes(aggregate: Aggregate<Type, Type>) =
+    private fun <Type> flatteningPreservingTypes(aggregate: Aggregate<Type, Type?>) =
         FlatteningAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             aggregator = aggregate,
@@ -140,7 +140,7 @@ internal object Aggregators {
      * Iterable<Column<Value?>>
      *     -> Iterable<Value> // flattened without nulls
      *     -> aggregator(Iterable<Value>, common colType)
-     *     -> Return?
+     *     -> Return
      * ```
      *
      * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
@@ -182,7 +182,7 @@ internal object Aggregators {
      *     -> aggregator(Iterable<specific Number>, unified number type of common colType) // called on each iterable
      *     -> Iterable<Return> // nulls filtered out
      *     -> aggregator(Iterable<specific Return>, unified number type of common valueType)
-     *     -> Return?
+     *     -> Return
      * ```
      *
      * @param name The name of this aggregator.
@@ -191,7 +191,7 @@ internal object Aggregators {
      *   While it takes a [Number] argument, you can assume that all values are of the same specific type, however,
      *   this type can be different for different calls to [aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.aggregator].
      */
-    private fun <Return : Number> twoStepForNumbers(
+    private fun <Return : Number?> twoStepForNumbers(
         getReturnTypeOrNull: CalculateReturnTypeOrNull,
         aggregate: Aggregate<Number, Return>,
     ) = TwoStepNumbersAggregator.Factory(

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt

@@ -18,7 +18,7 @@ import kotlin.reflect.full.withNullability
  * Iterable<Column<Value?>>
  *     -> Iterable<Value> // flattened without nulls
  *     -> aggregator(Iterable<Value>, common colType)
- *     -> Return?
+ *     -> Return
  * ```
  *
  * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
@@ -42,7 +42,7 @@ internal class FlatteningAggregator<in Value, out Return>(
      * The columns are flattened into a single list of values, filtering nulls as usual;
      * then the aggregation function is with the common type of the columns.
      */
-    override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return? {
+    override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return {
         val commonType = columns.map { it.type() }.commonType().withNullability(false)
         val allValues = columns.asSequence().flatMap { it.values() }.filterNotNull()
         return aggregate(allValues.asIterable(), commonType)

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt

@@ -23,7 +23,7 @@ import kotlin.reflect.full.withNullability
  *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
  *     -> Iterable<Return> // nulls filtered out
  *     -> stepTwoAggregator(Iterable<Return>, common valueType)
- *     -> Return?
+ *     -> Return
  * ```
  *
  * It can also be used as a "simple" aggregator by providing the same function for both steps.
@@ -40,7 +40,7 @@ internal class TwoStepAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     stepOneAggregator: Aggregate<Value, Return>,
-    private val stepTwoAggregator: Aggregate<Return, Return>,
+    private val stepTwoAggregator: Aggregate<Return & Any, Return>,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
     /**
@@ -50,7 +50,7 @@ internal class TwoStepAggregator<in Value, out Return>(
      *
      * Post-step-one types are calculated by [calculateReturnTypeOrNull].
      */
-    override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return? {
+    override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return {
         val (values, types) = columns.mapNotNull { col ->
             // uses stepOneAggregator
             val value = aggregate(col) ?: return@mapNotNull null
@@ -93,7 +93,7 @@ internal class TwoStepAggregator<in Value, out Return>(
     class Factory<in Value, out Return>(
         private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
         private val stepOneAggregator: Aggregate<Value, Return>,
-        private val stepTwoAggregator: Aggregate<Return, Return>,
+        private val stepTwoAggregator: Aggregate<Return & Any, Return>,
     ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
             TwoStepAggregator(
                 name = name,

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt

@@ -39,7 +39,7 @@ private val logger = KotlinLogging.logger { }
  *     -> aggregator(Iterable<specific Number>, unified number type of common colType) // called on each iterable
  *     -> Iterable<Return> // nulls filtered out
  *     -> aggregator(Iterable<specific Return>, unified number type of common valueType)
- *     -> Return?
+ *     -> Return
  * ```
  *
  * @param name The name of this aggregator.
@@ -48,7 +48,7 @@ private val logger = KotlinLogging.logger { }
  *   While it takes a [Number] argument, you can assume that all values are of the same specific type, however,
  *   this type can be different for different calls to [aggregator].
  */
-internal class TwoStepNumbersAggregator<out Return : Number>(
+internal class TwoStepNumbersAggregator<out Return : Number?>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     aggregator: Aggregate<Number, Return>,
@@ -62,7 +62,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      * After the first aggregation, the number types are found by [calculateReturnTypeOrNull] and then
      * unified using [aggregateCalculatingType].
      */
-    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return? {
+    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return {
         val (values, types) = columns.mapNotNull { col ->
             val value = aggregate(col) ?: return@mapNotNull null
             val type = calculateReturnTypeOrNull(
@@ -113,7 +113,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      *
      * When the exact [type] is unknown, use [aggregateCalculatingType].
      */
-    override fun aggregate(values: Iterable<Number>, type: KType): Return? {
+    override fun aggregate(values: Iterable<Number>, type: KType): Return {
         require(type.isSubtypeOf(typeOf<Number?>())) {
             "${TwoStepNumbersAggregator::class.simpleName}: Type $type is not a subtype of Number?"
         }
@@ -147,7 +147,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      *   If `null`, the types of [values] will be calculated at runtime (heavy!).
      */
     @Suppress("UNCHECKED_CAST")
-    override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
+    override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return {
         val valueTypes = valueTypes ?: values.types()
         val commonType = valueTypes
             .unifiedNumberType(PRIMITIVES_ONLY)
@@ -176,7 +176,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
      * @param aggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
      */
-    class Factory<out Return : Number>(
+    class Factory<out Return : Number?>(
         private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
         private val aggregate: Aggregate<Number, Return>,
     ) : AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt

@@ -19,19 +19,19 @@ import kotlin.reflect.typeOf
 internal inline fun <C, reified V, R> Aggregator<V, R>.aggregateOf(
     values: Iterable<C>,
     noinline transform: (C) -> V,
-): R? = aggregate(values.asSequence().map(transform).asIterable(), typeOf<V>())
+): R = aggregate(values.asSequence().map(transform).asIterable(), typeOf<V>())
 
 @PublishedApi
 internal inline fun <C, reified V, R> Aggregator<V, R>.aggregateOf(
     column: DataColumn<C>,
     noinline transform: (C) -> V,
-): R? = aggregateOf(column.values(), transform)
+): R = aggregateOf(column.values(), transform)
 
 @PublishedApi
 internal inline fun <T, reified C, R> Aggregator<*, R>.aggregateOf(
     frame: DataFrame<T>,
     crossinline expression: RowExpression<T, C>,
-): R? = (this as Aggregator<C, R>).aggregateOf(frame.rows()) { expression(it, it) }
+): R = (this as Aggregator<C, R>).aggregateOf(frame.rows()) { expression(it, it) }
 
 @PublishedApi
 internal fun <T, C, R> Aggregator<*, R>.aggregateOfDelegated(
@@ -47,10 +47,10 @@ internal fun <T, C, R> Aggregator<*, R>.aggregateOfDelegated(
 internal inline fun <T, reified C, R> Aggregator<*, R>.of(
     data: DataFrame<T>,
     crossinline expression: RowExpression<T, C>,
-): R? = aggregateOf(data as DataFrame<T>, expression)
+): R = aggregateOf(data as DataFrame<T>, expression)
 
 @PublishedApi
-internal inline fun <C, reified V, R> Aggregator<V, R>.of(data: DataColumn<C>, crossinline expression: (C) -> V): R? =
+internal inline fun <C, reified V, R> Aggregator<V, R>.of(data: DataColumn<C>, crossinline expression: (C) -> V): R =
     aggregateOf(data.values()) { expression(it) }
 
 @PublishedApi

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/row.kt

@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
  * @param columns selector of which columns inside the [row] to aggregate
  */
 @PublishedApi
-internal fun <V, R> Aggregator<V, R>.aggregateOfRow(row: AnyRow, columns: ColumnsSelector<*, V?>): R? {
+internal fun <V, R> Aggregator<V, R>.aggregateOfRow(row: AnyRow, columns: ColumnsSelector<*, V?>): R {
     val filteredColumns = row.df().getColumns(columns)
     return aggregateCalculatingType(
         values = filteredColumns.mapNotNull { row[it] },