Fixed AggregatorBase only filtering nulls when the type says they exist.

unifying numbers can now handle null/nothing in the input.

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

@@ -26,7 +26,8 @@ internal interface Aggregator<in Value, out Return> {
      * Base function of [Aggregator].
      *
      * Aggregates the given values, taking [type] into account,
-     * filtering nulls, and computes a single resulting value.
+     * 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.
      *
@@ -55,7 +56,7 @@ 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
 

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
 
@@ -26,7 +27,8 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * Base function of [Aggregator].
      *
      * Aggregates the given values, taking [type] into account,
-     * filtering nulls, and computes a single resulting value.
+     * 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.
      *
@@ -35,7 +37,13 @@ internal abstract class AggregatorBase<in Value, out Return>(
     @Suppress("UNCHECKED_CAST")
     override fun aggregate(values: Iterable<Value?>, type: KType): Return =
         aggregator(
-            values.asSequence().filterNotNull().asIterable(), // TODO make dependant on type's nullability
+            // values =
+            if (type.isMarkedNullable) {
+                values.asSequence().filterNotNull().asIterable()
+            } else {
+                values as Iterable<Value & Any>
+            },
+            // type =
             type.withNullability(false),
         )
 
@@ -66,7 +74,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
 
     /** @include [Aggregator.aggregateCalculatingType] */
     override fun aggregateCalculatingType(values: Iterable<Value?>, valueTypes: Set<KType>?): Return {
-        val commonType = if (valueTypes != null) {
+        val commonType = if (valueTypes != null && valueTypes.isNotEmpty()) {
             valueTypes.commonType(false)
         } else {
             var hasNulls = false
@@ -78,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/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
@@ -106,6 +107,8 @@ 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.
@@ -143,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.filterNotNull().types()
-        val commonType = valueTypes
-            .unifiedNumberType(PRIMITIVES_ONLY)
-            .withNullability(false)
+        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.",
             )