added calculateReturnTypeOrNull system to aggregators to avoid runtime value instance checks where we know the types already

Jolanrensen · Jolanrensen · commit 268d2387e00a · 2025-03-03T16:53:42.000+01:00
diff --git a/core/api/core.api b/core/api/core.api
@@ -10082,6 +10082,7 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/impl/aggregation
 	public abstract fun aggregate (Ljava/lang/Iterable;Lkotlin/reflect/KType;)Ljava/lang/Object;
 	public abstract fun aggregate (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Object;
 	public abstract fun aggregateCalculatingType (Ljava/lang/Iterable;Ljava/util/Set;)Ljava/lang/Object;
+	public abstract fun calculateReturnTypeOrNull (Lkotlin/reflect/KType;Z)Lkotlin/reflect/KType;
 	public abstract fun getName ()Ljava/lang/String;
 	public abstract fun getPreservesType ()Z
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
 /**
  * Base interface for all aggregators.
@@ -56,10 +57,19 @@ internal interface Aggregator<Value, Return> {
      * 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?
+
+    /**
+     * Function that can give the return type of [aggregate] as [KType], given the type of the input.
+     */
+    fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
 }
 
 @PublishedApi
 internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
 
 @PublishedApi
 internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>
+
+internal val preserveReturnTypeNullIfEmpty: (KType, Boolean) -> KType = { type, emptyInput ->
+    type.withNullability(emptyInput)
+}
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
@@ -18,6 +18,7 @@ import kotlin.reflect.full.withNullability
  */
 internal abstract class AggregatorBase<Value, Return>(
     override val name: String,
+    protected val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     protected val aggregator: (values: Iterable<Value>, type: KType) -> Return?,
 ) : Aggregator<Value, Return> {
 
@@ -29,6 +30,9 @@ internal abstract class AggregatorBase<Value, Return>(
      */
     override fun aggregate(values: Iterable<Value>, type: KType): Return? = aggregator(values, type)
 
+    override fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType? =
+        getReturnTypeOrNull(type, emptyInput)
+
     /**
      * 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].
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
@@ -1,12 +1,15 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.math.meanOrNull
+import org.jetbrains.kotlinx.dataframe.math.meanTypeResultOrNull
 import org.jetbrains.kotlinx.dataframe.math.median
 import org.jetbrains.kotlinx.dataframe.math.percentile
 import org.jetbrains.kotlinx.dataframe.math.std
 import org.jetbrains.kotlinx.dataframe.math.sum
 import java.math.BigDecimal
 import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
 @PublishedApi
 internal object Aggregators {
@@ -18,6 +21,7 @@ internal object Aggregators {
      */
     private fun <Type> twoStepPreservingType(aggregator: Iterable<Type>.(type: KType) -> Type?) =
         TwoStepAggregator.Factory(
+            getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             stepOneAggregator = aggregator,
             stepTwoAggregator = aggregator,
             preservesType = true,
@@ -29,9 +33,11 @@ internal object Aggregators {
      * @include [TwoStepAggregator]
      */
     private fun <Value, Return> twoStepChangingType(
+        getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         stepOneAggregator: Iterable<Value>.(type: KType) -> Return,
         stepTwoAggregator: Iterable<Return>.(type: KType) -> Return,
     ) = TwoStepAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
         stepOneAggregator = stepOneAggregator,
         stepTwoAggregator = stepTwoAggregator,
         preservesType = false,
@@ -44,6 +50,7 @@ internal object Aggregators {
      */
     private fun <Type> flatteningPreservingTypes(aggregate: Iterable<Type?>.(type: KType) -> Type?) =
         FlatteningAggregator.Factory(
+            getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             aggregator = aggregate,
             preservesType = true,
         )
@@ -53,19 +60,27 @@ internal object Aggregators {
      *
      * @include [FlatteningAggregator]
      */
-    private fun <Value, Return> flatteningChangingTypes(aggregate: Iterable<Value?>.(type: KType) -> Return?) =
-        FlatteningAggregator.Factory(
-            aggregator = aggregate,
-            preservesType = false,
-        )
+    private fun <Value, Return> flatteningChangingTypes(
+        getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
+        aggregate: Iterable<Value?>.(type: KType) -> Return?,
+    ) = FlatteningAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
+        aggregator = aggregate,
+        preservesType = false,
+    )
 
     /**
      * Factory for a two-step aggregator that works only with numbers.
      *
      * @include [TwoStepNumbersAggregator]
      */
-    private fun <Return : Number> twoStepForNumbers(aggregate: Iterable<Number>.(numberType: KType) -> Return?) =
-        TwoStepNumbersAggregator.Factory(aggregate)
+    private fun <Return : Number> twoStepForNumbers(
+        getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
+        aggregate: Iterable<Number>.(numberType: KType) -> Return?,
+    ) = TwoStepNumbersAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
+        aggregate = aggregate,
+    )
 
     /** @include [AggregatorOptionSwitch1] */
     private fun <Param1, AggregatorType : Aggregator<*, *>> withOneOption(
@@ -82,27 +97,29 @@ internal object Aggregators {
     val max by twoStepPreservingType<Comparable<Any?>> { maxOrNull() }
 
     val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
-        flatteningChangingTypes<Number, Double> { std(it, skipNA, ddof) }
+        flatteningChangingTypes<Number, Double>(
+            getReturnTypeOrNull = { type, emptyInput -> typeOf<Double>().withNullability(emptyInput) },
+        ) { std(it, skipNA, ddof) }
     }
 
     @Suppress("ClassName")
     object mean {
         val toNumber = withOneOption { skipNA: Boolean ->
-            twoStepForNumbers { meanOrNull(it, skipNA) }
+            twoStepForNumbers(::meanTypeResultOrNull) { meanOrNull(it, skipNA) }
         }.create(mean::class.simpleName!!)
 
         val toDouble = withOneOption { skipNA: Boolean ->
-            twoStepForNumbers { meanOrNull(it, skipNA) as Double? }
+            twoStepForNumbers(::meanTypeResultOrNull) { meanOrNull(it, skipNA) as Double? }
         }.create(mean::class.simpleName!!)
 
         val toBigDecimal =
-            twoStepForNumbers {
+            twoStepForNumbers(::meanTypeResultOrNull) {
                 meanOrNull(it) as BigDecimal?
             }.create(mean::class.simpleName!!)
     }
 
     val percentile by withOneOption { percentile: Double ->
-        flatteningChangingTypes<Comparable<Any?>, Comparable<Any?>> { type ->
+        flatteningChangingTypes<Comparable<Any?>, Comparable<Any?>>(preserveReturnTypeNullIfEmpty) { type ->
             percentile(percentile, type)
         }
     }
@@ -111,5 +128,7 @@ internal object Aggregators {
         median(it)
     }
 
-    val sum by twoStepForNumbers { sum(it) }
+    val sum by twoStepForNumbers(
+        getReturnTypeOrNull = { type, _ -> type.withNullability(false) },
+    ) { sum(it) }
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
@@ -32,9 +32,10 @@ import kotlin.reflect.full.withNullability
  */
 internal class FlatteningAggregator<Value, Return>(
     name: String,
+    getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     aggregator: (values: Iterable<Value>, type: KType) -> Return?,
     override val preservesType: Boolean,
-) : AggregatorBase<Value, Return>(name, aggregator) {
+) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, aggregator) {
 
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
@@ -54,9 +55,15 @@ internal class FlatteningAggregator<Value, Return>(
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<Value, Return>(
+        private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         private val aggregator: (Iterable<Value>, KType) -> Return?,
         private val preservesType: Boolean,
     ) : AggregatorProvider<FlatteningAggregator<Value, Return>> by AggregatorProvider({ name ->
-            FlatteningAggregator(name = name, aggregator = aggregator, preservesType = preservesType)
+            FlatteningAggregator(
+                name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
+                aggregator = aggregator,
+                preservesType = preservesType,
+            )
         })
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
@@ -1,9 +1,10 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.impl.classes
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import org.jetbrains.kotlinx.dataframe.size
 import kotlin.reflect.KType
+import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 
 /**
@@ -38,28 +39,30 @@ import kotlin.reflect.full.withNullability
  */
 internal class TwoStepAggregator<Value, Return>(
     name: String,
+    getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     stepOneAggregator: (values: Iterable<Value>, type: KType) -> Return?,
     private val stepTwoAggregator: (values: Iterable<Return>, type: KType) -> Return?,
     override val preservesType: Boolean,
-) : AggregatorBase<Value, Return>(name, stepOneAggregator) {
+) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
      *
      * This function calls [stepOneAggregator] on each column and then [stepTwoAggregator] on the results.
      */
     override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return? {
-        val columnValues = columns.mapNotNull {
+        val (values, types) = columns.mapNotNull { col ->
             // uses stepOneAggregator
-            aggregate(it)
-        }
-        val commonType = if (preservesType) {
-            columns.map { it.type() }.commonType().withNullability(false)
-        } else {
-            // heavy!
-            columnValues.classes().commonType(false)
-        }
-        return stepTwoAggregator(columnValues, commonType)
+            val value = aggregate(col) ?: return@mapNotNull null
+            val type = calculateReturnTypeOrNull(
+                type = col.type().withNullability(false),
+                emptyInput = col.size() == 0,
+            ) ?: value::class.starProjectedType // heavy fallback type calculation
+
+            value to type
+        }.unzip()
+        val commonType = types.commonType()
+        return stepTwoAggregator(values, commonType)
     }
 
     /**
@@ -71,12 +74,14 @@ internal class TwoStepAggregator<Value, Return>(
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<Value, Return>(
+        private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         private val stepOneAggregator: (Iterable<Value>, KType) -> Return?,
         private val stepTwoAggregator: (Iterable<Return>, KType) -> Return?,
         private val preservesType: Boolean,
     ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
             TwoStepAggregator(
                 name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
                 stepOneAggregator = stepOneAggregator,
                 stepTwoAggregator = stepTwoAggregator,
                 preservesType = preservesType,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.impl.types
 import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
 import kotlin.reflect.KType
 import kotlin.reflect.full.isSubtypeOf
+import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
@@ -38,8 +39,9 @@ import kotlin.reflect.typeOf
  */
 internal class TwoStepNumbersAggregator<Return : Number>(
     name: String,
+    getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     aggregator: (values: Iterable<Number>, numberType: KType) -> Return?,
-) : AggregatorBase<Number, Return>(name, aggregator) {
+) : AggregatorBase<Number, Return>(name, getReturnTypeOrNull, aggregator) {
 
     override fun aggregate(values: Iterable<Number>, type: KType): Return? {
         require(type.isSubtypeOf(typeOf<Number?>())) {
@@ -48,11 +50,22 @@ internal class TwoStepNumbersAggregator<Return : Number>(
         return super.aggregate(values, type)
     }
 
-    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return? =
-        aggregateCalculatingType(
-            values = columns.mapNotNull { aggregate(it) },
-            valueTypes = null, // makes the operation heavy
+    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return? {
+        val (values, types) = columns.mapNotNull { col ->
+            val value = aggregate(col) ?: return@mapNotNull null
+            val type = calculateReturnTypeOrNull(
+                type = col.type().withNullability(false),
+                emptyInput = col.size() == 0,
+            ) ?: value::class.starProjectedType // heavy fallback type calculation
+
+            value to type
+        }.unzip()
+
+        return aggregateCalculatingType(
+            values = values,
+            valueTypes = types.toSet(),
         )
+    }
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
@@ -71,8 +84,14 @@ internal class TwoStepNumbersAggregator<Return : Number>(
 
     override val preservesType = false
 
-    class Factory<Return : Number>(private val aggregate: Iterable<Number>.(numberType: KType) -> Return?) :
-        AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->
-            TwoStepNumbersAggregator(name = name, aggregator = aggregate)
+    class Factory<Return : Number>(
+        private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
+        private val aggregate: Iterable<Number>.(numberType: KType) -> Return?,
+    ) : AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->
+            TwoStepNumbersAggregator(
+                name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
+                aggregator = aggregate,
+            )
         })
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -3,6 +3,8 @@ package org.jetbrains.kotlinx.dataframe.math
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
 import org.jetbrains.kotlinx.dataframe.impl.api.toBigDecimal
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+import org.jetbrains.kotlinx.dataframe.impl.nothingType
+import org.jetbrains.kotlinx.dataframe.impl.nullableNothingType
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import org.jetbrains.kotlinx.dataframe.impl.types
 import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
@@ -79,6 +81,27 @@ internal fun <T : Number> Sequence<T>.meanOrNull(type: KType, skipNA: Boolean =
     }
 }
 
+internal fun meanTypeResultOrNull(type: KType, emptyInput: Boolean): KType? =
+    when (val type = type.withNullability(false)) {
+        typeOf<Double>(),
+        typeOf<Float>(),
+        typeOf<Int>(),
+        typeOf<Short>(),
+        typeOf<Byte>(),
+        typeOf<Long>(),
+        -> typeOf<Double>().withNullability(emptyInput)
+
+        typeOf<BigInteger>(),
+        typeOf<BigDecimal>(),
+        -> typeOf<BigDecimal>().withNullability(emptyInput)
+
+        nothingType -> nullableNothingType
+
+        typeOf<Number>() -> null
+
+        else -> throw IllegalArgumentException("Unable to compute the mean for type ${renderType(type)}")
+    }
+
 internal fun Sequence<Double>.meanOrNull(skipNA: Boolean = skipNA_default): Double? {
     var count = 0
     var sum: Double = 0.toDouble()

Original file line number	Diff line number	Diff line change
`@@ -10082,6 +10082,7 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/impl/aggregation`
`10082`	`10082`	`public abstract fun aggregate (Ljava/lang/Iterable;Lkotlin/reflect/KType;)Ljava/lang/Object;`
`10083`	`10083`	`public abstract fun aggregate (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Object;`
`10084`	`10084`	`public abstract fun aggregateCalculatingType (Ljava/lang/Iterable;Ljava/util/Set;)Ljava/lang/Object;`
	`10085`	`+ public abstract fun calculateReturnTypeOrNull (Lkotlin/reflect/KType;Z)Lkotlin/reflect/KType;`
`10085`	`10086`	`public abstract fun getName ()Ljava/lang/String;`
`10086`	`10087`	`public abstract fun getPreservesType ()Z`
`10087`	`10088`	`}`