aggregators now always filter out nulls, not just for columns

Author: Jolanrensen

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,18 @@ 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, 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].
      */
@@ -57,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.
@@ -92,8 +92,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

@@ -14,7 +14,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 +25,19 @@ 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, 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.asSequence().filterNotNull().asIterable(), // TODO make dependant on type's nullability
+            type.withNullability(false),
+        )
 
     /**
      * Function that can give the return type of [aggregate] as [KType], given the type of the input.
@@ -44,7 +50,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,17 +60,12 @@ 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 {
+    override fun aggregateCalculatingType(values: Iterable<Value?>, valueTypes: Set<KType>?): Return {
         val commonType = if (valueTypes != null) {
             valueTypes.commonType(false)
         } else {

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

@@ -24,7 +24,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:
@@ -113,11 +113,10 @@ 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?"
         }
-
         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
@@ -147,8 +146,8 @@ 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 {
-        val valueTypes = valueTypes ?: values.types()
+    override fun aggregateCalculatingType(values: Iterable<Number?>, valueTypes: Set<KType>?): Return {
+        val valueTypes = valueTypes ?: values.filterNotNull().types()
         val commonType = valueTypes
             .unifiedNumberType(PRIMITIVES_ONLY)
             .withNullability(false)

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(