Merge branch 'mean' into sum

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt

@@ -21,6 +21,8 @@ import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions
  * potentially losing a little precision, but a warning will be given.
  *
  * See [UnifiedNumberTypeOptions] for these settings.
+ *
+ * At the bottom of the graph is [Nothing]. This can be interpreted as `null`.
  */
 internal interface UnifyingNumbers {
 
@@ -40,6 +42,9 @@ internal interface UnifyingNumbers {
      *      |    /   |
      *      |   /    |
      *    UByte     Byte
+     *        \\     /
+     *        \\    /
+     *       Nothing?
      * ```
      */
     interface Graph

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

@@ -82,6 +82,9 @@ internal fun getUnifiedNumberTypeGraph(
 
             addEdge(typeOf<Short>(), typeOf<UByte>())
             addEdge(typeOf<Short>(), typeOf<Byte>())
+
+            addEdge(typeOf<UByte>(), nothingType)
+            addEdge(typeOf<Byte>(), nothingType)
         }
     }
 
@@ -121,7 +124,11 @@ internal fun getUnifiedNumberType(
             ?: error("Can not find common number type for $first and $second")
     }
 
-    return if (first.isMarkedNullable || second.isMarkedNullable) result.withNullability(true) else result
+    return if (first.isMarkedNullable || second.isMarkedNullable) {
+        result.withNullability(true)
+    } else {
+        result
+    }
 }
 
 /** @include [getUnifiedNumberType] */
@@ -184,7 +191,7 @@ internal fun Iterable<Number?>.convertToUnifiedNumberType(
     options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
     commonNumberType: KType? = null,
 ): Iterable<Number?> {
-    val commonNumberType = commonNumberType ?: this.filterNotNull().types().unifiedNumberType(options)
+    val commonNumberType = commonNumberType ?: this.types().unifiedNumberType(options)
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
         if (it == null) return@map null
@@ -209,7 +216,7 @@ internal fun Sequence<Number?>.convertToUnifiedNumberType(
     options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
     commonNumberType: KType? = null,
 ): Sequence<Number?> {
-    val commonNumberType = commonNumberType ?: this.filterNotNull().asIterable().types().unifiedNumberType(options)
+    val commonNumberType = commonNumberType ?: this.asIterable().types().unifiedNumberType(options)
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
         if (it == null) return@map null

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

@@ -667,19 +667,26 @@ internal fun Any.isBigNumber(): Boolean = this is BigInteger || this is BigDecim
  *
  * The [KClass] is determined by retrieving the runtime class of each element.
  *
+ * [Nothing::class][Nothing] is used for elements that are `null`.
+ *
  * @return A set of [KClass] objects representing the runtime types of elements in the iterable.
  */
-internal fun Iterable<Any>.classes(): Set<KClass<*>> = mapTo(mutableSetOf()) { it::class }
+internal fun Iterable<Any?>.classes(): Set<KClass<*>> =
+    mapTo(mutableSetOf()) {
+        if (it == null) Nothing::class else it::class
+    }
 
 /**
  * Returns a set of [KType] objects representing the star-projected types of the runtime classes
  * of all unique elements in the iterable.
  *
- * The method internally relies on the [classes] function to collect the runtime classes of the
- * elements in the iterable and then maps each class to its star-projected type.
- *
  * This can be a heavy operation!
  *
+ * [typeOf<Nothing?>()][nullableNothingType] is used for elements that are `null`.
+ *
  * @return A set of [KType] objects corresponding to the star-projected runtime types of elements in the iterable.
  */
-internal fun Iterable<Any>.types(): Set<KType> = classes().mapTo(mutableSetOf()) { it.createStarProjectedType(false) }
+internal fun Iterable<Any?>.types(): Set<KType> =
+    mapTo(mutableSetOf()) {
+        if (it == null) nullableNothingType else it::class.createStarProjectedType(false)
+    }

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

@@ -13,8 +13,7 @@ import kotlin.reflect.full.withNullability
  * 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.
+ *   The input can always have nulls, they are filtered out.
  * @param Return The type of the resulting value. Can optionally be nullable.
  */
 @PublishedApi
@@ -26,17 +25,19 @@ internal interface Aggregator<in Value, out Return> {
     /**
      * Base function of [Aggregator].
      *
-     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     * Aggregates the given values, taking [type] 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 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.
+     * Calls [aggregate] (with [Iterable] and [KType]).
      *
      * See [AggregatorBase.aggregate].
      */
@@ -55,9 +56,9 @@ internal interface Aggregator<in Value, out Return> {
      * @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`, the types of [values] will be calculated at runtime (heavy!).
+     *   If `null` or empty, 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.
@@ -92,8 +93,11 @@ internal fun <Value, Return> Aggregator<*, *>.cast2(): 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
+/**
+ * Type alias for the argument for [Aggregator.aggregate].
+ * Nulls have already been filtered out when this argument is called.
+ */
+internal typealias Aggregate<Value, Return> = Iterable<Value & Any>.(type: KType) -> Return
 
 /** Common case for [CalculateReturnTypeOrNull], preserves return type, but makes it nullable for empty inputs. */
 internal val preserveReturnTypeNullIfEmpty: CalculateReturnTypeOrNull = { type, emptyInput ->

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

@@ -4,6 +4,7 @@ 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
 
@@ -14,7 +15,7 @@ import kotlin.reflect.full.withNullability
  * or multiple [DataColumns][DataColumn].
  *
  * @param name The name of this aggregator.
- * @param aggregator Functional argument for the [aggregate] function.
+ * @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,
@@ -25,13 +26,26 @@ internal abstract class AggregatorBase<in Value, out Return>(
     /**
      * Base function of [Aggregator].
      *
-     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     * Aggregates the given values, taking [type] into account,
+     * filtering nulls (only if [type.isMarkedNullable][KType.isMarkedNullable]),
+     * and computes a single resulting value.
      *
-     * Uses [aggregator] to compute the result.
+     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregator] argument.
      *
      * When the exact [type] is unknown, use [aggregateCalculatingType].
      */
-    override fun aggregate(values: Iterable<Value>, type: KType): Return = aggregator(values, type)
+    @Suppress("UNCHECKED_CAST")
+    override fun aggregate(values: Iterable<Value?>, type: KType): Return =
+        aggregator(
+            // values =
+            if (type.isMarkedNullable) {
+                values.asSequence().filterNotNull().asIterable()
+            } else {
+                values as Iterable<Value & Any>
+            },
+            // type =
+            type.withNullability(false),
+        )
 
     /**
      * Function that can give the return type of [aggregate] as [KType], given the type of the input.
@@ -44,7 +58,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * @return The return type of [aggregate] as [KType].
      */
     override fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType? =
-        getReturnTypeOrNull(type, emptyInput)
+        getReturnTypeOrNull(type.withNullability(false), emptyInput)
 
     /**
      * Aggregates the data in the given column and computes a single resulting value.
@@ -54,18 +68,13 @@ internal abstract class AggregatorBase<in Value, out Return>(
     @Suppress("UNCHECKED_CAST")
     override fun aggregate(column: DataColumn<Value?>): Return =
         aggregate(
-            values =
-                if (column.hasNulls()) {
-                    column.asSequence().filterNotNull().asIterable()
-                } else {
-                    column.asIterable() as Iterable<Value>
-                },
-            type = column.type().withNullability(false),
+            values = column.asIterable(),
+            type = column.type(),
         )
 
     /** @include [Aggregator.aggregateCalculatingType] */
-    override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return {
-        val commonType = if (valueTypes != null) {
+    override fun aggregateCalculatingType(values: Iterable<Value?>, valueTypes: Set<KType>?): Return {
+        val commonType = if (valueTypes != null && valueTypes.isNotEmpty()) {
             valueTypes.commonType(false)
         } else {
             var hasNulls = false
@@ -77,7 +86,11 @@ internal abstract class AggregatorBase<in Value, out Return>(
                     it.javaClass.kotlin
                 }
             }
-            classes.commonType(hasNulls)
+            if (classes.isEmpty()) {
+                nothingType(hasNulls)
+            } else {
+                classes.commonType(hasNulls)
+            }
         }
         return aggregate(values, commonType)
     }

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

@@ -8,7 +8,7 @@ import kotlin.reflect.full.withNullability
 /**
  * Simple [Aggregator] implementation with flattening behavior for multiple columns.
  *
- * Nulls are filtered from columns.
+ * Nulls are filtered out.
  *
  * When called on multiple columns,
  * the columns are flattened into a single list of values, filtering nulls as usual;

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

@@ -10,7 +10,7 @@ import kotlin.reflect.full.withNullability
 /**
  * A slightly more advanced [Aggregator] implementation.
  *
- * Nulls are filtered from columns.
+ * Nulls are filtered out.
  *
  * When called on multiple columns, this [Aggregator] works in two steps:
  * First, it aggregates within a [DataColumn]/[Iterable] ([stepOneAggregator]) with their (given) type,

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

@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
 import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY
 import org.jetbrains.kotlinx.dataframe.impl.anyNull
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+import org.jetbrains.kotlinx.dataframe.impl.isNothing
 import org.jetbrains.kotlinx.dataframe.impl.nothingType
 import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
 import org.jetbrains.kotlinx.dataframe.impl.renderType
@@ -24,7 +25,7 @@ private val logger = KotlinLogging.logger { }
  * [Aggregator] made specifically for number calculations.
  * Mixed number types are [unified][UnifyingNumbers] to [primitives][PRIMITIVES_ONLY].
  *
- * Nulls are filtered from columns.
+ * Nulls are filtered out.
  *
  * When called on multiple columns (with potentially mixed [Number] types),
  * this [Aggregator] works in two steps:
@@ -106,18 +107,19 @@ internal class TwoStepNumbersAggregator<out Return : Number?>(
      *
      * Aggregates the given values, taking [type] into account, and computes a single resulting value.
      *
+     * Nulls are filtered out (only if [type.isMarkedNullable][KType.isMarkedNullable]).
+     *
      * Uses [aggregator] to compute the result.
      *
      * This function is modified to call [aggregateCalculatingType] when it encounters mixed number types.
      * This is not optimal and should be avoided by calling [aggregateCalculatingType] with known number types directly.
      *
      * 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?"
         }
-
         return when (type.withNullability(false)) {
             // If the type is not a specific number, but rather a mixed Number, we unify the types first.
             // This is heavy and could be avoided by calling aggregate with a specific number type
@@ -144,21 +146,21 @@ internal class TwoStepNumbersAggregator<out Return : Number?>(
      * @param valueTypes The types of the numbers.
      *   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`, the types of [values] will be calculated at runtime (heavy!).
+     *   If `null` or empty, the types of [values] will be calculated at runtime (heavy!).
      */
     @Suppress("UNCHECKED_CAST")
-    override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return {
-        val valueTypes = valueTypes ?: values.types()
-        val commonType = valueTypes
-            .unifiedNumberType(PRIMITIVES_ONLY)
-            .withNullability(false)
+    override fun aggregateCalculatingType(values: Iterable<Number?>, valueTypes: Set<KType>?): Return {
+        val valueTypes = valueTypes?.takeUnless { it.isEmpty() } ?: values.types()
+        val commonType = valueTypes.unifiedNumberType(PRIMITIVES_ONLY)
 
-        if (commonType == typeOf<Double>() && (typeOf<ULong>() in valueTypes || typeOf<Long>() in valueTypes)) {
+        if (commonType.isSubtypeOf(typeOf<Double?>()) &&
+            (typeOf<ULong>() in valueTypes || typeOf<Long>() in valueTypes)
+        ) {
             logger.warn {
                 "Number unification of Long -> Double happened during aggregation. Loss of precision may have occurred."
             }
         }
-        if (commonType !in primitiveNumberTypes && commonType != nothingType) {
+        if (commonType.withNullability(false) !in primitiveNumberTypes && !commonType.isNothing) {
             throw IllegalArgumentException(
                 "Cannot calculate $name of ${renderType(commonType)}, only primitive numbers are supported.",
             )

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

@@ -13,13 +13,18 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.internal
 import org.jetbrains.kotlinx.dataframe.impl.emptyPath
+import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
 @PublishedApi
 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 = values.asSequence().mapNotNull(transform).asIterable(),
+        type = typeOf<V>().withNullability(false),
+    )
 
 @PublishedApi
 internal inline fun <C, reified V, R> Aggregator<V, R>.aggregateOf(