From 2209c1da9c64d05295201b03c9946c18ab68114f Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Tue, 25 Feb 2025 16:16:32 +0100
Subject: [PATCH 01/18] WIP rework of mean functions

---
 .../jetbrains/kotlinx/dataframe/api/mean.kt   | 190 ++++++++++++++++--
 .../kotlinx/dataframe/impl/ExceptionUtils.kt  |  15 ++
 .../kotlinx/dataframe/impl/TypeUtils.kt       |  10 +
 .../aggregation/aggregators/Aggregators.kt    |  19 +-
 .../jetbrains/kotlinx/dataframe/math/mean.kt  | 116 +++++++----
 .../dataframe/examples/titanic/ml/titanic.kt  |   7 +-
 6 files changed, 291 insertions(+), 66 deletions(-)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
index 994cbf27db..63fc249950 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
@@ -21,30 +21,181 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.numberColumns
 import org.jetbrains.kotlinx.dataframe.impl.columns.toNumberColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull
 import org.jetbrains.kotlinx.dataframe.math.mean
+import java.math.BigDecimal
+import java.math.BigInteger
+import kotlin.experimental.ExperimentalTypeInference
 import kotlin.reflect.KProperty
 import kotlin.reflect.typeOf
 
 // region DataColumn
 
-public fun <T : Number> DataColumn<T?>.mean(skipNA: Boolean = skipNA_default): Double =
-    meanOrNull(skipNA).suggestIfNull("mean")
+// region mean
 
-public fun <T : Number> DataColumn<T?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
-    Aggregators.mean(skipNA).aggregate(this)
+@JvmName("meanInt")
+public fun DataColumn<Int?>.mean(): Double = meanOrNull().suggestIfNull("mean")
 
-public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
-    skipNA: Boolean = skipNA_default,
-    noinline expression: (T) -> R?,
-): Double = Aggregators.mean(skipNA).cast2<R?, Double>().aggregateOf(this, expression) ?: Double.NaN
+@JvmName("meanShort")
+public fun DataColumn<Short?>.mean(): Double = meanOrNull().suggestIfNull("mean")
+
+@JvmName("meanByte")
+public fun DataColumn<Byte?>.mean(): Double = meanOrNull().suggestIfNull("mean")
+
+@JvmName("meanLong")
+public fun DataColumn<Long?>.mean(): Double = meanOrNull().suggestIfNull("mean")
+
+@JvmName("meanDouble")
+public fun DataColumn<Double?>.mean(skipNA: Boolean = skipNA_default): Double = meanOrNull(skipNA).suggestIfNull("mean")
+
+@JvmName("meanFloat")
+public fun DataColumn<Float?>.mean(skipNA: Boolean = skipNA_default): Double = meanOrNull(skipNA).suggestIfNull("mean")
+
+@JvmName("meanBigInteger")
+public fun DataColumn<BigInteger?>.mean(): BigDecimal = meanOrNull().suggestIfNull("mean")
+
+@JvmName("meanBigDecimal")
+public fun DataColumn<BigDecimal?>.mean(): BigDecimal = meanOrNull().suggestIfNull("mean")
+
+@JvmName("meanNumber")
+public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Number? = meanOrNull(skipNA)
 
 // endregion
 
-// region DataRow
+// region meanOrNull
+
+@JvmName("meanOrNullInt")
+public fun DataColumn<Int?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
+
+@JvmName("meanOrNullShort")
+public fun DataColumn<Short?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
+
+@JvmName("meanOrNullByte")
+public fun DataColumn<Byte?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
+
+@JvmName("meanOrNullLong")
+public fun DataColumn<Long?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
 
+@JvmName("meanOrNullDouble")
+public fun DataColumn<Double?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
+    Aggregators.mean.toDouble(skipNA).aggregate(this)
+
+@JvmName("meanOrNullFloat")
+public fun DataColumn<Float?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
+    Aggregators.mean.toDouble(skipNA).aggregate(this)
+
+@JvmName("meanOrNullBigInteger")
+public fun DataColumn<BigInteger?>.meanOrNull(): BigDecimal? = Aggregators.mean.toBigDecimal.aggregate(this)
+
+@JvmName("meanOrNullBigDecimal")
+public fun DataColumn<BigDecimal?>.meanOrNull(): BigDecimal? = Aggregators.mean.toBigDecimal.aggregate(this)
+
+@JvmName("meanOrNullNumber")
+public fun DataColumn<Number?>.meanOrNull(skipNA: Boolean = skipNA_default): Number? =
+    Aggregators.mean.toNumber(skipNA).aggregate(this)
+
+// endregion
+
+// region meanOf
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfInt")
+//@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Int?): Double =
+    Aggregators.mean.toDouble(skipNA_default)
+        .cast2<Int?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfShort")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Short?): Double =
+    Aggregators.mean.toDouble(skipNA_default)
+        .cast2<Short?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfByte")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Byte?): Double =
+    Aggregators.mean.toDouble(skipNA_default)
+        .cast2<Byte?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfLong")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Long?): Double =
+    Aggregators.mean.toDouble(skipNA_default)
+        .cast2<Long?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfDouble")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double =
+    Aggregators.mean.toDouble(skipNA)
+        .cast2<Double?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfFloat")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double =
+    Aggregators.mean.toDouble(skipNA)
+        .cast2<Float?, Double>()
+        .aggregateOf(this, expression)
+        ?: Double.NaN
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfBigInteger")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigInteger?): BigDecimal? =
+    Aggregators.mean.toBigDecimal
+        .cast2<BigInteger?, BigDecimal?>()
+        .aggregateOf(this, expression)
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfBigDecimal")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigDecimal?): BigDecimal? =
+    Aggregators.mean.toBigDecimal
+        .cast2<BigDecimal?, BigDecimal?>()
+        .aggregateOf(this, expression)
+
+@OptIn(ExperimentalTypeInference::class)
+@JvmName("meanOfNumber")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number? =
+    Aggregators.mean.toNumber(skipNA)
+        .cast2<Number?, Number?>()
+        .aggregateOf(this, expression)
+
+public fun main() {
+    val data = (1..10).toList()
+    val df = data.toDataFrame()
+
+    val mean = df.value.meanOf { if (true) it.toLong() else it.toDouble() }
+    val mean2 = df.value.meanOf { it.toBigInteger() }
+
+    println(mean)
+    println(mean!!::class)
+}
+
+// endregion
+
+// endregion
+
+// region DataRow
+// todo
 public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Double =
     values().filterIsInstance<Number>().map { it.toDouble() }.mean(skipNA)
 
-public inline fun <reified T : Number> AnyRow.rowMeanOf(): Double = values().filterIsInstance<T>().mean(typeOf<T>())
+public inline fun <reified T : Number> AnyRow.rowMeanOf(): Double =
+    values().filterIsInstance<T>().mean(typeOf<T>()) as Double
 
 // endregion
 
@@ -55,7 +206,7 @@ public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default): DataRow<T> =
 public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataRow<T> = Aggregators.mean(skipNA).aggregateFor(this, columns)
+): DataRow<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, columns)
 
 public fun <T> DataFrame<T>.meanFor(vararg columns: String, skipNA: Boolean = skipNA_default): DataRow<T> =
     meanFor(skipNA) { columns.toNumberColumns() }
@@ -72,10 +223,11 @@ public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
 ): DataRow<T> = meanFor(skipNA) { columns.toColumnSet() }
 
+// todo
 public fun <T, C : Number> DataFrame<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): Double = Aggregators.mean(skipNA).aggregateAll(this, columns) as Double? ?: Double.NaN
+): Double = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns) as Double? ?: Double.NaN
 
 public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Double =
     mean(skipNA) { columns.toNumberColumns() }
@@ -93,7 +245,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.toNumber(skipNA).of(this, expression) as Double? ?: Double.NaN
 
 // endregion
 
@@ -104,7 +256,7 @@ public fun <T> Grouped<T>.mean(skipNA: Boolean = skipNA_default): DataFrame<T> =
 public fun <T, C : Number> Grouped<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateFor(this, columns)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, columns)
 
 public fun <T> Grouped<T>.meanFor(vararg columns: String, skipNA: Boolean = skipNA_default): DataFrame<T> =
     meanFor(skipNA) { columns.toNumberColumns() }
@@ -125,7 +277,7 @@ public fun <T, C : Number> Grouped<T>.mean(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateAll(this, name, columns)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateAll(this, name, columns)
 
 public fun <T> Grouped<T>.mean(
     vararg columns: String,
@@ -151,7 +303,7 @@ public inline fun <T, reified R : Number> Grouped<T>.meanOf(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateOf(this, name, expression)
 
 // endregion
 
@@ -207,7 +359,7 @@ public fun <T, C : Number> PivotGroupBy<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     separate: Boolean = false,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateFor(this, separate, columns)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, separate, columns)
 
 public fun <T> PivotGroupBy<T>.meanFor(
     vararg columns: String,
@@ -232,7 +384,7 @@ public fun <T, C : Number> PivotGroupBy<T>.meanFor(
 public fun <T, R : Number> PivotGroupBy<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, R?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateAll(this, columns)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns)
 
 public fun <T> PivotGroupBy<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): DataFrame<T> =
     mean(skipNA) { columns.toColumnsSetOf() }
@@ -252,6 +404,6 @@ public fun <T, R : Number> PivotGroupBy<T>.mean(
 public inline fun <T, reified R : Number> PivotGroupBy<T>.meanOf(
     skipNA: Boolean = skipNA_default,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.mean(skipNA).aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateOf(this, expression)
 
 // endregion
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
index 6da0acdbac..93bbb5950d 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
@@ -1,7 +1,22 @@
 package org.jetbrains.kotlinx.dataframe.impl
 
+import java.math.BigDecimal
+import java.math.BigInteger
+
 internal fun <T> T?.throwIfNull(message: String): T = this ?: throw NoSuchElementException(message)
 
 @PublishedApi
 internal fun <T> T?.suggestIfNull(operation: String): T =
     throwIfNull("No elements for `$operation` operation. Use `${operation}OrNull` instead.")
+
+@PublishedApi
+internal fun BigInteger?.suggestIfNull(operation: String): BigInteger =
+    throwIfNull(
+        "The `$operation` operation either had no elements, or the result is NaN. Use `${operation}OrNull` instead.",
+    )
+
+@PublishedApi
+internal fun BigDecimal?.suggestIfNull(operation: String): BigDecimal =
+    throwIfNull(
+        "The `$operation` operation either had no elements, or the result is NaN. Use `${operation}OrNull` instead.",
+    )
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
index 38be1760bf..7841a70b2d 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
@@ -643,3 +643,13 @@ internal fun Iterable<Any>.classes(): Set<KClass<*>> = mapTo(mutableSetOf()) { i
  * @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) }
+
+/**
+ * Casts [this]: [Number] to a [Double]. If [this] is `null`, returns [Double.NaN].
+ */
+internal fun Number?.asDoubleOrNaN(): Double = this as Double? ?: Double.NaN
+
+/**
+ * Casts [this]: [Number] to a [Float]. If [this] is `null`, returns [Float.NaN].
+ */
+internal fun Number?.asFloatOrNaN(): Float = this as Float? ?: Float.NaN
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
index 4c90f286d8..5cde09801f 100644
--- 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
@@ -38,8 +38,23 @@ internal object Aggregators {
         mergedValuesChangingTypes { std(it, skipNA, ddof) }
     }
 
-    val mean by withOption<Boolean, Number, Double> { skipNA ->
-        changesType({ mean(it, skipNA) }) { mean(skipNA) }
+    @Suppress("ClassName")
+    object mean {
+        val toNumber = withOption { skipNA: Boolean ->
+            extendsNumbers { mean(it, skipNA) }
+        }.create("meanToNumber")
+
+        val toDouble = withOption { skipNA: Boolean ->
+            changesType(
+                aggregateWithType = { mean(it, skipNA).asDoubleOrNaN() },
+                aggregateWithValues = { mean(skipNA) },
+            )
+        }.create("meanToDouble")
+
+        val toBigDecimal = changesType(
+            aggregateWithType = { mean(it) as BigDecimal? },
+            aggregateWithValues = { filterNotNull().mean() },
+        ).create("meanToBigDecimal")
     }
 
     val percentile by withOption<Double, Comparable<Any?>, Comparable<Any?>> { percentile ->
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
index a1ab845624..3786feb7be 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -1,43 +1,79 @@
 package org.jetbrains.kotlinx.dataframe.math
 
+import org.jetbrains.kotlinx.dataframe.api.isNaN
 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.renderType
+import org.jetbrains.kotlinx.dataframe.impl.types
+import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
+/** @include [Sequence.mean] */
 @PublishedApi
-internal fun <T : Number> Iterable<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Double =
+internal fun <T : Number> Iterable<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Number? =
     asSequence().mean(type, skipNA)
 
+/**
+ * Returns the mean of the numbers in [this].
+ *
+ * If the input is empty, the return value will be `null`.
+ *
+ * If the [type] given or input consists of only [Int], [Short], [Byte], [Long], [Double], or [Float],
+ * the return type will be [Double]`?` (Never `NaN`).
+ *
+ * If the [type] given or the input contains [BigInteger] or [BigDecimal], the return type will be [BigDecimal]`?`.
+ * @param type The type of the numbers in the sequence.
+ * @param skipNA Whether to skip `NaN` values (default: `false`). Only relevant for [Double] and [Float].
+ */
 @Suppress("UNCHECKED_CAST")
-internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Double {
+internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Number? {
     if (type.isMarkedNullable) {
         return filterNotNull().mean(type.withNullability(false), skipNA)
     }
     return when (type.classifier) {
-        Double::class -> (this as Sequence<Double>).mean(skipNA)
+        // Double -> Double?
+        Double::class -> (this as Sequence<Double>).mean(skipNA).takeUnless { it.isNaN }
 
-        Float::class -> (this as Sequence<Float>).mean(skipNA)
+        // Float -> Double?
+        Float::class -> (this as Sequence<Float>).mean(skipNA).takeUnless { it.isNaN }
 
-        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.mean(false)
+        // Int -> Double?
+        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
 
-        // for integer values NA is not possible
-        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.mean(false)
+        // Short -> Double?
+        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
 
-        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.mean(false)
+        // Byte -> Double?
+        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
 
-        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.mean(false)
+        // Long -> Double?
+        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
 
-        BigInteger::class -> (this as Sequence<BigInteger>).map { it.toDouble() }.mean(false)
+        // BigInteger -> BigDecimal?
+        BigInteger::class -> (this as Sequence<BigInteger>).mean()
 
-        BigDecimal::class -> (this as Sequence<BigDecimal>).map { it.toDouble() }.mean(skipNA)
+        // BigDecimal -> BigDecimal?
+        BigDecimal::class -> (this as Sequence<BigDecimal>).mean()
 
-        Number::class -> (this as Sequence<Number>).map { it.toDouble() }.mean(skipNA)
+        // Number -> Conversion(Common number type) -> Number? (Double or BigDecimal?)
+        // fallback case, heavy as it needs to collect all types at runtime
+        Number::class -> {
+            val numberTypes = (this as Sequence<Number>).asIterable().types()
+            val unifiedType = numberTypes.unifiedNumberType()
+            if (unifiedType.withNullability(false) == typeOf<Number>()) {
+                error("Cannot find unified number type for $numberTypes")
+            }
+            this.convertToUnifiedNumberType(unifiedType)
+                .mean(unifiedType, skipNA)
+        }
 
         // this means the sequence is empty
-        Nothing::class -> Double.NaN
+        Nothing::class -> null
 
         else -> throw IllegalArgumentException("Unable to compute the mean for type ${renderType(type)}")
     }
@@ -78,12 +114,38 @@ internal fun Sequence<Float>.mean(skipNA: Boolean = skipNA_default): Double {
     return if (count > 0) sum / count else Double.NaN
 }
 
+@JvmName("bigIntegerMean")
+internal fun Sequence<BigInteger>.mean(): BigDecimal? {
+    var count = 0
+    val sum = sumOf {
+        count++
+        it
+    }
+    return if (count > 0) sum.toBigDecimal() / count.toBigDecimal() else null
+}
+
+@JvmName("bigDecimalMean")
+internal fun Sequence<BigDecimal>.mean(): BigDecimal? {
+    var count = 0
+    val sum = sumOf {
+        count++
+        it
+    }
+    return if (count > 0) sum.toBigDecimal() / count.toBigDecimal() else null
+}
+
 @JvmName("doubleMean")
 internal fun Iterable<Double>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
 
 @JvmName("floatMean")
 internal fun Iterable<Float>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
 
+@JvmName("bigDecimalMean")
+internal fun Iterable<BigDecimal>.mean(): BigDecimal? = asSequence().mean()
+
+@JvmName("bigIntegerMean")
+internal fun Iterable<BigInteger>.mean(): BigDecimal? = asSequence().mean()
+
 @JvmName("intMean")
 internal fun Iterable<Int>.mean(): Double =
     if (this is Collection) {
@@ -135,31 +197,3 @@ internal fun Iterable<Long>.mean(): Double =
         }
         if (count > 0) sum / count else Double.NaN
     }
-
-// TODO result is Double, but should be BigDecimal, Issue #558
-@JvmName("bigIntegerMean")
-internal fun Iterable<BigInteger>.mean(): Double =
-    if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
-    } else {
-        var count = 0
-        val sum = sumOf {
-            count++
-            it.toDouble()
-        }
-        if (count > 0) sum / count else Double.NaN
-    }
-
-// TODO result is Double, but should be BigDecimal, Issue #558
-@JvmName("bigDecimalMean")
-internal fun Iterable<BigDecimal>.mean(): Double =
-    if (this is Collection) {
-        if (size > 0) sum().toDouble() / size else Double.NaN
-    } else {
-        var count = 0
-        val sum = sumOf {
-            count++
-            it.toDouble()
-        }
-        if (count > 0) sum / count else Double.NaN
-    }
diff --git a/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt b/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
index b29fd1625a..64d0e9423c 100644
--- a/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
+++ b/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
@@ -24,11 +24,10 @@ private val model = Sequential.of(
     Input(9),
     Dense(50, Activations.Relu, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
     Dense(50, Activations.Relu, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
-    Dense(2, Activations.Linear, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros())
+    Dense(2, Activations.Linear, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
 )
 
 fun main() {
-
     // Set Locale for correct number parsing
     Locale.setDefault(Locale.FRANCE)
 
@@ -37,7 +36,7 @@ fun main() {
     // Calculating imputing values
     val (train, test) = df
         // imputing
-        .fillNulls { sibsp and parch and age and fare }.perCol { it.mean() }
+        .fillNulls { sibsp and parch and age and fare }.perCol { it.mean()?.toDouble() }
         .fillNulls { sex }.with { "female" }
         // one hot encoding
         .pivotMatches { pclass and sex }
@@ -50,7 +49,7 @@ fun main() {
         it.compile(
             optimizer = Adam(),
             loss = Losses.SOFT_MAX_CROSS_ENTROPY_WITH_LOGITS,
-            metric = Metrics.ACCURACY
+            metric = Metrics.ACCURACY,
         )
 
         it.summary()

From 756473ea6be61c7ca97113c3d7037eda2ef23c1b Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Wed, 26 Feb 2025 16:19:41 +0100
Subject: [PATCH 02/18] WIP rework of aggregator implementation

---
 .../jetbrains/kotlinx/dataframe/api/median.kt |   5 +-
 .../jetbrains/kotlinx/dataframe/api/sum.kt    |   4 +-
 .../kotlinx/dataframe/impl/NumberTypeUtils.kt |  12 ++
 .../aggregation/aggregators/Aggregator.kt     |  55 +++++++--
 .../aggregation/aggregators/AggregatorBase.kt |  64 ++++++++--
 .../aggregators/AggregatorOptionSwitch.kt     |  73 ++++++++---
 .../aggregators/AggregatorProvider.kt         |  26 +++-
 .../aggregation/aggregators/Aggregators.kt    | 114 +++++++++++++-----
 .../aggregators/FlatteningAggregator.kt       |  62 ++++++++++
 .../aggregators/MergedValuesAggregator.kt     |  42 -------
 .../aggregators/NumbersAggregator.kt          |  37 ------
 .../aggregators/TwoStepAggregator.kt          |  86 ++++++++++---
 .../aggregators/TwoStepNumbersAggregator.kt   |  78 ++++++++++++
 13 files changed, 494 insertions(+), 164 deletions(-)
 create mode 100644 core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
 delete mode 100644 core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
 delete mode 100644 core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
 create mode 100644 core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index f2cdbb390e..f939609362 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -39,8 +39,9 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOf(noinline
 // region DataRow
 
 public fun AnyRow.rowMedianOrNull(): Any? =
-    Aggregators.median.aggregateMixed(
-        values().filterIsInstance<Comparable<Any?>>().asIterable(),
+    Aggregators.median.aggregateCalculatingType(
+        values = values().filterIsInstance<Comparable<Any?>>().asIterable(),
+        valueTypes = df().columns().filter { it.valuesAreComparable() }.map { it.type() }.toSet(),
     )
 
 public fun AnyRow.rowMedian(): Any = rowMedianOrNull().suggestIfNull("rowMedian")
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
index af9bea3657..9e84d74622 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
@@ -44,9 +44,9 @@ public inline fun <T, reified R : Number> DataColumn<T>.sumOf(crossinline expres
 // region DataRow
 
 public fun AnyRow.rowSum(): Number =
-    Aggregators.sum.aggregateMixed(
+    Aggregators.sum.aggregateCalculatingType(
         values = values().filterIsInstance<Number>(),
-        types = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
+        valueTypes = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
     ) ?: 0
 
 public inline fun <reified T : Number> AnyRow.rowSumOf(): T = values().filterIsInstance<T>().sum(typeOf<T>())
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
index 06f2a92a74..1745147f47 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
@@ -137,3 +137,15 @@ internal fun Iterable<Number>.convertToUnifiedNumberType(
         converter(it) ?: error("Can not convert $it to $commonNumberType")
     }
 }
+
+/** @include [Iterable.convertToUnifiedNumberType] */
+@JvmName("convertToUnifiedNumberTypeSequence")
+@Suppress("UNCHECKED_CAST")
+internal fun Sequence<Number>.convertToUnifiedNumberType(
+    commonNumberType: KType = asIterable().types().unifiedNumberType(),
+): Sequence<Number> {
+    val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
+    return map {
+        converter(it) ?: error("Can not convert $it to $commonNumberType")
+    }
+}
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
index dcd88a15a7..bc5b82b0ce 100644
--- 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
@@ -3,22 +3,63 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import kotlin.reflect.KType
 
+/**
+ * Base interface for all aggregators.
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * 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.
+ * @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]`?`.
+ */
 @PublishedApi
-internal interface Aggregator<C, R> {
+internal interface Aggregator<Value, Return> {
 
+    /** The name of this aggregator. */
     val name: String
 
-    fun aggregate(column: DataColumn<C?>): R?
-
+    /** If `true`, [Value][Value]`  ==  ` [Return][Return]. */
     val preservesType: Boolean
 
-    fun aggregate(columns: Iterable<DataColumn<C?>>): R?
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * When using [AggregatorBase], this can be supplied by the [AggregatorBase.aggregator] argument.
+     */
+    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 the aggregation function (with [Iterable] and [KType]) is called.
+     *
+     * See [AggregatorBase.aggregate].
+     */
+    fun aggregate(column: DataColumn<Value?>): Return?
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     *
+     * Must be overridden when using [AggregatorBase].
+     */
+    fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
 
-    fun aggregate(values: Iterable<C>, type: KType): R?
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     * 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?
 }
 
 @PublishedApi
-internal fun <T> Aggregator<*, *>.cast(): Aggregator<T, T> = this as Aggregator<T, T>
+internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
 
 @PublishedApi
-internal fun <T, P> Aggregator<*, *>.cast2(): Aggregator<T, P> = this as Aggregator<T, P>
+internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>
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
index 1deb052b2f..0683a5e9ec 100644
--- 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
@@ -3,19 +3,69 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 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 kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
-internal abstract class AggregatorBase<C, R>(
+/**
+ * Base class for [aggregators][Aggregator].
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * @param name The name of this aggregator.
+ * @param aggregator Functional argument for the [aggregate] function.
+ */
+internal abstract class AggregatorBase<Value, Return>(
     override val name: String,
-    protected val aggregator: (Iterable<C>, KType) -> R?,
-) : Aggregator<C, R> {
+    protected val aggregator: (values: Iterable<Value>, type: KType) -> Return?,
+) : Aggregator<Value, Return> {
 
-    override fun aggregate(column: DataColumn<C?>): R? =
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     * Uses [aggregator] to compute the result.
+     */
+    override fun aggregate(values: Iterable<Value>, type: KType): Return? = aggregator(values, type)
+
+    /**
+     * 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].
+     */
+    override fun aggregate(column: DataColumn<Value?>): Return? =
         if (column.hasNulls()) {
-            aggregate(column.asSequence().filterNotNull().asIterable(), column.type())
+            aggregate(column.asSequence().filterNotNull().asIterable(), column.type().withNullability(false))
+        } else {
+            aggregate(column.asIterable() as Iterable<Value>, column.type().withNullability(false))
+        }
+
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     */
+    override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return? {
+        val commonType = if (valueTypes != null) {
+            valueTypes.commonType(false)
         } else {
-            aggregate(column.asIterable() as Iterable<C>, column.type())
+            var hasNulls = false
+            val classes = values.mapNotNull {
+                if (it == null) {
+                    hasNulls = true
+                    null
+                } else {
+                    it.javaClass.kotlin
+                }
+            }
+            classes.commonType(hasNulls)
         }
+        return aggregate(values, commonType)
+    }
 
-    override fun aggregate(values: Iterable<C>, type: KType): R? = aggregator(values, type)
+    /**
+     * 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?
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index 45cb01be19..d16def6dcb 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -1,33 +1,72 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
-import kotlin.reflect.KProperty
-
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require a single parameter.
+ *
+ * Aggregators are cached by their parameter value.
+ * @see AggregatorOptionSwitch2
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch<P, C, R>(val name: String, val getAggregator: (P) -> AggregatorProvider<C, R>) {
+internal class AggregatorOptionSwitch1<Param1, AggregatorType : Aggregator<*, *>>(
+    val name: String,
+    val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
+) {
 
-    private val cache = mutableMapOf<P, Aggregator<C, R>>()
+    private val cache: MutableMap<Param1, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option: P) = cache.getOrPut(option) { getAggregator(option).create(name) }
+    operator fun invoke(param1: Param1): AggregatorType =
+        cache.getOrPut(param1) {
+            getAggregator(param1).create(name)
+        }
 
-    class Factory<P, C, R>(val getAggregator: (P) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch1].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch1.Factory { param1: Param1 ->
+     *   MyAggregator.Factory(param1)
+     * }
+     */
+    class Factory<Param1, AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch1(name, getAggregator)
+        })
 }
 
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require two parameters.
+ *
+ * Aggregators are cached by their parameter values.
+ * @see AggregatorOptionSwitch1
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch2<P1, P2, C, R>(
+internal class AggregatorOptionSwitch2<Param1, Param2, AggregatorType : Aggregator<*, *>>(
     val name: String,
-    val getAggregator: (P1, P2) -> AggregatorProvider<C, R>,
+    val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
 ) {
 
-    private val cache = mutableMapOf<Pair<P1, P2>, Aggregator<C, R>>()
+    private val cache: MutableMap<Pair<Param1, Param2>, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option1: P1, option2: P2) =
-        cache.getOrPut(option1 to option2) {
-            getAggregator(option1, option2).create(name)
+    operator fun invoke(param1: Param1, param2: Param2): AggregatorType =
+        cache.getOrPut(param1 to param2) {
+            getAggregator(param1, param2).create(name)
         }
 
-    class Factory<P1, P2, C, R>(val getAggregator: (P1, P2) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch2(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch2].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch2.Factory { param1: Param1, param2: Param2 ->
+     *   MyAggregator.Factory(param1, param2)
+     * }
+     * ```
+     */
+    class Factory<Param1, Param2, AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch2(name, getAggregator)
+        })
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
index a8265a8175..a0cbea44fd 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
@@ -2,9 +2,27 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import kotlin.reflect.KProperty
 
-internal interface AggregatorProvider<C, R> {
+/**
+ * Common interface for providers or "factory" objects that create anything of type [T].
+ *
+ * When implemented, this allows the object to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myNamedValue by MyFactory
+ * ```
+ */
+internal fun interface Provider<T> {
 
-    operator fun getValue(obj: Any?, property: KProperty<*>): Aggregator<C, R> = create(property.name)
-
-    fun create(name: String): Aggregator<C, R>
+    fun create(name: String): T
 }
+
+internal operator fun <T> Provider<T>.getValue(obj: Any?, property: KProperty<*>): T = create(property.name)
+
+/**
+ * Common interface for providers of [Aggregators][Aggregator] or "factory" objects that create aggregators.
+ *
+ * When implemented, this allows an aggregator to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myAggregator by MyAggregator.Factory
+ * ```
+ */
+internal fun interface AggregatorProvider<AggregatorType : Aggregator<*, *>> : Provider<AggregatorType>
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
index 4c90f286d8..27eef93b91 100644
--- 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
@@ -10,43 +10,93 @@ import kotlin.reflect.KType
 @PublishedApi
 internal object Aggregators {
 
-    private fun <C> preservesType(aggregate: Iterable<C>.(KType) -> C?) =
-        TwoStepAggregator.Factory(aggregate, aggregate, true)
-
-    private fun <C : Any, R> mergedValues(aggregate: Iterable<C?>.(KType) -> R?) =
-        MergedValuesAggregator.Factory(aggregate, true)
-
-    private fun <C : Any, R> mergedValuesChangingTypes(aggregate: Iterable<C?>.(KType) -> R?) =
-        MergedValuesAggregator.Factory(aggregate, false)
-
-    private fun <C, R> changesType(aggregate1: Iterable<C>.(KType) -> R, aggregate2: Iterable<R>.(KType) -> R) =
-        TwoStepAggregator.Factory(aggregate1, aggregate2, false)
-
-    private fun extendsNumbers(aggregate: Iterable<Number>.(KType) -> Number?) = NumbersAggregator.Factory(aggregate)
-
-    private fun <P, C, R> withOption(getAggregator: (P) -> AggregatorProvider<C, R>) =
-        AggregatorOptionSwitch.Factory(getAggregator)
-
-    private fun <P1, P2, C, R> withOption2(getAggregator: (P1, P2) -> AggregatorProvider<C, R>) =
-        AggregatorOptionSwitch2.Factory(getAggregator)
-
-    val min by preservesType<Comparable<Any?>> { minOrNull() }
-
-    val max by preservesType<Comparable<Any?>> { maxOrNull() }
-
-    val std by withOption2<Boolean, Int, Number, Double> { skipNA, ddof ->
-        mergedValuesChangingTypes { std(it, skipNA, ddof) }
+    /**
+     * Factory for a simple aggregator that preserves the type of the input values.
+     *
+     * @include [TwoStepAggregator]
+     */
+    private fun <Type> twoStepPreservingType(aggregator: Iterable<Type>.(type: KType) -> Type?) =
+        TwoStepAggregator.Factory(
+            stepOneAggregator = aggregator,
+            stepTwoAggregator = aggregator,
+            preservesType = true,
+        )
+
+    /**
+     * Factory for a simple aggregator that changes the type of the input values.
+     *
+     * @include [TwoStepAggregator]
+     */
+    private fun <Value, Return> twoStepChangingType(
+        stepOneAggregator: Iterable<Value>.(type: KType) -> Return,
+        stepTwoAggregator: Iterable<Return>.(type: KType) -> Return,
+    ) = TwoStepAggregator.Factory(
+        stepOneAggregator = stepOneAggregator,
+        stepTwoAggregator = stepTwoAggregator,
+        preservesType = false,
+    )
+
+    /**
+     * Factory for a flattening aggregator that preserves the type of the input values.
+     *
+     * @include [FlatteningAggregator]
+     */
+    private fun <Type> flatteningPreservingTypes(aggregate: Iterable<Type?>.(type: KType) -> Type?) =
+        FlatteningAggregator.Factory(
+            aggregator = aggregate,
+            preservesType = true,
+        )
+
+    /**
+     * Factory for a flattening aggregator that changes the type of the input values.
+     *
+     * @include [FlatteningAggregator]
+     */
+    private fun <Value, Return> flatteningChangingTypes(aggregate: Iterable<Value?>.(type: KType) -> Return?) =
+        FlatteningAggregator.Factory(
+            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)
+
+    /** @include [AggregatorOptionSwitch1] */
+    private fun <Param1, AggregatorType : Aggregator<*, *>> withOneOption(
+        getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
+    ) = AggregatorOptionSwitch1.Factory(getAggregator)
+
+    /** @include [AggregatorOptionSwitch2] */
+    private fun <Param1, Param2, AggregatorType : Aggregator<*, *>> withTwoOptions(
+        getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
+    ) = AggregatorOptionSwitch2.Factory(getAggregator)
+
+    val min by twoStepPreservingType<Comparable<Any?>> { minOrNull() }
+
+    val max by twoStepPreservingType<Comparable<Any?>> { maxOrNull() }
+
+    val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
+        flatteningChangingTypes<Number, Double> { std(it, skipNA, ddof) }
     }
 
-    val mean by withOption<Boolean, Number, Double> { skipNA ->
-        changesType({ mean(it, skipNA) }) { mean(skipNA) }
+    val mean by withOneOption { skipNA: Boolean ->
+        twoStepChangingType({ mean(it, skipNA) }) { mean(skipNA) }
     }
 
-    val percentile by withOption<Double, Comparable<Any?>, Comparable<Any?>> { percentile ->
-        mergedValuesChangingTypes { type -> percentile(percentile, type) }
+    val percentile by withOneOption { percentile: Double ->
+        flatteningChangingTypes<Comparable<Any?>, Comparable<Any?>> { type ->
+            percentile(percentile, type)
+        }
     }
 
-    val median by mergedValues<Comparable<Any?>, Comparable<Any?>> { median(it) }
+    val median by flatteningPreservingTypes<Comparable<Any?>> {
+        median(it)
+    }
 
-    val sum by extendsNumbers { sum(it) }
+    val sum by twoStepForNumbers { 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
new file mode 100644
index 0000000000..4561ac4991
--- /dev/null
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
@@ -0,0 +1,62 @@
+package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
+
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
+
+/**
+ * Simple [Aggregator] implementation with flattening behavior for multiple columns.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns,
+ * the columns are flattened into a single list of values, filtering nulls as usual;
+ * then the aggregation function is called with their common type.
+ *
+ * ```
+ * Iterable<Column<Value?>>
+ *     -> Iterable<Value> // flattened without nulls
+ *     -> aggregator(Iterable<Value>, common colType)
+ *     -> Return?
+ * ```
+ *
+ * This is essential for aggregators that depend on the distribution of all values across the dataframe.
+ *
+ * See [TwoStepAggregator] for different behavior for multiple columns.
+ *
+ * @param name The name of this aggregator.
+ * @param aggregator Functional argument for the [aggregate] function.
+ *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
+ * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+ */
+internal class FlatteningAggregator<Value, Return>(
+    name: String,
+    aggregator: (values: Iterable<Value>, type: KType) -> Return?,
+    override val preservesType: Boolean,
+) : AggregatorBase<Value, Return>(name, aggregator) {
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     * 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? {
+        val commonType = columns.map { it.type() }.commonType().withNullability(false)
+        val allValues = columns.asSequence().flatMap { it.values() }.filterNotNull()
+        return aggregate(allValues.asIterable(), commonType)
+    }
+
+    /**
+     * Creates [FlatteningAggregator].
+     *
+     * @param aggregator Functional argument for the [aggregate] function.
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    class Factory<Value, Return>(
+        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)
+        })
+}
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
deleted file mode 100644
index 135ba0a5ec..0000000000
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
+++ /dev/null
@@ -1,42 +0,0 @@
-package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
-
-import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.impl.commonType
-import kotlin.reflect.KProperty
-import kotlin.reflect.KType
-
-internal class MergedValuesAggregator<C : Any, R>(
-    name: String,
-    val aggregateWithType: (Iterable<C?>, KType) -> R?,
-    override val preservesType: Boolean,
-) : AggregatorBase<C, R>(name, aggregateWithType) {
-
-    override fun aggregate(columns: Iterable<DataColumn<C?>>): R? {
-        val commonType = columns.map { it.type() }.commonType()
-        val allValues = columns.flatMap { it.values() }
-        return aggregateWithType(allValues, commonType)
-    }
-
-    fun aggregateMixed(values: Iterable<C?>): R? {
-        var hasNulls = false
-        val classes = values.mapNotNull {
-            if (it == null) {
-                hasNulls = true
-                null
-            } else {
-                it.javaClass.kotlin
-            }
-        }
-        return aggregateWithType(values, classes.commonType(hasNulls))
-    }
-
-    class Factory<C : Any, R>(
-        private val aggregateWithType: (Iterable<C?>, KType) -> R?,
-        private val preservesType: Boolean,
-    ) : AggregatorProvider<C, R> {
-        override fun create(name: String) = MergedValuesAggregator(name, aggregateWithType, preservesType)
-
-        override operator fun getValue(obj: Any?, property: KProperty<*>): MergedValuesAggregator<C, R> =
-            create(property.name)
-    }
-}
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
deleted file mode 100644
index 00ef22febe..0000000000
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
+++ /dev/null
@@ -1,37 +0,0 @@
-package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
-
-import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
-import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
-import kotlin.reflect.KProperty
-import kotlin.reflect.KType
-
-internal class NumbersAggregator(name: String, aggregate: (Iterable<Number>, KType) -> Number?) :
-    AggregatorBase<Number, Number>(name, aggregate) {
-
-    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Number? =
-        aggregateMixed(
-            values = columns.mapNotNull { aggregate(it) },
-            types = columns.map { it.type() }.toSet(),
-        )
-
-    class Factory(private val aggregate: Iterable<Number>.(KType) -> Number?) : AggregatorProvider<Number, Number> {
-        override fun create(name: String) = NumbersAggregator(name, aggregate)
-
-        override operator fun getValue(obj: Any?, property: KProperty<*>): NumbersAggregator = create(property.name)
-    }
-
-    /**
-     * Can aggregate numbers with different types by first converting them to a compatible type.
-     */
-    @Suppress("UNCHECKED_CAST")
-    fun aggregateMixed(values: Iterable<Number>, types: Set<KType>): Number? {
-        val commonType = types.unifiedNumberType()
-        return aggregate(
-            values = values.convertToUnifiedNumberType(commonType),
-            type = commonType,
-        )
-    }
-
-    override val preservesType = false
-}
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
index 9d01169d02..b42597ea8a 100644
--- 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,27 +1,85 @@
 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 kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
-internal class TwoStepAggregator<C, R>(
+/**
+ * A slightly more advanced [Aggregator] implementation.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns, this [Aggregator] works in two steps:
+ * First, it aggregates within a [DataColumn]/[Iterable] ([stepOneAggregator]) with their (given) type,
+ * and then in between different columns ([stepTwoAggregator]) using the results of the first and the newly
+ * calculated common type of those results.
+ *
+ * ```
+ * Iterable<Column<Value?>>
+ *     -> Iterable<Iterable<Value>> // nulls filtered out
+ *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
+ *     -> Iterable<Return> // nulls filtered out
+ *     -> stepTwoAggregator(Iterable<Return>, common valueType)
+ *     -> Return?
+ * ```
+ *
+ * It can also be used as a "simple" aggregator by providing the same function for both steps,
+ * requires [preservesType] be set to `true`.
+ *
+ * See [FlatteningAggregator] for different behavior for multiple columns.
+ *
+ * @param name The name of this aggregator.
+ * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+ * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+ *   It is run on the results of [stepOneAggregator].
+ * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+ */
+internal class TwoStepAggregator<Value, Return>(
     name: String,
-    aggregateWithType: (Iterable<C>, KType) -> R?,
-    private val aggregateValues: (Iterable<R>, KType) -> R?,
+    stepOneAggregator: (values: Iterable<Value>, type: KType) -> Return?,
+    private val stepTwoAggregator: (values: Iterable<Return>, type: KType) -> Return?,
     override val preservesType: Boolean,
-) : AggregatorBase<C, R>(name, aggregateWithType) {
+) : AggregatorBase<Value, Return>(name, stepOneAggregator) {
 
-    override fun aggregate(columns: Iterable<DataColumn<C?>>): R? {
-        val columnValues = columns.mapNotNull { aggregate(it) }
-        val commonType = columnValues.map { it.javaClass.kotlin }.commonType(false)
-        return aggregateValues(columnValues, commonType)
+    /**
+     * 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 {
+            // 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)
     }
 
-    class Factory<C, R>(
-        private val aggregateWithType: (Iterable<C>, KType) -> R?,
-        private val aggregateValues: (Iterable<R>, KType) -> R?,
+    /**
+     * Creates [TwoStepAggregator].
+     *
+     * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+     * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+     *   It is run on the results of [stepOneAggregator].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    class Factory<Value, Return>(
+        private val stepOneAggregator: (Iterable<Value>, KType) -> Return?,
+        private val stepTwoAggregator: (Iterable<Return>, KType) -> Return?,
         private val preservesType: Boolean,
-    ) : AggregatorProvider<C, R> {
-        override fun create(name: String) = TwoStepAggregator(name, aggregateWithType, aggregateValues, preservesType)
-    }
+    ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
+            TwoStepAggregator(
+                name = name,
+                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
new file mode 100644
index 0000000000..74cfa32618
--- /dev/null
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
@@ -0,0 +1,78 @@
+package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
+
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
+import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+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.withNullability
+import kotlin.reflect.typeOf
+
+/**
+ * [Aggregator] made specifically for number calculations.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns (with potentially different [Number] types),
+ * this [Aggregator] works in two steps:
+ *
+ * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type,
+ * and then between different columns
+ * using the results of the first and the newly calculated [unified number][UnifyingNumbers] type of those results.
+ *
+ * ```
+ * Iterable<Column<Number?>>
+ *     -> Iterable<Iterable<Number>> // nulls filtered out
+ *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+ *     -> Iterable<Return> // nulls filtered out
+ *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+ *     -> Return?
+ * ```
+ *
+ * @param name The name of this aggregator.
+ * @param aggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+ *   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<Return : Number>(
+    name: String,
+    aggregator: (values: Iterable<Number>, numberType: KType) -> Return?,
+) : AggregatorBase<Number, Return>(name, aggregator) {
+
+    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 super.aggregate(values, type)
+    }
+
+    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return? =
+        aggregateCalculatingType(
+            values = columns.mapNotNull { aggregate(it) },
+            valueTypes = null, // makes the operation heavy
+        )
+
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
+     * of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     */
+    @Suppress("UNCHECKED_CAST")
+    override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
+        val commonType = (valueTypes ?: values.types()).unifiedNumberType().withNullability(false)
+        return aggregate(
+            values = values.convertToUnifiedNumberType(commonType),
+            type = commonType,
+        )
+    }
+
+    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)
+        })
+}

From 3ea01678fb08a9d0a8a486d5ebfc9ca73417005c Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Fri, 28 Feb 2025 16:55:37 +0100
Subject: [PATCH 03/18] mean rework: returns null for no values regardless of
 the type. Added orNull overloads for each mean function. Added specific
 overloads for each primitive type -> Double(?) and big number ->
 BigDecimal(?)

---
 core/api/core.api                             | 175 ++++++--
 .../kotlinx/dataframe/api/describe.kt         |   2 +-
 .../jetbrains/kotlinx/dataframe/api/mean.kt   | 374 +++++++++++++++---
 .../aggregation/aggregators/Aggregators.kt    |  30 +-
 .../aggregators/TwoStepNumbersAggregator.kt   |   4 +-
 .../impl/aggregation/modes/ofRowExpression.kt |   7 +-
 .../kotlinx/dataframe/impl/api/describe.kt    |   3 +-
 .../jetbrains/kotlinx/dataframe/math/mean.kt  |  98 ++---
 .../kotlinx/dataframe/api/describe.kt         |   6 +-
 .../kotlinx/dataframe/puzzles/BasicTests.kt   |   6 +-
 .../kotlinx/dataframe/puzzles/MediumTests.kt  |   2 +-
 .../dataframe/statistics/BasicMathTests.kt    |  17 +-
 .../testSets/animals/AnimalsTests.kt          |   6 +-
 .../testSets/person/DataFrameTests.kt         |  15 +-
 14 files changed, 562 insertions(+), 183 deletions(-)

diff --git a/core/api/core.api b/core/api/core.api
index 85a3020563..7409c6c769 100644
--- a/core/api/core.api
+++ b/core/api/core.api
@@ -2633,7 +2633,7 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/api/ColumnDescri
 	public abstract fun getCount ()I
 	public abstract fun getFreq ()I
 	public abstract fun getMax ()Ljava/lang/Object;
-	public abstract fun getMean ()D
+	public abstract fun getMean ()Ljava/lang/Number;
 	public abstract fun getMedian ()Ljava/lang/Object;
 	public abstract fun getMin ()Ljava/lang/Object;
 	public abstract fun getName ()Ljava/lang/String;
@@ -2655,7 +2655,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/ColumnDescription_Extensi
 	public static final fun ColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun ColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
-	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)D
+	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Number;
 	public static final fun ColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun ColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun ColumnDescription_min (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
@@ -2685,7 +2685,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/ColumnDescription_Extensi
 	public static final fun NullableColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun NullableColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
-	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Double;
+	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Number;
 	public static final fun NullableColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun NullableColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun NullableColumnDescription_min (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
@@ -5995,12 +5995,10 @@ public final class org/jetbrains/kotlinx/dataframe/api/MaxKt {
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)D
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)D
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)D
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Ljava/lang/Number;
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Ljava/lang/Number;
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Ljava/lang/Number;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Ljava/lang/String;ZLkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;[Ljava/lang/String;Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
@@ -6013,12 +6011,10 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataRow;
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)D
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)D
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)D
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Ljava/lang/Number;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Ljava/lang/String;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;[Ljava/lang/String;Ljava/lang/String;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
@@ -6031,6 +6027,20 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
+	public static final fun meanBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
+	public static final fun meanBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
+	public static final fun meanBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
+	public static final fun meanByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
+	public static final fun meanDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
+	public static final fun meanDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
+	public static synthetic fun meanDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
+	public static final fun meanFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
+	public static synthetic fun meanFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataRow;
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
@@ -6063,10 +6073,100 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
-	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
-	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun rowMean (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)D
-	public static synthetic fun rowMean$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)D
+	public static final fun meanInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
+	public static final fun meanInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
+	public static final fun meanLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
+	public static final fun meanLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
+	public static final fun meanNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Number;
+	public static final fun meanNumber (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Number;
+	public static synthetic fun meanNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun meanNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
+	public static final fun meanOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
+	public static final fun meanOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOfByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanOfByte$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)D
+	public static synthetic fun meanOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)D
+	public static synthetic fun meanOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanOfInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
+	public static final fun meanOfLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanOfLong$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanOfNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Number;
+	public static synthetic fun meanOfNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOfOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
+	public static final fun meanOfOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
+	public static final fun meanOfOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static final fun meanOfOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static synthetic fun meanOfOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOfOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static synthetic fun meanOfOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOfOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static final fun meanOfOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static final fun meanOfOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Number;
+	public static synthetic fun meanOfOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOfOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
+	public static final fun meanOfShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
+	public static final fun meanOfShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static synthetic fun meanOfShort$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Ljava/lang/Number;
+	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Ljava/lang/Number;
+	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Ljava/lang/Number;
+	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
+	public static final fun meanOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static final fun meanOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
+	public static final fun meanOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
+	public static final fun meanOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
+	public static final fun meanOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static final fun meanOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
+	public static final fun meanOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static final fun meanOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Number;
+	public static final fun meanOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Number;
+	public static synthetic fun meanOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun meanOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOrNullOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
+	public static final fun meanOrNullOfByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullOfByte$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullOfInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static final fun meanOrNullOfLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullOfLong$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullOfShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static synthetic fun meanOrNullOfShort$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun meanOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
+	public static final fun meanOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
+	public static final fun meanShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
+	public static final fun meanShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
+	public static final fun rowMean (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)Ljava/lang/Number;
+	public static synthetic fun rowMean$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun rowMeanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)Ljava/lang/Number;
+	public static synthetic fun rowMeanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)Ljava/lang/Number;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/MedianKt {
@@ -9930,6 +10030,8 @@ public final class org/jetbrains/kotlinx/dataframe/impl/DataFrameSize {
 
 public final class org/jetbrains/kotlinx/dataframe/impl/ExceptionUtilsKt {
 	public static final fun suggestIfNull (Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/Object;
+	public static final fun suggestIfNull (Ljava/math/BigDecimal;Ljava/lang/String;)Ljava/math/BigDecimal;
+	public static final fun suggestIfNull (Ljava/math/BigInteger;Ljava/lang/String;)Ljava/math/BigInteger;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/TypeUtilsKt {
@@ -9979,26 +10081,32 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/impl/aggregation
 	public abstract fun aggregate (Ljava/lang/Iterable;)Ljava/lang/Object;
 	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 getName ()Ljava/lang/String;
 	public abstract fun getPreservesType ()Z
 }
 
+public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator$DefaultImpls {
+	public static synthetic fun aggregateCalculatingType$default (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;Ljava/lang/Iterable;Ljava/util/Set;ILjava/lang/Object;)Ljava/lang/Object;
+}
+
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorKt {
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
 	public static final fun cast2 (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch {
+public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1 {
 	public fun <init> (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
 	public final fun getGetAggregator ()Lkotlin/jvm/functions/Function1;
 	public final fun getName ()Ljava/lang/String;
 	public final fun invoke (Ljava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch$Factory {
+public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1$Factory : org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Provider {
 	public fun <init> (Lkotlin/jvm/functions/Function1;)V
+	public synthetic fun create (Ljava/lang/String;)Ljava/lang/Object;
+	public fun create (Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
 	public final fun getGetAggregator ()Lkotlin/jvm/functions/Function1;
-	public final fun getValue (Ljava/lang/Object;Lkotlin/reflect/KProperty;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2 {
@@ -10008,21 +10116,28 @@ public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/
 	public final fun invoke (Ljava/lang/Object;Ljava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2$Factory {
+public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2$Factory : org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Provider {
 	public fun <init> (Lkotlin/jvm/functions/Function2;)V
+	public synthetic fun create (Ljava/lang/String;)Ljava/lang/Object;
+	public fun create (Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2;
 	public final fun getGetAggregator ()Lkotlin/jvm/functions/Function2;
-	public final fun getValue (Ljava/lang/Object;Lkotlin/reflect/KProperty;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators {
 	public static final field INSTANCE Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators;
-	public final fun getMax ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
-	public final fun getMean ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch;
-	public final fun getMedian ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator;
-	public final fun getMin ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
-	public final fun getPercentile ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch;
+	public final fun getMax ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator;
+	public final fun getMedian ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator;
+	public final fun getMin ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator;
+	public final fun getPercentile ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
 	public final fun getStd ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch2;
-	public final fun getSum ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator;
+	public final fun getSum ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator;
+}
+
+public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators$mean {
+	public static final field INSTANCE Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators$mean;
+	public final fun getToBigDecimal ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator;
+	public final fun getToDouble ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
+	public final fun getToNumber ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/NoAggregationKt {
@@ -10807,8 +10922,8 @@ public final class org/jetbrains/kotlinx/dataframe/jupyter/RenderedContent$Compa
 }
 
 public final class org/jetbrains/kotlinx/dataframe/math/MeanKt {
-	public static final fun mean (Ljava/lang/Iterable;Lkotlin/reflect/KType;Z)D
-	public static synthetic fun mean$default (Ljava/lang/Iterable;Lkotlin/reflect/KType;ZILjava/lang/Object;)D
+	public static final fun meanOrNull (Ljava/lang/Iterable;Lkotlin/reflect/KType;Z)Ljava/lang/Number;
+	public static synthetic fun meanOrNull$default (Ljava/lang/Iterable;Lkotlin/reflect/KType;ZILjava/lang/Object;)Ljava/lang/Number;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/math/PercentileKt {
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
index 9dc1a5b1c7..7c3c2e0c95 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
@@ -25,7 +25,7 @@ public interface ColumnDescription {
     public val nulls: Int
     public val top: Any
     public val freq: Int
-    public val mean: Double
+    public val mean: Number?
     public val std: Double
     public val min: Any
     public val p25: Any
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
index 63fc249950..d7dcc7f048 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
@@ -1,3 +1,5 @@
+@file:OptIn(ExperimentalTypeInference::class)
+
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyRow
@@ -20,11 +22,12 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.numberColumns
 import org.jetbrains.kotlinx.dataframe.impl.columns.toNumberColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull
-import org.jetbrains.kotlinx.dataframe.math.mean
+import org.jetbrains.kotlinx.dataframe.math.meanOrNull
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.experimental.ExperimentalTypeInference
 import kotlin.reflect.KProperty
+import kotlin.reflect.full.isSubtypeOf
 import kotlin.reflect.typeOf
 
 // region DataColumn
@@ -56,7 +59,7 @@ public fun DataColumn<BigInteger?>.mean(): BigDecimal = meanOrNull().suggestIfNu
 public fun DataColumn<BigDecimal?>.mean(): BigDecimal = meanOrNull().suggestIfNull("mean")
 
 @JvmName("meanNumber")
-public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Number? = meanOrNull(skipNA)
+public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Number = meanOrNull(skipNA).suggestIfNull("mean")
 
 // endregion
 
@@ -96,112 +99,139 @@ public fun DataColumn<Number?>.meanOrNull(skipNA: Boolean = skipNA_default): Num
 
 // region meanOf
 
-@OptIn(ExperimentalTypeInference::class)
 @JvmName("meanOfInt")
-//@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Int?): Double =
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Int?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfShort")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Short?): Double =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfByte")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Byte?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfLong")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> Long?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfDouble")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfFloat")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfBigInteger")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigInteger?): BigDecimal =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfBigDecimal")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigDecimal?): BigDecimal =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@JvmName("meanOfNumber")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+// endregion
+
+// region meanOfOrNull
+
+@JvmName("meanOfOrNullInt")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Int?): Double? =
     Aggregators.mean.toDouble(skipNA_default)
         .cast2<Int?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfShort")
+@JvmName("meanOfOrNullShort")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Short?): Double =
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Short?): Double? =
     Aggregators.mean.toDouble(skipNA_default)
         .cast2<Short?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfByte")
+@JvmName("meanOfOrNullByte")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Byte?): Double =
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Byte?): Double? =
     Aggregators.mean.toDouble(skipNA_default)
         .cast2<Byte?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfLong")
+@JvmName("meanOfOrNullLong")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Long?): Double =
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Long?): Double? =
     Aggregators.mean.toDouble(skipNA_default)
         .cast2<Long?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfDouble")
+@JvmName("meanOfOrNullDouble")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double =
+public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double? =
     Aggregators.mean.toDouble(skipNA)
         .cast2<Double?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfFloat")
+@JvmName("meanOfOrNullFloat")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double =
+public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double? =
     Aggregators.mean.toDouble(skipNA)
         .cast2<Float?, Double>()
         .aggregateOf(this, expression)
-        ?: Double.NaN
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfBigInteger")
+@JvmName("meanOfOrNullBigInteger")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigInteger?): BigDecimal? =
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> BigInteger?): BigDecimal? =
     Aggregators.mean.toBigDecimal
         .cast2<BigInteger?, BigDecimal?>()
         .aggregateOf(this, expression)
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfBigDecimal")
+@JvmName("meanOfOrNullBigDecimal")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigDecimal?): BigDecimal? =
+public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> BigDecimal?): BigDecimal? =
     Aggregators.mean.toBigDecimal
         .cast2<BigDecimal?, BigDecimal?>()
         .aggregateOf(this, expression)
 
-@OptIn(ExperimentalTypeInference::class)
-@JvmName("meanOfNumber")
+@JvmName("meanOfOrNullNumber")
 @OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number? =
+public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number? =
     Aggregators.mean.toNumber(skipNA)
         .cast2<Number?, Number?>()
         .aggregateOf(this, expression)
 
-public fun main() {
-    val data = (1..10).toList()
-    val df = data.toDataFrame()
+// endregion
 
-    val mean = df.value.meanOf { if (true) it.toLong() else it.toDouble() }
-    val mean2 = df.value.meanOf { it.toBigInteger() }
+// endregion
 
-    println(mean)
-    println(mean!!::class)
-}
+// region DataRow - rowMean
 
-// endregion
+public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Number = rowMeanOrNull(skipNA).suggestIfNull("rowMean")
 
-// endregion
+public fun AnyRow.rowMeanOrNull(skipNA: Boolean = skipNA_default): Number? =
+    Aggregators.mean.toNumber(skipNA).aggregateCalculatingType(
+        values().filterIsInstance<Number>(),
+        columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
+    )
 
-// region DataRow
-// todo
-public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Double =
-    values().filterIsInstance<Number>().map { it.toDouble() }.mean(skipNA)
+public inline fun <reified T : Number> AnyRow.rowMeanOf(): Number = rowMeanOfOrNull<T>().suggestIfNull("rowMeanOf")
 
-public inline fun <reified T : Number> AnyRow.rowMeanOf(): Double =
-    values().filterIsInstance<T>().mean(typeOf<T>()) as Double
+public inline fun <reified T : Number> AnyRow.rowMeanOfOrNull(): Number? =
+    values().filterIsInstance<T>().meanOrNull(typeOf<T>())
 
 // endregion
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default): DataRow<T> = meanFor(skipNA, numberColumns())
+// region meanFor
 
 public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
@@ -223,29 +253,249 @@ public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
 ): DataRow<T> = meanFor(skipNA) { columns.toColumnSet() }
 
-// todo
+// endregion
+
+// region mean
+
+public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default): DataRow<T> = meanFor(skipNA, numberColumns())
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanInt")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Int?>): Double = meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanShort")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Short?>): Double =
+    meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanByte")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Byte?>): Double = meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanLong")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Long?>): Double = meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanDouble")
+public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Double?>): Double =
+    meanOrNull(skipNA, columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanFloat")
+public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Float?>): Double =
+    meanOrNull(skipNA, columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanBigInteger")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, BigInteger?>): BigDecimal =
+    meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanBigDecimal")
+public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, BigDecimal?>): BigDecimal =
+    meanOrNull(columns).suggestIfNull("mean")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanNumber")
 public fun <T, C : Number> DataFrame<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): Double = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns) as Double? ?: Double.NaN
+): Number = meanOrNull(skipNA, columns).suggestIfNull("mean")
 
-public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Double =
-    mean(skipNA) { columns.toNumberColumns() }
+public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Number =
+    meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
 
 @AccessApiOverload
 public fun <T, C : Number> DataFrame<T>.mean(
     vararg columns: ColumnReference<C?>,
     skipNA: Boolean = skipNA_default,
-): Double = mean(skipNA) { columns.toColumnSet() }
+): Number = meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
 
 @AccessApiOverload
-public fun <T, C : Number> DataFrame<T>.mean(vararg columns: KProperty<C?>, skipNA: Boolean = skipNA_default): Double =
-    mean(skipNA) { columns.toColumnSet() }
+public fun <T, C : Number> DataFrame<T>.mean(vararg columns: KProperty<C?>, skipNA: Boolean = skipNA_default): Number =
+    meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
 
+// endregion
+
+// region meanOrNull
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullInt")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Int?>): Double? =
+    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullShort")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Short?>): Double? =
+    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullByte")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Byte?>): Double? =
+    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullLong")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Long?>): Double? =
+    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullDouble")
+public fun <T> DataFrame<T>.meanOrNull(
+    skipNA: Boolean = skipNA_default,
+    columns: ColumnsSelector<T, Double?>,
+): Double? = Aggregators.mean.toDouble(skipNA).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullFloat")
+public fun <T> DataFrame<T>.meanOrNull(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Float?>): Double? =
+    Aggregators.mean.toDouble(skipNA).aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullBigInteger")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, BigInteger?>): BigDecimal? =
+    Aggregators.mean.toBigDecimal.aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullBigDecimal")
+public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, BigDecimal?>): BigDecimal? =
+    Aggregators.mean.toBigDecimal.aggregateAll(this, columns)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullNumber")
+public fun <T, C : Number> DataFrame<T>.meanOrNull(
+    skipNA: Boolean = skipNA_default,
+    columns: ColumnsSelector<T, C?>,
+): Number? = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns)
+
+public fun <T> DataFrame<T>.meanOrNull(vararg columns: String, skipNA: Boolean = skipNA_default): Number? =
+    meanOrNull(skipNA) { columns.toNumberColumns() }
+
+@AccessApiOverload
+public fun <T, C : Number> DataFrame<T>.meanOrNull(
+    vararg columns: ColumnReference<C?>,
+    skipNA: Boolean = skipNA_default,
+): Number? = meanOrNull(skipNA) { columns.toColumnSet() }
+
+@AccessApiOverload
+public fun <T, C : Number> DataFrame<T>.meanOrNull(
+    vararg columns: KProperty<C?>,
+    skipNA: Boolean = skipNA_default,
+): Number? = meanOrNull(skipNA) { columns.toColumnSet() }
+
+// endregion
+
+// region meanOf
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfInt")
+public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, Int?>): Double =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfShort")
+public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Short?>): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfByte")
+public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Byte?>): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfLong")
+public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Long?>): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfDouble")
+public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Double?>): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfFloat")
+public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Float?>): Double =
+    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfBigInteger")
+public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, BigInteger?>): BigDecimal =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfBigDecimal")
+public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, BigDecimal?>): BigDecimal =
+    meanOfOrNull(expression).suggestIfNull("meanOf")
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOfNumber")
 public inline fun <T, reified D : Number> DataFrame<T>.meanOf(
     skipNA: Boolean = skipNA_default,
     noinline expression: RowExpression<T, D?>,
-): Double = Aggregators.mean.toNumber(skipNA).of(this, expression) as Double? ?: Double.NaN
+): Number = meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
+
+// endregion
+
+// region meanOfOrNull
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfInt")
+public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, Int?>): Double? =
+    Aggregators.mean.toDouble(skipNA_default).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfShort")
+public fun <T> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    expression: RowExpression<T, Short?>,
+): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfByte")
+public fun <T> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    expression: RowExpression<T, Byte?>,
+): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfLong")
+public fun <T> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    expression: RowExpression<T, Long?>,
+): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfDouble")
+public fun <T> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    expression: RowExpression<T, Double?>,
+): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfFloat")
+public fun <T> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    expression: RowExpression<T, Float?>,
+): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfBigInteger")
+public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, BigInteger?>): BigDecimal? =
+    Aggregators.mean.toBigDecimal.of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfBigDecimal")
+public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, BigDecimal?>): BigDecimal? =
+    Aggregators.mean.toBigDecimal.of(this, expression)
+
+@OverloadResolutionByLambdaReturnType
+@JvmName("meanOrNullOfNumber")
+public inline fun <T, reified D : Number> DataFrame<T>.meanOfOrNull(
+    skipNA: Boolean = skipNA_default,
+    noinline expression: RowExpression<T, D?>,
+): Number? = Aggregators.mean.toNumber(skipNA).of(this, expression)
+
+// endregion
 
 // endregion
 
@@ -303,7 +553,9 @@ public inline fun <T, reified R : Number> Grouped<T>.meanOf(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateOf(this, name, expression)
+): DataFrame<T> =
+    Aggregators.mean.toNumber(skipNA)
+        .aggregateOf(this, name, expression)
 
 // endregion
 
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
index ffb1815ddb..fa545db4f2 100644
--- 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,10 +1,11 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
-import org.jetbrains.kotlinx.dataframe.math.mean
+import org.jetbrains.kotlinx.dataframe.math.meanOrNull
 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
 
 @PublishedApi
@@ -86,21 +87,18 @@ internal object Aggregators {
 
     @Suppress("ClassName")
     object mean {
-        val toNumber = withOption { skipNA: Boolean ->
-            extendsNumbers { mean(it, skipNA) }
-        }.create("meanToNumber")
-
-        val toDouble = withOption { skipNA: Boolean ->
-            changesType(
-                aggregateWithType = { mean(it, skipNA).asDoubleOrNaN() },
-                aggregateWithValues = { mean(skipNA) },
-            )
-        }.create("meanToDouble")
-
-        val toBigDecimal = changesType(
-            aggregateWithType = { mean(it) as BigDecimal? },
-            aggregateWithValues = { filterNotNull().mean() },
-        ).create("meanToBigDecimal")
+        val toNumber = withOneOption { skipNA: Boolean ->
+            twoStepForNumbers { meanOrNull(it, skipNA) }
+        }.create(mean::class.simpleName!!)
+
+        val toDouble = withOneOption { skipNA: Boolean ->
+            twoStepForNumbers { meanOrNull(it, skipNA) as Double? }
+        }.create(mean::class.simpleName!!)
+
+        val toBigDecimal =
+            twoStepForNumbers {
+                meanOrNull(it) as BigDecimal?
+            }.create(mean::class.simpleName!!)
     }
 
     val percentile by withOneOption { percentile: Double ->
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
index 74cfa32618..b5dbdab50f 100644
--- 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
@@ -42,8 +42,8 @@ internal class TwoStepNumbersAggregator<Return : Number>(
 ) : AggregatorBase<Number, Return>(name, aggregator) {
 
     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"
+        require(type.isSubtypeOf(typeOf<Number?>())) {
+            "${TwoStepNumbersAggregator::class.simpleName}: Type $type is not a subtype of Number?"
         }
         return super.aggregate(values, type)
     }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
index 4d43fb6128..299587a3b9 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
@@ -30,7 +30,7 @@ internal inline fun <C, reified V, R> Aggregator<V, R>.aggregateOf(
 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) } // TODO: inline
+): R? = (this as Aggregator<C, R>).aggregateOf(frame.rows()) { expression(it, it) }
 
 @PublishedApi
 internal fun <T, C, R> Aggregator<*, R>.aggregateOfDelegated(
@@ -50,7 +50,7 @@ internal inline fun <T, reified C, R> Aggregator<*, R>.of(
 
 @PublishedApi
 internal inline fun <C, reified V, R> Aggregator<V, R>.of(data: DataColumn<C>, crossinline expression: (C) -> V): R? =
-    aggregateOf(data.values()) { expression(it) } // TODO: inline
+    aggregateOf(data.values()) { expression(it) }
 
 @PublishedApi
 internal inline fun <T, reified C, reified R> Aggregator<*, R>.aggregateOf(
@@ -75,7 +75,8 @@ internal inline fun <T, reified C, reified R> Grouped<T>.aggregateOf(
     val type = typeOf<R>()
     return aggregateInternal {
         val value = aggregator.aggregateOf(df, expression)
-        yield(path, value, type, null, false)
+        val inferType = !aggregator.preservesType
+        yield(path, value, type, null, inferType)
     }
 }
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
index 254a598a86..1b734cf8f9 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
@@ -16,6 +16,7 @@ import org.jetbrains.kotlinx.dataframe.api.isNumber
 import org.jetbrains.kotlinx.dataframe.api.map
 import org.jetbrains.kotlinx.dataframe.api.maxOrNull
 import org.jetbrains.kotlinx.dataframe.api.mean
+import org.jetbrains.kotlinx.dataframe.api.meanOrNull
 import org.jetbrains.kotlinx.dataframe.api.medianOrNull
 import org.jetbrains.kotlinx.dataframe.api.minOrNull
 import org.jetbrains.kotlinx.dataframe.api.move
@@ -56,7 +57,7 @@ internal fun describeImpl(cols: List<AnyCol>): DataFrame<ColumnDescription> {
                 ?.key
         }
         if (hasNumericCols) {
-            ColumnDescription::mean from { if (it.isNumber()) it.asNumbers().mean() else null }
+            ColumnDescription::mean from { if (it.isNumber()) it.asNumbers().meanOrNull() else null }
             ColumnDescription::std from { if (it.isNumber()) it.asNumbers().std() else null }
         }
         if (hasComparableCols || hasNumericCols) {
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
index 3786feb7be..36c67ac161 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -1,6 +1,5 @@
 package org.jetbrains.kotlinx.dataframe.math
 
-import org.jetbrains.kotlinx.dataframe.api.isNaN
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
 import org.jetbrains.kotlinx.dataframe.impl.api.toBigDecimal
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
@@ -13,10 +12,10 @@ import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
-/** @include [Sequence.mean] */
+/** @include [Sequence.meanOrNull] */
 @PublishedApi
-internal fun <T : Number> Iterable<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Number? =
-    asSequence().mean(type, skipNA)
+internal fun <T : Number> Iterable<T>.meanOrNull(type: KType, skipNA: Boolean = skipNA_default): Number? =
+    asSequence().meanOrNull(type, skipNA)
 
 /**
  * Returns the mean of the numbers in [this].
@@ -24,41 +23,42 @@ internal fun <T : Number> Iterable<T>.mean(type: KType, skipNA: Boolean = skipNA
  * If the input is empty, the return value will be `null`.
  *
  * If the [type] given or input consists of only [Int], [Short], [Byte], [Long], [Double], or [Float],
- * the return type will be [Double]`?` (Never `NaN`).
+ * the return type will be [Double].
  *
- * If the [type] given or the input contains [BigInteger] or [BigDecimal], the return type will be [BigDecimal]`?`.
+ * If the [type] given or the input contains [BigInteger] or [BigDecimal],
+ * the return type will be [BigDecimal].
  * @param type The type of the numbers in the sequence.
  * @param skipNA Whether to skip `NaN` values (default: `false`). Only relevant for [Double] and [Float].
  */
 @Suppress("UNCHECKED_CAST")
-internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Number? {
+internal fun <T : Number> Sequence<T>.meanOrNull(type: KType, skipNA: Boolean = skipNA_default): Number? {
     if (type.isMarkedNullable) {
-        return filterNotNull().mean(type.withNullability(false), skipNA)
+        return filterNotNull().meanOrNull(type.withNullability(false), skipNA)
     }
     return when (type.classifier) {
-        // Double -> Double?
-        Double::class -> (this as Sequence<Double>).mean(skipNA).takeUnless { it.isNaN }
+        // Double -> Double
+        Double::class -> (this as Sequence<Double>).meanOrNull(skipNA)
 
-        // Float -> Double?
-        Float::class -> (this as Sequence<Float>).mean(skipNA).takeUnless { it.isNaN }
+        // Float -> Double
+        Float::class -> (this as Sequence<Float>).meanOrNull(skipNA)
 
-        // Int -> Double?
-        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
+        // Int -> Double
+        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.meanOrNull(false)
 
-        // Short -> Double?
-        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
+        // Short -> Double
+        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.meanOrNull(false)
 
-        // Byte -> Double?
-        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
+        // Byte -> Double
+        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.meanOrNull(false)
 
-        // Long -> Double?
-        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.mean(false).takeUnless { it.isNaN }
+        // Long -> Double
+        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.meanOrNull(false)
 
-        // BigInteger -> BigDecimal?
-        BigInteger::class -> (this as Sequence<BigInteger>).mean()
+        // BigInteger -> BigDecimal
+        BigInteger::class -> (this as Sequence<BigInteger>).meanOrNull()
 
-        // BigDecimal -> BigDecimal?
-        BigDecimal::class -> (this as Sequence<BigDecimal>).mean()
+        // BigDecimal -> BigDecimal
+        BigDecimal::class -> (this as Sequence<BigDecimal>).meanOrNull()
 
         // Number -> Conversion(Common number type) -> Number? (Double or BigDecimal?)
         // fallback case, heavy as it needs to collect all types at runtime
@@ -69,7 +69,7 @@ internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA
                 error("Cannot find unified number type for $numberTypes")
             }
             this.convertToUnifiedNumberType(unifiedType)
-                .mean(unifiedType, skipNA)
+                .meanOrNull(unifiedType, skipNA)
         }
 
         // this means the sequence is empty
@@ -79,7 +79,7 @@ internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA
     }
 }
 
-internal fun Sequence<Double>.mean(skipNA: Boolean = skipNA_default): Double {
+internal fun Sequence<Double>.meanOrNull(skipNA: Boolean = skipNA_default): Double? {
     var count = 0
     var sum: Double = 0.toDouble()
     for (element in this) {
@@ -87,17 +87,17 @@ internal fun Sequence<Double>.mean(skipNA: Boolean = skipNA_default): Double {
             if (skipNA) {
                 continue
             } else {
-                return Double.NaN
+                return null
             }
         }
         sum += element
         count++
     }
-    return if (count > 0) sum / count else Double.NaN
+    return if (count > 0) sum / count else null
 }
 
 @JvmName("meanFloat")
-internal fun Sequence<Float>.mean(skipNA: Boolean = skipNA_default): Double {
+internal fun Sequence<Float>.meanOrNull(skipNA: Boolean = skipNA_default): Double? {
     var count = 0
     var sum: Double = 0.toDouble()
     for (element in this) {
@@ -105,17 +105,17 @@ internal fun Sequence<Float>.mean(skipNA: Boolean = skipNA_default): Double {
             if (skipNA) {
                 continue
             } else {
-                return Double.NaN
+                return null
             }
         }
         sum += element
         count++
     }
-    return if (count > 0) sum / count else Double.NaN
+    return if (count > 0) sum / count else null
 }
 
 @JvmName("bigIntegerMean")
-internal fun Sequence<BigInteger>.mean(): BigDecimal? {
+internal fun Sequence<BigInteger>.meanOrNull(): BigDecimal? {
     var count = 0
     val sum = sumOf {
         count++
@@ -125,7 +125,7 @@ internal fun Sequence<BigInteger>.mean(): BigDecimal? {
 }
 
 @JvmName("bigDecimalMean")
-internal fun Sequence<BigDecimal>.mean(): BigDecimal? {
+internal fun Sequence<BigDecimal>.meanOrNull(): BigDecimal? {
     var count = 0
     val sum = sumOf {
         count++
@@ -135,65 +135,65 @@ internal fun Sequence<BigDecimal>.mean(): BigDecimal? {
 }
 
 @JvmName("doubleMean")
-internal fun Iterable<Double>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
+internal fun Iterable<Double>.meanOrNull(skipNA: Boolean = skipNA_default): Double? = asSequence().meanOrNull(skipNA)
 
 @JvmName("floatMean")
-internal fun Iterable<Float>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
+internal fun Iterable<Float>.meanOrNull(skipNA: Boolean = skipNA_default): Double? = asSequence().meanOrNull(skipNA)
 
 @JvmName("bigDecimalMean")
-internal fun Iterable<BigDecimal>.mean(): BigDecimal? = asSequence().mean()
+internal fun Iterable<BigDecimal>.meanOrNull(): BigDecimal? = asSequence().meanOrNull()
 
 @JvmName("bigIntegerMean")
-internal fun Iterable<BigInteger>.mean(): BigDecimal? = asSequence().mean()
+internal fun Iterable<BigInteger>.meanOrNull(): BigDecimal? = asSequence().meanOrNull()
 
 @JvmName("intMean")
-internal fun Iterable<Int>.mean(): Double =
+internal fun Iterable<Int>.meanOrNull(): Double? =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+        if (size > 0) sumOf { it.toDouble() } / size else null
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else Double.NaN
+        if (count > 0) sum / count else null
     }
 
 @JvmName("shortMean")
-internal fun Iterable<Short>.mean(): Double =
+internal fun Iterable<Short>.meanOrNull(): Double? =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+        if (size > 0) sumOf { it.toDouble() } / size else null
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else Double.NaN
+        if (count > 0) sum / count else null
     }
 
 @JvmName("byteMean")
-internal fun Iterable<Byte>.mean(): Double =
+internal fun Iterable<Byte>.meanOrNull(): Double? =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+        if (size > 0) sumOf { it.toDouble() } / size else null
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else Double.NaN
+        if (count > 0) sum / count else null
     }
 
 @JvmName("longMean")
-internal fun Iterable<Long>.mean(): Double =
+internal fun Iterable<Long>.meanOrNull(): Double? =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+        if (size > 0) sumOf { it.toDouble() } / size else null
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else Double.NaN
+        if (count > 0) sum / count else null
     }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
index 04694ad901..188c6673f9 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
@@ -39,7 +39,7 @@ class DescribeTests {
             nulls shouldBe 1
             top shouldBe 1
             freq shouldBe 1
-            mean shouldBe 4.5
+            mean shouldBe 4.5.toBigDecimal()
             std shouldBe 2.449489742783178
             min shouldBe 1.toBigDecimal()
             (p25 as BigDecimal).setScale(2) shouldBe 2.75.toBigDecimal()
@@ -64,8 +64,8 @@ class DescribeTests {
             nulls shouldBe 0
             top shouldBe 1
             freq shouldBe 1
-            mean.isNaN() shouldBe true
-            std.isNaN() shouldBe true
+            mean shouldBe null
+            std.isNaN shouldBe true
             min shouldBe 1.0 // TODO should be NaN too?
             p25 shouldBe 1.75
             median shouldBe 3.0
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
index 4241f7570a..71e55b21d6 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
@@ -142,9 +142,9 @@ class BasicTests {
     @Test
     fun `calculate mean age for each animal`() {
         val expected = dataFrameOf("animal", "age")(
-            "cat", Double.NaN,
+            "cat", null,
             "snake", 2.5,
-            "dog", Double.NaN,
+            "dog", null,
         )
 
         df.groupBy { animal }.mean { age } shouldBe expected
@@ -213,7 +213,7 @@ class BasicTests {
         val expected = dataFrameOf("animal", "1", "3", "2")(
             "cat", 2.5, 2.5, null,
             "snake", 4.5, null, 0.5,
-            "dog", 3.0, Double.NaN, 6.0,
+            "dog", 3.0, null, 6.0,
         )
 
         val actualDfAcc = df.pivot(inward = false) { visits }.groupBy { animal }.mean(skipNA = true) { age }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
index 72fd1dc4dc..338be0fd59 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
@@ -66,7 +66,7 @@ class MediumTests {
             -1, 0, 1,
         )
 
-        df.convert { colsOf<Double>() }.with { (it - rowMean()).roundToInt() } shouldBe expected
+        df.convert { colsOf<Double>() }.with { (it - rowMean().toDouble()).roundToInt() } shouldBe expected
     }
 
     @Test
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
index 0c0c156f26..80c571e95a 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
@@ -2,9 +2,11 @@ package org.jetbrains.kotlinx.dataframe.statistics
 
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.columnOf
 import org.jetbrains.kotlinx.dataframe.api.mean
-import org.jetbrains.kotlinx.dataframe.impl.nothingType
+import org.jetbrains.kotlinx.dataframe.api.meanOrNull
+import org.jetbrains.kotlinx.dataframe.impl.nullableNothingType
 import org.junit.Test
 import kotlin.reflect.typeOf
 
@@ -18,10 +20,17 @@ class BasicMathTests {
 
     @Test
     fun `mean with nans and nulls`() {
-        columnOf(10, 20, Double.NaN, null).mean() shouldBe Double.NaN
+        columnOf(10, 20, Double.NaN, null).meanOrNull() shouldBe null
         columnOf(10, 20, Double.NaN, null).mean(skipNA = true) shouldBe 15
 
-        DataColumn.createValueColumn("", emptyList<Nothing>(), nothingType(false)).mean() shouldBe Double.NaN
-        DataColumn.createValueColumn("", listOf(null), nothingType(true)).mean() shouldBe Double.NaN
+        DataColumn.createValueColumn("", emptyList<Nothing>(), nullableNothingType)
+            .cast<Number?>()
+            .meanOrNull() shouldBe null
+        DataColumn.createValueColumn("", emptyList<Nothing>(), typeOf<Double?>())
+            .cast<Double?>()
+            .meanOrNull() shouldBe null
+        DataColumn.createValueColumn("", listOf(null), typeOf<Double?>())
+            .cast<Double?>()
+            .meanOrNull() shouldBe null
     }
 }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
index f8ee4b37d5..3283d4e9d3 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
@@ -33,7 +33,7 @@ class AnimalsTests {
         mean.columnsCount() shouldBe 2
         mean.rowsCount() shouldBe 2
         mean.name.values() shouldBe listOf("age", "visits")
-        mean.value.type() shouldBe typeOf<Double>()
+        mean.value.type() shouldBe typeOf<Double?>()
     }
 
     @Test
@@ -42,7 +42,7 @@ class AnimalsTests {
             .update { age }.with { Double.NaN }
             .update { visits }.withNull()
         val mean = cleared.mean()
-        mean[age] shouldBe Double.NaN
-        (mean[visits.name()] as Double).isNaN() shouldBe true
+        mean[age] shouldBe null
+        mean[visits.name()] shouldBe null
     }
 }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
index 5bc715e078..831d826d71 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.testSets.person
 
 import io.kotest.assertions.throwables.shouldThrow
+import io.kotest.matchers.collections.shouldBeIn
 import io.kotest.matchers.doubles.ToleranceMatcher
 import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
@@ -122,6 +123,7 @@ import org.jetbrains.kotlinx.dataframe.api.rename
 import org.jetbrains.kotlinx.dataframe.api.reorderColumnsByName
 import org.jetbrains.kotlinx.dataframe.api.replace
 import org.jetbrains.kotlinx.dataframe.api.rows
+import org.jetbrains.kotlinx.dataframe.api.schema
 import org.jetbrains.kotlinx.dataframe.api.select
 import org.jetbrains.kotlinx.dataframe.api.single
 import org.jetbrains.kotlinx.dataframe.api.sortBy
@@ -177,11 +179,12 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.isMissingColumn
 import org.jetbrains.kotlinx.dataframe.impl.emptyPath
 import org.jetbrains.kotlinx.dataframe.impl.getColumnsImpl
 import org.jetbrains.kotlinx.dataframe.impl.nothingType
+import org.jetbrains.kotlinx.dataframe.impl.nullableNothingType
 import org.jetbrains.kotlinx.dataframe.impl.trackColumnAccess
 import org.jetbrains.kotlinx.dataframe.index
 import org.jetbrains.kotlinx.dataframe.io.renderValueForStdout
 import org.jetbrains.kotlinx.dataframe.kind
-import org.jetbrains.kotlinx.dataframe.math.mean
+import org.jetbrains.kotlinx.dataframe.math.meanOrNull
 import org.jetbrains.kotlinx.dataframe.name
 import org.jetbrains.kotlinx.dataframe.ncol
 import org.jetbrains.kotlinx.dataframe.nrow
@@ -817,8 +820,8 @@ class DataFrameTests : BaseTest() {
 
     @Test
     fun `groupBy meanOf`() {
-        typed.groupBy { name }.meanOf { age * 2 } shouldBe typed
-            .groupBy { name }.aggregate { mean { age } * 2 into "mean" }
+        typed.groupBy { name }.meanOf { age * 2 }.schema() shouldBe typed
+            .groupBy { name }.aggregate { mean { age } * 2 into "mean" }.schema()
     }
 
     @Test
@@ -1485,7 +1488,7 @@ class DataFrameTests : BaseTest() {
 
     @Test
     fun `column stats`() {
-        typed.age.mean() shouldBe typed.age.toList().mean()
+        typed.age.mean() shouldBe typed.age.toList().meanOrNull()
         typed.age.min() shouldBe typed.age.toList().minOrNull()
         typed.age.max() shouldBe typed.age.toList().maxOrNull()
         typed.age.sum() shouldBe typed.age.toList().sum()
@@ -2110,8 +2113,8 @@ class DataFrameTests : BaseTest() {
             it.kind() shouldBe ColumnKind.Group
             val group = it.asColumnGroup()
             group.columnNames() shouldBe listOf("age", "weight")
-            group.columns().forEach {
-                it.type() shouldBe typeOf<Double?>()
+            group.columnTypes().forEach {
+                it shouldBeIn setOf(typeOf<Double?>(), nullableNothingType)
             }
         }
     }

From 268d2387e00a9cff6d9dcdbe9289014a3a26a78a Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 3 Mar 2025 16:53:42 +0100
Subject: [PATCH 04/18] added calculateReturnTypeOrNull system to aggregators
 to avoid runtime value instance checks where we know the types already

---
 core/api/core.api                             |  1 +
 .../aggregation/aggregators/Aggregator.kt     | 10 +++++
 .../aggregation/aggregators/AggregatorBase.kt |  4 ++
 .../aggregation/aggregators/Aggregators.kt    | 45 +++++++++++++------
 .../aggregators/FlatteningAggregator.kt       | 11 ++++-
 .../aggregators/TwoStepAggregator.kt          | 29 +++++++-----
 .../aggregators/TwoStepNumbersAggregator.kt   | 35 +++++++++++----
 .../jetbrains/kotlinx/dataframe/math/mean.kt  | 23 ++++++++++
 8 files changed, 123 insertions(+), 35 deletions(-)

diff --git a/core/api/core.api b/core/api/core.api
index 7409c6c769..a84a99fa79 100644
--- 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
index bc5b82b0ce..6bb788e694 100644
--- 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,6 +57,11 @@ 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
@@ -63,3 +69,7 @@ internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Ag
 
 @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
index 0683a5e9ec..3860eedf5b 100644
--- 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
index fa545db4f2..09c14f4a95 100644
--- 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
index 4561ac4991..96e722f05f 100644
--- 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
index b42597ea8a..a000e6e662 100644
--- 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,10 +39,11 @@ 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.
@@ -49,17 +51,18 @@ internal class TwoStepAggregator<Value, Return>(
      * 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
index b5dbdab50f..be5fd22e1a 100644
--- 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
index 36c67ac161..6f95432fa9 100644
--- 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()

From c72335f6db66124aabb1c877efea0a8aa87a56f7 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Wed, 5 Mar 2025 13:39:02 +0100
Subject: [PATCH 05/18] renamed interComparable to intraComparable. Language is
 hard

---
 .../kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt  |  8 ++++----
 .../org/jetbrains/kotlinx/dataframe/api/median.kt      | 10 +++++-----
 .../kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt  |  8 ++++----
 .../org/jetbrains/kotlinx/dataframe/api/percentile.kt  | 10 +++++-----
 .../kotlinx/dataframe/impl/aggregation/getColumns.kt   |  2 +-
 5 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
index a265a20f9e..fc58dc626c 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMaxOf(): T = rowMaxOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(interComparableColumns())
+public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.max.aggregateFor(this, columns)
@@ -134,7 +134,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.maxByOrNull(column: KProperty<C?>
 
 // region GroupBy
 
-public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(interComparableColumns())
+public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.max.aggregateFor(this, columns)
@@ -246,7 +246,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.maxBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.maxFor(
     separate: Boolean = false,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index f939609362..b69f76b18c 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -12,7 +12,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMedianOf(): T =
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(interComparableColumns())
+public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.median.aggregateFor(this, columns)
@@ -105,7 +105,7 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.medianOf(
 
 // region GroupBy
 
-public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(interComparableColumns())
+public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.median.aggregateFor(this, columns)
@@ -147,7 +147,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.medianOf(
 
 // region Pivot
 
-public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, interComparableColumns())
+public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Pivot<T>.medianFor(
     separate: Boolean = false,
@@ -191,7 +191,7 @@ public inline fun <T, reified R : Comparable<R>> Pivot<T>.medianOf(
 // region PivotGroupBy
 
 public fun <T> PivotGroupBy<T>.median(separate: Boolean = false): DataFrame<T> =
-    medianFor(separate, interComparableColumns())
+    medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> PivotGroupBy<T>.medianFor(
     separate: Boolean = false,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
index d1cae852aa..01f74ce41e 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMinOf(): T = rowMinOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(interComparableColumns())
+public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.min.aggregateFor(this, columns)
@@ -134,7 +134,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.minByOrNull(column: KProperty<C?>
 
 // region GroupBy
 
-public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(interComparableColumns())
+public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.min.aggregateFor(this, columns)
@@ -247,7 +247,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.minBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.minFor(
     separate: Boolean = false,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
index 34c482612d..1ced2969e5 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
@@ -12,7 +12,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
@@ -63,7 +63,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowPercentileOf(percentile:
 // region DataFrame
 
 public fun <T> DataFrame<T>.percentile(percentile: Double): DataRow<T> =
-    percentileFor(percentile, interComparableColumns())
+    percentileFor(percentile, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.percentileFor(
     percentile: Double,
@@ -128,7 +128,7 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.percentileOf(
 // region GroupBy
 
 public fun <T> Grouped<T>.percentile(percentile: Double): DataFrame<T> =
-    percentileFor(percentile, interComparableColumns())
+    percentileFor(percentile, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.percentileFor(
     percentile: Double,
@@ -184,7 +184,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.percentileOf(
 // region Pivot
 
 public fun <T> Pivot<T>.percentile(percentile: Double, separate: Boolean = false): DataRow<T> =
-    percentileFor(percentile, separate, interComparableColumns())
+    percentileFor(percentile, separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Pivot<T>.percentileFor(
     percentile: Double,
@@ -238,7 +238,7 @@ public inline fun <T, reified R : Comparable<R>> Pivot<T>.percentileOf(
 // region PivotGroupBy
 
 public fun <T> PivotGroupBy<T>.percentile(percentile: Double, separate: Boolean = false): DataFrame<T> =
-    percentileFor(percentile, separate, interComparableColumns())
+    percentileFor(percentile, separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> PivotGroupBy<T>.percentileFor(
     percentile: Double,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
index 6f514d95eb..b7b2c1052d 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
@@ -14,7 +14,7 @@ internal inline fun <T> Aggregatable<T>.remainingColumns(
     crossinline predicate: (AnyCol) -> Boolean,
 ): ColumnsSelector<T, Any?> = remainingColumnsSelector().filter { predicate(it.data) }
 
-internal fun <T> Aggregatable<T>.interComparableColumns() =
+internal fun <T> Aggregatable<T>.intraComparableColumns() =
     remainingColumns { it.valuesAreComparable() } as ColumnsSelector<T, Comparable<Any?>>
 
 internal fun <T> Aggregatable<T>.numberColumns() = remainingColumns { it.isNumber() } as ColumnsSelector<T, Number?>

From 2df7d557000ceaddd3161b3604b96195614aecb6 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Wed, 5 Mar 2025 19:06:09 +0100
Subject: [PATCH 06/18] rollback of changes to mean

---
 core/api/core.api                             | 148 +-----
 .../kotlinx/dataframe/api/describe.kt         |   2 +-
 .../jetbrains/kotlinx/dataframe/api/mean.kt   | 460 ++----------------
 .../kotlinx/dataframe/impl/ExceptionUtils.kt  |  15 -
 .../kotlinx/dataframe/impl/TypeUtils.kt       |  10 -
 .../aggregation/aggregators/Aggregator.kt     |   2 +-
 .../aggregation/aggregators/AggregatorBase.kt |   2 +-
 .../aggregators/AggregatorOptionSwitch.kt     |   8 +-
 .../aggregators/AggregatorProvider.kt         |   4 +-
 .../aggregation/aggregators/Aggregators.kt    |  56 ++-
 .../aggregators/FlatteningAggregator.kt       |   4 +-
 .../aggregators/TwoStepAggregator.kt          |   4 +-
 .../aggregators/TwoStepNumbersAggregator.kt   |   4 +-
 .../kotlinx/dataframe/impl/api/describe.kt    |   3 +-
 .../jetbrains/kotlinx/dataframe/math/mean.kt  | 185 +++----
 .../kotlinx/dataframe/api/describe.kt         |   7 +-
 .../kotlinx/dataframe/puzzles/BasicTests.kt   |   6 +-
 .../kotlinx/dataframe/puzzles/MediumTests.kt  |   2 +-
 .../dataframe/statistics/BasicMathTests.kt    |  17 +-
 .../testSets/animals/AnimalsTests.kt          |   6 +-
 .../testSets/person/DataFrameTests.kt         |  15 +-
 .../dataframe/examples/titanic/ml/titanic.kt  |   7 +-
 22 files changed, 184 insertions(+), 783 deletions(-)

diff --git a/core/api/core.api b/core/api/core.api
index a84a99fa79..900d7b03f5 100644
--- a/core/api/core.api
+++ b/core/api/core.api
@@ -2633,7 +2633,7 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/api/ColumnDescri
 	public abstract fun getCount ()I
 	public abstract fun getFreq ()I
 	public abstract fun getMax ()Ljava/lang/Object;
-	public abstract fun getMean ()Ljava/lang/Number;
+	public abstract fun getMean ()D
 	public abstract fun getMedian ()Ljava/lang/Object;
 	public abstract fun getMin ()Ljava/lang/Object;
 	public abstract fun getName ()Ljava/lang/String;
@@ -2655,7 +2655,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/ColumnDescription_Extensi
 	public static final fun ColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun ColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
-	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Number;
+	public static final fun ColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)D
 	public static final fun ColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun ColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun ColumnDescription_min (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
@@ -2685,7 +2685,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/ColumnDescription_Extensi
 	public static final fun NullableColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun NullableColumnDescription_max (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
-	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Number;
+	public static final fun NullableColumnDescription_mean (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Double;
 	public static final fun NullableColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun NullableColumnDescription_median (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Ljava/lang/Object;
 	public static final fun NullableColumnDescription_min (Lorg/jetbrains/kotlinx/dataframe/ColumnsScope;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
@@ -5995,10 +5995,12 @@ public final class org/jetbrains/kotlinx/dataframe/api/MaxKt {
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Ljava/lang/Number;
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Ljava/lang/Number;
-	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Ljava/lang/Number;
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)D
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)D
+	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)D
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Ljava/lang/String;ZLkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;[Ljava/lang/String;Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
@@ -6011,10 +6013,12 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun mean (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataRow;
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)D
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)D
+	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)D
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Ljava/lang/String;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;[Ljava/lang/String;Ljava/lang/String;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
@@ -6027,20 +6031,6 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun mean$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
-	public static final fun meanBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
-	public static final fun meanBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
-	public static final fun meanBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
-	public static final fun meanByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
-	public static final fun meanDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
-	public static final fun meanDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
-	public static synthetic fun meanDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)D
-	public static final fun meanFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)D
-	public static synthetic fun meanFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataRow;
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
 	public static final fun meanFor (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
@@ -6073,100 +6063,10 @@ public final class org/jetbrains/kotlinx/dataframe/api/MeanKt {
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Ljava/lang/String;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lkotlin/reflect/KProperty;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static synthetic fun meanFor$default (Lorg/jetbrains/kotlinx/dataframe/api/PivotGroupBy;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
-	public static final fun meanInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
-	public static final fun meanInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
-	public static final fun meanLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
-	public static final fun meanLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
-	public static final fun meanNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Number;
-	public static final fun meanNumber (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Number;
-	public static synthetic fun meanNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun meanNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun meanOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
-	public static final fun meanOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
-	public static final fun meanOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOfByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanOfByte$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)D
-	public static synthetic fun meanOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)D
-	public static synthetic fun meanOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanOfInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
-	public static final fun meanOfLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanOfLong$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanOfNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Number;
-	public static synthetic fun meanOfNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun meanOfOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
-	public static final fun meanOfOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/math/BigDecimal;
-	public static final fun meanOfOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static final fun meanOfOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static synthetic fun meanOfOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOfOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static synthetic fun meanOfOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOfOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static final fun meanOfOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static final fun meanOfOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;)Ljava/lang/Number;
-	public static synthetic fun meanOfOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZLkotlin/jvm/functions/Function1;ILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun meanOfOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)Ljava/lang/Double;
-	public static final fun meanOfShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Lkotlin/jvm/functions/Function1;)D
-	public static final fun meanOfShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)D
-	public static synthetic fun meanOfShort$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)D
-	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;Z)Ljava/lang/Number;
-	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;Z)Ljava/lang/Number;
-	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;Z)Ljava/lang/Number;
-	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lkotlin/reflect/KProperty;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun meanOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
-	public static final fun meanOrNullByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static final fun meanOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
-	public static final fun meanOrNullDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
-	public static final fun meanOrNullFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
-	public static final fun meanOrNullInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static final fun meanOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
-	public static final fun meanOrNullLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static final fun meanOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Number;
-	public static final fun meanOrNullNumber (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Number;
-	public static synthetic fun meanOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static synthetic fun meanOrNullNumber$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun meanOrNullOfBigDecimal (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullOfBigInteger (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/math/BigDecimal;
-	public static final fun meanOrNullOfByte (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullOfByte$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullOfDouble (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullOfDouble$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullOfFloat (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullOfFloat$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullOfInt (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static final fun meanOrNullOfLong (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullOfLong$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullOfShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static synthetic fun meanOrNullOfShort$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Double;
-	public static final fun meanOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Ljava/lang/Double;
-	public static final fun meanOrNullShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Ljava/lang/Double;
-	public static final fun meanShort (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)D
-	public static final fun meanShort (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)D
-	public static final fun rowMean (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)Ljava/lang/Number;
-	public static synthetic fun rowMean$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)Ljava/lang/Number;
-	public static final fun rowMeanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)Ljava/lang/Number;
-	public static synthetic fun rowMeanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun meanOrNull (Lorg/jetbrains/kotlinx/dataframe/DataColumn;Z)Ljava/lang/Double;
+	public static synthetic fun meanOrNull$default (Lorg/jetbrains/kotlinx/dataframe/DataColumn;ZILjava/lang/Object;)Ljava/lang/Double;
+	public static final fun rowMean (Lorg/jetbrains/kotlinx/dataframe/DataRow;Z)D
+	public static synthetic fun rowMean$default (Lorg/jetbrains/kotlinx/dataframe/DataRow;ZILjava/lang/Object;)D
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/MedianKt {
@@ -10030,8 +9930,6 @@ public final class org/jetbrains/kotlinx/dataframe/impl/DataFrameSize {
 
 public final class org/jetbrains/kotlinx/dataframe/impl/ExceptionUtilsKt {
 	public static final fun suggestIfNull (Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/Object;
-	public static final fun suggestIfNull (Ljava/math/BigDecimal;Ljava/lang/String;)Ljava/math/BigDecimal;
-	public static final fun suggestIfNull (Ljava/math/BigInteger;Ljava/lang/String;)Ljava/math/BigInteger;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/TypeUtilsKt {
@@ -10127,6 +10025,7 @@ public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators {
 	public static final field INSTANCE Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators;
 	public final fun getMax ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator;
+	public final fun getMean ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
 	public final fun getMedian ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator;
 	public final fun getMin ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator;
 	public final fun getPercentile ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
@@ -10134,13 +10033,6 @@ public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/
 	public final fun getSum ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators$mean {
-	public static final field INSTANCE Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators$mean;
-	public final fun getToBigDecimal ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator;
-	public final fun getToDouble ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
-	public final fun getToNumber ()Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch1;
-}
-
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/NoAggregationKt {
 	public static final fun aggregateValue (Lorg/jetbrains/kotlinx/dataframe/api/Grouped;Ljava/lang/String;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 }
@@ -10923,8 +10815,8 @@ public final class org/jetbrains/kotlinx/dataframe/jupyter/RenderedContent$Compa
 }
 
 public final class org/jetbrains/kotlinx/dataframe/math/MeanKt {
-	public static final fun meanOrNull (Ljava/lang/Iterable;Lkotlin/reflect/KType;Z)Ljava/lang/Number;
-	public static synthetic fun meanOrNull$default (Ljava/lang/Iterable;Lkotlin/reflect/KType;ZILjava/lang/Object;)Ljava/lang/Number;
+	public static final fun mean (Ljava/lang/Iterable;Lkotlin/reflect/KType;Z)D
+	public static synthetic fun mean$default (Ljava/lang/Iterable;Lkotlin/reflect/KType;ZILjava/lang/Object;)D
 }
 
 public final class org/jetbrains/kotlinx/dataframe/math/PercentileKt {
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
index 7c3c2e0c95..9dc1a5b1c7 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
@@ -25,7 +25,7 @@ public interface ColumnDescription {
     public val nulls: Int
     public val top: Any
     public val freq: Int
-    public val mean: Number?
+    public val mean: Double
     public val std: Double
     public val min: Any
     public val p25: Any
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
index d7dcc7f048..994cbf27db 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/mean.kt
@@ -1,5 +1,3 @@
-@file:OptIn(ExperimentalTypeInference::class)
-
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyRow
@@ -22,221 +20,42 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.numberColumns
 import org.jetbrains.kotlinx.dataframe.impl.columns.toNumberColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull
-import org.jetbrains.kotlinx.dataframe.math.meanOrNull
-import java.math.BigDecimal
-import java.math.BigInteger
-import kotlin.experimental.ExperimentalTypeInference
+import org.jetbrains.kotlinx.dataframe.math.mean
 import kotlin.reflect.KProperty
-import kotlin.reflect.full.isSubtypeOf
 import kotlin.reflect.typeOf
 
 // region DataColumn
 
-// region mean
-
-@JvmName("meanInt")
-public fun DataColumn<Int?>.mean(): Double = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanShort")
-public fun DataColumn<Short?>.mean(): Double = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanByte")
-public fun DataColumn<Byte?>.mean(): Double = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanLong")
-public fun DataColumn<Long?>.mean(): Double = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanDouble")
-public fun DataColumn<Double?>.mean(skipNA: Boolean = skipNA_default): Double = meanOrNull(skipNA).suggestIfNull("mean")
-
-@JvmName("meanFloat")
-public fun DataColumn<Float?>.mean(skipNA: Boolean = skipNA_default): Double = meanOrNull(skipNA).suggestIfNull("mean")
-
-@JvmName("meanBigInteger")
-public fun DataColumn<BigInteger?>.mean(): BigDecimal = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanBigDecimal")
-public fun DataColumn<BigDecimal?>.mean(): BigDecimal = meanOrNull().suggestIfNull("mean")
-
-@JvmName("meanNumber")
-public fun DataColumn<Number?>.mean(skipNA: Boolean = skipNA_default): Number = meanOrNull(skipNA).suggestIfNull("mean")
-
-// endregion
-
-// region meanOrNull
-
-@JvmName("meanOrNullInt")
-public fun DataColumn<Int?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
-
-@JvmName("meanOrNullShort")
-public fun DataColumn<Short?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
-
-@JvmName("meanOrNullByte")
-public fun DataColumn<Byte?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
-
-@JvmName("meanOrNullLong")
-public fun DataColumn<Long?>.meanOrNull(): Double? = Aggregators.mean.toDouble(skipNA_default).aggregate(this)
-
-@JvmName("meanOrNullDouble")
-public fun DataColumn<Double?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
-    Aggregators.mean.toDouble(skipNA).aggregate(this)
-
-@JvmName("meanOrNullFloat")
-public fun DataColumn<Float?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
-    Aggregators.mean.toDouble(skipNA).aggregate(this)
-
-@JvmName("meanOrNullBigInteger")
-public fun DataColumn<BigInteger?>.meanOrNull(): BigDecimal? = Aggregators.mean.toBigDecimal.aggregate(this)
-
-@JvmName("meanOrNullBigDecimal")
-public fun DataColumn<BigDecimal?>.meanOrNull(): BigDecimal? = Aggregators.mean.toBigDecimal.aggregate(this)
+public fun <T : Number> DataColumn<T?>.mean(skipNA: Boolean = skipNA_default): Double =
+    meanOrNull(skipNA).suggestIfNull("mean")
 
-@JvmName("meanOrNullNumber")
-public fun DataColumn<Number?>.meanOrNull(skipNA: Boolean = skipNA_default): Number? =
-    Aggregators.mean.toNumber(skipNA).aggregate(this)
+public fun <T : Number> DataColumn<T?>.meanOrNull(skipNA: Boolean = skipNA_default): Double? =
+    Aggregators.mean(skipNA).aggregate(this)
 
-// endregion
-
-// region meanOf
-
-@JvmName("meanOfInt")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Int?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfShort")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Short?): Double =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfByte")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Byte?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfLong")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> Long?): Double = meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfDouble")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfFloat")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfBigInteger")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigInteger?): BigDecimal =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfBigDecimal")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(expression: (T) -> BigDecimal?): BigDecimal =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@JvmName("meanOfNumber")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOf(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-// endregion
-
-// region meanOfOrNull
-
-@JvmName("meanOfOrNullInt")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Int?): Double? =
-    Aggregators.mean.toDouble(skipNA_default)
-        .cast2<Int?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullShort")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Short?): Double? =
-    Aggregators.mean.toDouble(skipNA_default)
-        .cast2<Short?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullByte")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Byte?): Double? =
-    Aggregators.mean.toDouble(skipNA_default)
-        .cast2<Byte?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullLong")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> Long?): Double? =
-    Aggregators.mean.toDouble(skipNA_default)
-        .cast2<Long?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullDouble")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Double?): Double? =
-    Aggregators.mean.toDouble(skipNA)
-        .cast2<Double?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullFloat")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Float?): Double? =
-    Aggregators.mean.toDouble(skipNA)
-        .cast2<Float?, Double>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullBigInteger")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> BigInteger?): BigDecimal? =
-    Aggregators.mean.toBigDecimal
-        .cast2<BigInteger?, BigDecimal?>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullBigDecimal")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(expression: (T) -> BigDecimal?): BigDecimal? =
-    Aggregators.mean.toBigDecimal
-        .cast2<BigDecimal?, BigDecimal?>()
-        .aggregateOf(this, expression)
-
-@JvmName("meanOfOrNullNumber")
-@OverloadResolutionByLambdaReturnType
-public fun <T> DataColumn<T>.meanOfOrNull(skipNA: Boolean = skipNA_default, expression: (T) -> Number?): Number? =
-    Aggregators.mean.toNumber(skipNA)
-        .cast2<Number?, Number?>()
-        .aggregateOf(this, expression)
-
-// endregion
+public inline fun <T, reified R : Number> DataColumn<T>.meanOf(
+    skipNA: Boolean = skipNA_default,
+    noinline expression: (T) -> R?,
+): Double = Aggregators.mean(skipNA).cast2<R?, Double>().aggregateOf(this, expression) ?: Double.NaN
 
 // endregion
 
-// region DataRow - rowMean
-
-public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Number = rowMeanOrNull(skipNA).suggestIfNull("rowMean")
+// region DataRow
 
-public fun AnyRow.rowMeanOrNull(skipNA: Boolean = skipNA_default): Number? =
-    Aggregators.mean.toNumber(skipNA).aggregateCalculatingType(
-        values().filterIsInstance<Number>(),
-        columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
-    )
+public fun AnyRow.rowMean(skipNA: Boolean = skipNA_default): Double =
+    values().filterIsInstance<Number>().map { it.toDouble() }.mean(skipNA)
 
-public inline fun <reified T : Number> AnyRow.rowMeanOf(): Number = rowMeanOfOrNull<T>().suggestIfNull("rowMeanOf")
-
-public inline fun <reified T : Number> AnyRow.rowMeanOfOrNull(): Number? =
-    values().filterIsInstance<T>().meanOrNull(typeOf<T>())
+public inline fun <reified T : Number> AnyRow.rowMeanOf(): Double = values().filterIsInstance<T>().mean(typeOf<T>())
 
 // endregion
 
 // region DataFrame
 
-// region meanFor
+public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default): DataRow<T> = meanFor(skipNA, numberColumns())
 
 public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataRow<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, columns)
+): DataRow<T> = Aggregators.mean(skipNA).aggregateFor(this, columns)
 
 public fun <T> DataFrame<T>.meanFor(vararg columns: String, skipNA: Boolean = skipNA_default): DataRow<T> =
     meanFor(skipNA) { columns.toNumberColumns() }
@@ -253,249 +72,28 @@ public fun <T, C : Number> DataFrame<T>.meanFor(
     skipNA: Boolean = skipNA_default,
 ): DataRow<T> = meanFor(skipNA) { columns.toColumnSet() }
 
-// endregion
-
-// region mean
-
-public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default): DataRow<T> = meanFor(skipNA, numberColumns())
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanInt")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Int?>): Double = meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanShort")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Short?>): Double =
-    meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanByte")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Byte?>): Double = meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanLong")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, Long?>): Double = meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanDouble")
-public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Double?>): Double =
-    meanOrNull(skipNA, columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanFloat")
-public fun <T> DataFrame<T>.mean(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Float?>): Double =
-    meanOrNull(skipNA, columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanBigInteger")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, BigInteger?>): BigDecimal =
-    meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanBigDecimal")
-public fun <T> DataFrame<T>.mean(columns: ColumnsSelector<T, BigDecimal?>): BigDecimal =
-    meanOrNull(columns).suggestIfNull("mean")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanNumber")
 public fun <T, C : Number> DataFrame<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): Number = meanOrNull(skipNA, columns).suggestIfNull("mean")
+): Double = Aggregators.mean(skipNA).aggregateAll(this, columns) as Double? ?: Double.NaN
 
-public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Number =
-    meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
+public fun <T> DataFrame<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): Double =
+    mean(skipNA) { columns.toNumberColumns() }
 
 @AccessApiOverload
 public fun <T, C : Number> DataFrame<T>.mean(
     vararg columns: ColumnReference<C?>,
     skipNA: Boolean = skipNA_default,
-): Number = meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
-
-@AccessApiOverload
-public fun <T, C : Number> DataFrame<T>.mean(vararg columns: KProperty<C?>, skipNA: Boolean = skipNA_default): Number =
-    meanOrNull(columns = columns, skipNA = skipNA).suggestIfNull("mean")
-
-// endregion
-
-// region meanOrNull
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullInt")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Int?>): Double? =
-    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullShort")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Short?>): Double? =
-    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullByte")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Byte?>): Double? =
-    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullLong")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, Long?>): Double? =
-    Aggregators.mean.toDouble(skipNA_default).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullDouble")
-public fun <T> DataFrame<T>.meanOrNull(
-    skipNA: Boolean = skipNA_default,
-    columns: ColumnsSelector<T, Double?>,
-): Double? = Aggregators.mean.toDouble(skipNA).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullFloat")
-public fun <T> DataFrame<T>.meanOrNull(skipNA: Boolean = skipNA_default, columns: ColumnsSelector<T, Float?>): Double? =
-    Aggregators.mean.toDouble(skipNA).aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullBigInteger")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, BigInteger?>): BigDecimal? =
-    Aggregators.mean.toBigDecimal.aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullBigDecimal")
-public fun <T> DataFrame<T>.meanOrNull(columns: ColumnsSelector<T, BigDecimal?>): BigDecimal? =
-    Aggregators.mean.toBigDecimal.aggregateAll(this, columns)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullNumber")
-public fun <T, C : Number> DataFrame<T>.meanOrNull(
-    skipNA: Boolean = skipNA_default,
-    columns: ColumnsSelector<T, C?>,
-): Number? = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns)
-
-public fun <T> DataFrame<T>.meanOrNull(vararg columns: String, skipNA: Boolean = skipNA_default): Number? =
-    meanOrNull(skipNA) { columns.toNumberColumns() }
-
-@AccessApiOverload
-public fun <T, C : Number> DataFrame<T>.meanOrNull(
-    vararg columns: ColumnReference<C?>,
-    skipNA: Boolean = skipNA_default,
-): Number? = meanOrNull(skipNA) { columns.toColumnSet() }
+): Double = mean(skipNA) { columns.toColumnSet() }
 
 @AccessApiOverload
-public fun <T, C : Number> DataFrame<T>.meanOrNull(
-    vararg columns: KProperty<C?>,
-    skipNA: Boolean = skipNA_default,
-): Number? = meanOrNull(skipNA) { columns.toColumnSet() }
+public fun <T, C : Number> DataFrame<T>.mean(vararg columns: KProperty<C?>, skipNA: Boolean = skipNA_default): Double =
+    mean(skipNA) { columns.toColumnSet() }
 
-// endregion
-
-// region meanOf
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfInt")
-public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, Int?>): Double =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfShort")
-public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Short?>): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfByte")
-public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Byte?>): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfLong")
-public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Long?>): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfDouble")
-public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Double?>): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfFloat")
-public fun <T> DataFrame<T>.meanOf(skipNA: Boolean = skipNA_default, expression: RowExpression<T, Float?>): Double =
-    meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfBigInteger")
-public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, BigInteger?>): BigDecimal =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfBigDecimal")
-public fun <T> DataFrame<T>.meanOf(expression: RowExpression<T, BigDecimal?>): BigDecimal =
-    meanOfOrNull(expression).suggestIfNull("meanOf")
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOfNumber")
 public inline fun <T, reified D : Number> DataFrame<T>.meanOf(
     skipNA: Boolean = skipNA_default,
     noinline expression: RowExpression<T, D?>,
-): Number = meanOfOrNull(skipNA, expression).suggestIfNull("meanOf")
-
-// endregion
-
-// region meanOfOrNull
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfInt")
-public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, Int?>): Double? =
-    Aggregators.mean.toDouble(skipNA_default).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfShort")
-public fun <T> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    expression: RowExpression<T, Short?>,
-): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfByte")
-public fun <T> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    expression: RowExpression<T, Byte?>,
-): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfLong")
-public fun <T> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    expression: RowExpression<T, Long?>,
-): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfDouble")
-public fun <T> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    expression: RowExpression<T, Double?>,
-): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfFloat")
-public fun <T> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    expression: RowExpression<T, Float?>,
-): Double? = Aggregators.mean.toDouble(skipNA).of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfBigInteger")
-public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, BigInteger?>): BigDecimal? =
-    Aggregators.mean.toBigDecimal.of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfBigDecimal")
-public fun <T> DataFrame<T>.meanOfOrNull(expression: RowExpression<T, BigDecimal?>): BigDecimal? =
-    Aggregators.mean.toBigDecimal.of(this, expression)
-
-@OverloadResolutionByLambdaReturnType
-@JvmName("meanOrNullOfNumber")
-public inline fun <T, reified D : Number> DataFrame<T>.meanOfOrNull(
-    skipNA: Boolean = skipNA_default,
-    noinline expression: RowExpression<T, D?>,
-): Number? = Aggregators.mean.toNumber(skipNA).of(this, expression)
-
-// endregion
+): Double = Aggregators.mean(skipNA).of(this, expression) ?: Double.NaN
 
 // endregion
 
@@ -506,7 +104,7 @@ public fun <T> Grouped<T>.mean(skipNA: Boolean = skipNA_default): DataFrame<T> =
 public fun <T, C : Number> Grouped<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, columns)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateFor(this, columns)
 
 public fun <T> Grouped<T>.meanFor(vararg columns: String, skipNA: Boolean = skipNA_default): DataFrame<T> =
     meanFor(skipNA) { columns.toNumberColumns() }
@@ -527,7 +125,7 @@ public fun <T, C : Number> Grouped<T>.mean(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateAll(this, name, columns)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateAll(this, name, columns)
 
 public fun <T> Grouped<T>.mean(
     vararg columns: String,
@@ -553,9 +151,7 @@ public inline fun <T, reified R : Number> Grouped<T>.meanOf(
     name: String? = null,
     skipNA: Boolean = skipNA_default,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> =
-    Aggregators.mean.toNumber(skipNA)
-        .aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateOf(this, name, expression)
 
 // endregion
 
@@ -611,7 +207,7 @@ public fun <T, C : Number> PivotGroupBy<T>.meanFor(
     skipNA: Boolean = skipNA_default,
     separate: Boolean = false,
     columns: ColumnsForAggregateSelector<T, C?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateFor(this, separate, columns)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateFor(this, separate, columns)
 
 public fun <T> PivotGroupBy<T>.meanFor(
     vararg columns: String,
@@ -636,7 +232,7 @@ public fun <T, C : Number> PivotGroupBy<T>.meanFor(
 public fun <T, R : Number> PivotGroupBy<T>.mean(
     skipNA: Boolean = skipNA_default,
     columns: ColumnsSelector<T, R?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateAll(this, columns)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateAll(this, columns)
 
 public fun <T> PivotGroupBy<T>.mean(vararg columns: String, skipNA: Boolean = skipNA_default): DataFrame<T> =
     mean(skipNA) { columns.toColumnsSetOf() }
@@ -656,6 +252,6 @@ public fun <T, R : Number> PivotGroupBy<T>.mean(
 public inline fun <T, reified R : Number> PivotGroupBy<T>.meanOf(
     skipNA: Boolean = skipNA_default,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.mean.toNumber(skipNA).aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.mean(skipNA).aggregateOf(this, expression)
 
 // endregion
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
index 93bbb5950d..6da0acdbac 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ExceptionUtils.kt
@@ -1,22 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.impl
 
-import java.math.BigDecimal
-import java.math.BigInteger
-
 internal fun <T> T?.throwIfNull(message: String): T = this ?: throw NoSuchElementException(message)
 
 @PublishedApi
 internal fun <T> T?.suggestIfNull(operation: String): T =
     throwIfNull("No elements for `$operation` operation. Use `${operation}OrNull` instead.")
-
-@PublishedApi
-internal fun BigInteger?.suggestIfNull(operation: String): BigInteger =
-    throwIfNull(
-        "The `$operation` operation either had no elements, or the result is NaN. Use `${operation}OrNull` instead.",
-    )
-
-@PublishedApi
-internal fun BigDecimal?.suggestIfNull(operation: String): BigDecimal =
-    throwIfNull(
-        "The `$operation` operation either had no elements, or the result is NaN. Use `${operation}OrNull` instead.",
-    )
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
index 7841a70b2d..38be1760bf 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt
@@ -643,13 +643,3 @@ internal fun Iterable<Any>.classes(): Set<KClass<*>> = mapTo(mutableSetOf()) { i
  * @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) }
-
-/**
- * Casts [this]: [Number] to a [Double]. If [this] is `null`, returns [Double.NaN].
- */
-internal fun Number?.asDoubleOrNaN(): Double = this as Double? ?: Double.NaN
-
-/**
- * Casts [this]: [Number] to a [Float]. If [this] is `null`, returns [Float.NaN].
- */
-internal fun Number?.asFloatOrNaN(): Float = this as Float? ?: Float.NaN
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
index 6bb788e694..3693c24af9 100644
--- 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
@@ -19,7 +19,7 @@ import kotlin.reflect.full.withNullability
  *   will always return a [Return]`?`.
  */
 @PublishedApi
-internal interface Aggregator<Value, Return> {
+internal interface Aggregator<in Value, out Return> {
 
     /** The name of this aggregator. */
     val name: String
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
index 3860eedf5b..f791e50111 100644
--- 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
@@ -16,7 +16,7 @@ import kotlin.reflect.full.withNullability
  * @param name The name of this aggregator.
  * @param aggregator Functional argument for the [aggregate] function.
  */
-internal abstract class AggregatorBase<Value, Return>(
+internal abstract class AggregatorBase<in Value, out Return>(
     override val name: String,
     protected val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     protected val aggregator: (values: Iterable<Value>, type: KType) -> Return?,
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index d16def6dcb..11a7054a28 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -7,7 +7,7 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
  * @see AggregatorOptionSwitch2
  */
 @PublishedApi
-internal class AggregatorOptionSwitch1<Param1, AggregatorType : Aggregator<*, *>>(
+internal class AggregatorOptionSwitch1<in Param1, out AggregatorType : Aggregator<*, *>>(
     val name: String,
     val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
 ) {
@@ -28,7 +28,7 @@ internal class AggregatorOptionSwitch1<Param1, AggregatorType : Aggregator<*, *>
      *   MyAggregator.Factory(param1)
      * }
      */
-    class Factory<Param1, AggregatorType : Aggregator<*, *>>(
+    class Factory<in Param1, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
     ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
             AggregatorOptionSwitch1(name, getAggregator)
@@ -42,7 +42,7 @@ internal class AggregatorOptionSwitch1<Param1, AggregatorType : Aggregator<*, *>
  * @see AggregatorOptionSwitch1
  */
 @PublishedApi
-internal class AggregatorOptionSwitch2<Param1, Param2, AggregatorType : Aggregator<*, *>>(
+internal class AggregatorOptionSwitch2<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
     val name: String,
     val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
 ) {
@@ -64,7 +64,7 @@ internal class AggregatorOptionSwitch2<Param1, Param2, AggregatorType : Aggregat
      * }
      * ```
      */
-    class Factory<Param1, Param2, AggregatorType : Aggregator<*, *>>(
+    class Factory<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
     ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
             AggregatorOptionSwitch2(name, getAggregator)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
index a0cbea44fd..9c16fcdb59 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
@@ -10,7 +10,7 @@ import kotlin.reflect.KProperty
  * val myNamedValue by MyFactory
  * ```
  */
-internal fun interface Provider<T> {
+internal fun interface Provider<out T> {
 
     fun create(name: String): T
 }
@@ -25,4 +25,4 @@ internal operator fun <T> Provider<T>.getValue(obj: Any?, property: KProperty<*>
  * val myAggregator by MyAggregator.Factory
  * ```
  */
-internal fun interface AggregatorProvider<AggregatorType : Aggregator<*, *>> : Provider<AggregatorType>
+internal fun interface AggregatorProvider<out AggregatorType : Aggregator<*, *>> : Provider<AggregatorType>
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
index 09c14f4a95..9c6e441139 100644
--- 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,10 @@
 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.mean
 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
@@ -92,43 +90,51 @@ internal object Aggregators {
         getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
     ) = AggregatorOptionSwitch2.Factory(getAggregator)
 
-    val min by twoStepPreservingType<Comparable<Any?>> { minOrNull() }
+    // T: Comparable<T> -> T?
+    val min by twoStepPreservingType<Comparable<Any?>> {
+        minOrNull()
+    }
 
-    val max by twoStepPreservingType<Comparable<Any?>> { maxOrNull() }
+    // T: Comparable<T> -> T?
+    val max by twoStepPreservingType<Comparable<Any?>> {
+        maxOrNull()
+    }
 
+    // T: Number? -> Double
     val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
         flatteningChangingTypes<Number, Double>(
-            getReturnTypeOrNull = { type, emptyInput -> typeOf<Double>().withNullability(emptyInput) },
-        ) { std(it, skipNA, ddof) }
+            getReturnTypeOrNull = { _, emptyInput -> typeOf<Double>().withNullability(emptyInput) },
+        ) { type ->
+            std(type, skipNA, ddof)
+        }
     }
 
-    @Suppress("ClassName")
-    object mean {
-        val toNumber = withOneOption { skipNA: Boolean ->
-            twoStepForNumbers(::meanTypeResultOrNull) { meanOrNull(it, skipNA) }
-        }.create(mean::class.simpleName!!)
-
-        val toDouble = withOneOption { skipNA: Boolean ->
-            twoStepForNumbers(::meanTypeResultOrNull) { meanOrNull(it, skipNA) as Double? }
-        }.create(mean::class.simpleName!!)
-
-        val toBigDecimal =
-            twoStepForNumbers(::meanTypeResultOrNull) {
-                meanOrNull(it) as BigDecimal?
-            }.create(mean::class.simpleName!!)
+    // step one: T: Number? -> Double
+    // step two: Double -> Double
+    val mean by withOneOption { skipNA: Boolean ->
+        twoStepChangingType(
+            getReturnTypeOrNull = { _, _ -> typeOf<Double>() },
+            stepOneAggregator = { type -> mean(type, skipNA) },
+            stepTwoAggregator = { mean(skipNA) },
+        )
     }
 
+    // T: Comparable<T>? -> T
     val percentile by withOneOption { percentile: Double ->
-        flatteningChangingTypes<Comparable<Any?>, Comparable<Any?>>(preserveReturnTypeNullIfEmpty) { type ->
+        flatteningPreservingTypes<Comparable<Any?>> { type ->
             percentile(percentile, type)
         }
     }
 
-    val median by flatteningPreservingTypes<Comparable<Any?>> {
-        median(it)
+    // T: Comparable<T>? -> T
+    val median by flatteningPreservingTypes<Comparable<Any?>> { type ->
+        median(type)
     }
 
+    // T: Number -> T
     val sum by twoStepForNumbers(
         getReturnTypeOrNull = { type, _ -> type.withNullability(false) },
-    ) { sum(it) }
+    ) { type ->
+        sum(type)
+    }
 }
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
index 96e722f05f..bd4a03a9ea 100644
--- 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
@@ -30,7 +30,7 @@ import kotlin.reflect.full.withNullability
  *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
  * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
-internal class FlatteningAggregator<Value, Return>(
+internal class FlatteningAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     aggregator: (values: Iterable<Value>, type: KType) -> Return?,
@@ -54,7 +54,7 @@ internal class FlatteningAggregator<Value, Return>(
      * @param aggregator Functional argument for the [aggregate] function.
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
-    class Factory<Value, Return>(
+    class Factory<in Value, out Return>(
         private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         private val aggregator: (Iterable<Value>, KType) -> Return?,
         private val preservesType: Boolean,
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
index a000e6e662..8cb4ee81a1 100644
--- 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
@@ -37,7 +37,7 @@ import kotlin.reflect.full.withNullability
  *   It is run on the results of [stepOneAggregator].
  * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
-internal class TwoStepAggregator<Value, Return>(
+internal class TwoStepAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     stepOneAggregator: (values: Iterable<Value>, type: KType) -> Return?,
@@ -73,7 +73,7 @@ internal class TwoStepAggregator<Value, Return>(
      *   It is run on the results of [stepOneAggregator].
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
-    class Factory<Value, Return>(
+    class Factory<in Value, out Return>(
         private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         private val stepOneAggregator: (Iterable<Value>, KType) -> Return?,
         private val stepTwoAggregator: (Iterable<Return>, KType) -> Return?,
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
index be5fd22e1a..d7bd0a8c8c 100644
--- 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
@@ -37,7 +37,7 @@ import kotlin.reflect.typeOf
  *   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<Return : Number>(
+internal class TwoStepNumbersAggregator<out Return : Number>(
     name: String,
     getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
     aggregator: (values: Iterable<Number>, numberType: KType) -> Return?,
@@ -84,7 +84,7 @@ internal class TwoStepNumbersAggregator<Return : Number>(
 
     override val preservesType = false
 
-    class Factory<Return : Number>(
+    class Factory<out Return : Number>(
         private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
         private val aggregate: Iterable<Number>.(numberType: KType) -> Return?,
     ) : AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
index 1b734cf8f9..254a598a86 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt
@@ -16,7 +16,6 @@ import org.jetbrains.kotlinx.dataframe.api.isNumber
 import org.jetbrains.kotlinx.dataframe.api.map
 import org.jetbrains.kotlinx.dataframe.api.maxOrNull
 import org.jetbrains.kotlinx.dataframe.api.mean
-import org.jetbrains.kotlinx.dataframe.api.meanOrNull
 import org.jetbrains.kotlinx.dataframe.api.medianOrNull
 import org.jetbrains.kotlinx.dataframe.api.minOrNull
 import org.jetbrains.kotlinx.dataframe.api.move
@@ -57,7 +56,7 @@ internal fun describeImpl(cols: List<AnyCol>): DataFrame<ColumnDescription> {
                 ?.key
         }
         if (hasNumericCols) {
-            ColumnDescription::mean from { if (it.isNumber()) it.asNumbers().meanOrNull() else null }
+            ColumnDescription::mean from { if (it.isNumber()) it.asNumbers().mean() else null }
             ColumnDescription::std from { if (it.isNumber()) it.asNumbers().std() else null }
         }
         if (hasComparableCols || hasNumericCols) {
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
index 6f95432fa9..bb4c1e2f21 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -1,108 +1,49 @@
 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
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
-import kotlin.reflect.typeOf
 
-/** @include [Sequence.meanOrNull] */
 @PublishedApi
-internal fun <T : Number> Iterable<T>.meanOrNull(type: KType, skipNA: Boolean = skipNA_default): Number? =
-    asSequence().meanOrNull(type, skipNA)
-
-/**
- * Returns the mean of the numbers in [this].
- *
- * If the input is empty, the return value will be `null`.
- *
- * If the [type] given or input consists of only [Int], [Short], [Byte], [Long], [Double], or [Float],
- * the return type will be [Double].
- *
- * If the [type] given or the input contains [BigInteger] or [BigDecimal],
- * the return type will be [BigDecimal].
- * @param type The type of the numbers in the sequence.
- * @param skipNA Whether to skip `NaN` values (default: `false`). Only relevant for [Double] and [Float].
- */
+internal fun <T : Number> Iterable<T?>.mean(type: KType, skipNA: Boolean = skipNA_default): Double =
+    asSequence().mean(type, skipNA)
+
 @Suppress("UNCHECKED_CAST")
-internal fun <T : Number> Sequence<T>.meanOrNull(type: KType, skipNA: Boolean = skipNA_default): Number? {
+internal fun <T : Number> Sequence<T?>.mean(type: KType, skipNA: Boolean = skipNA_default): Double {
     if (type.isMarkedNullable) {
-        return filterNotNull().meanOrNull(type.withNullability(false), skipNA)
+        return filterNotNull().mean(type.withNullability(false), skipNA)
     }
     return when (type.classifier) {
-        // Double -> Double
-        Double::class -> (this as Sequence<Double>).meanOrNull(skipNA)
+        Double::class -> (this as Sequence<Double>).mean(skipNA)
 
-        // Float -> Double
-        Float::class -> (this as Sequence<Float>).meanOrNull(skipNA)
+        Float::class -> (this as Sequence<Float>).mean(skipNA)
 
-        // Int -> Double
-        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.meanOrNull(false)
+        Int::class -> (this as Sequence<Int>).map { it.toDouble() }.mean(false)
 
-        // Short -> Double
-        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.meanOrNull(false)
+        // for integer values NA is not possible
+        Short::class -> (this as Sequence<Short>).map { it.toDouble() }.mean(false)
 
-        // Byte -> Double
-        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.meanOrNull(false)
+        Byte::class -> (this as Sequence<Byte>).map { it.toDouble() }.mean(false)
 
-        // Long -> Double
-        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.meanOrNull(false)
+        Long::class -> (this as Sequence<Long>).map { it.toDouble() }.mean(false)
 
-        // BigInteger -> BigDecimal
-        BigInteger::class -> (this as Sequence<BigInteger>).meanOrNull()
+        BigInteger::class -> (this as Sequence<BigInteger>).map { it.toDouble() }.mean(false)
 
-        // BigDecimal -> BigDecimal
-        BigDecimal::class -> (this as Sequence<BigDecimal>).meanOrNull()
+        BigDecimal::class -> (this as Sequence<BigDecimal>).map { it.toDouble() }.mean(skipNA)
 
-        // Number -> Conversion(Common number type) -> Number? (Double or BigDecimal?)
-        // fallback case, heavy as it needs to collect all types at runtime
-        Number::class -> {
-            val numberTypes = (this as Sequence<Number>).asIterable().types()
-            val unifiedType = numberTypes.unifiedNumberType()
-            if (unifiedType.withNullability(false) == typeOf<Number>()) {
-                error("Cannot find unified number type for $numberTypes")
-            }
-            this.convertToUnifiedNumberType(unifiedType)
-                .meanOrNull(unifiedType, skipNA)
-        }
+        Number::class -> (this as Sequence<Number>).map { it.toDouble() }.mean(skipNA)
 
         // this means the sequence is empty
-        Nothing::class -> null
+        Nothing::class -> Double.NaN
 
         else -> throw IllegalArgumentException("Unable to compute the mean for type ${renderType(type)}")
     }
 }
 
-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? {
+internal fun Sequence<Double>.mean(skipNA: Boolean = skipNA_default): Double {
     var count = 0
     var sum: Double = 0.toDouble()
     for (element in this) {
@@ -110,17 +51,17 @@ internal fun Sequence<Double>.meanOrNull(skipNA: Boolean = skipNA_default): Doub
             if (skipNA) {
                 continue
             } else {
-                return null
+                return Double.NaN
             }
         }
         sum += element
         count++
     }
-    return if (count > 0) sum / count else null
+    return if (count > 0) sum / count else Double.NaN
 }
 
 @JvmName("meanFloat")
-internal fun Sequence<Float>.meanOrNull(skipNA: Boolean = skipNA_default): Double? {
+internal fun Sequence<Float>.mean(skipNA: Boolean = skipNA_default): Double {
     var count = 0
     var sum: Double = 0.toDouble()
     for (element in this) {
@@ -128,95 +69,97 @@ internal fun Sequence<Float>.meanOrNull(skipNA: Boolean = skipNA_default): Doubl
             if (skipNA) {
                 continue
             } else {
-                return null
+                return Double.NaN
             }
         }
         sum += element
         count++
     }
-    return if (count > 0) sum / count else null
-}
-
-@JvmName("bigIntegerMean")
-internal fun Sequence<BigInteger>.meanOrNull(): BigDecimal? {
-    var count = 0
-    val sum = sumOf {
-        count++
-        it
-    }
-    return if (count > 0) sum.toBigDecimal() / count.toBigDecimal() else null
-}
-
-@JvmName("bigDecimalMean")
-internal fun Sequence<BigDecimal>.meanOrNull(): BigDecimal? {
-    var count = 0
-    val sum = sumOf {
-        count++
-        it
-    }
-    return if (count > 0) sum.toBigDecimal() / count.toBigDecimal() else null
+    return if (count > 0) sum / count else Double.NaN
 }
 
 @JvmName("doubleMean")
-internal fun Iterable<Double>.meanOrNull(skipNA: Boolean = skipNA_default): Double? = asSequence().meanOrNull(skipNA)
+internal fun Iterable<Double>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
 
 @JvmName("floatMean")
-internal fun Iterable<Float>.meanOrNull(skipNA: Boolean = skipNA_default): Double? = asSequence().meanOrNull(skipNA)
-
-@JvmName("bigDecimalMean")
-internal fun Iterable<BigDecimal>.meanOrNull(): BigDecimal? = asSequence().meanOrNull()
-
-@JvmName("bigIntegerMean")
-internal fun Iterable<BigInteger>.meanOrNull(): BigDecimal? = asSequence().meanOrNull()
+internal fun Iterable<Float>.mean(skipNA: Boolean = skipNA_default): Double = asSequence().mean(skipNA)
 
 @JvmName("intMean")
-internal fun Iterable<Int>.meanOrNull(): Double? =
+internal fun Iterable<Int>.mean(): Double =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else null
+        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else null
+        if (count > 0) sum / count else Double.NaN
     }
 
 @JvmName("shortMean")
-internal fun Iterable<Short>.meanOrNull(): Double? =
+internal fun Iterable<Short>.mean(): Double =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else null
+        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else null
+        if (count > 0) sum / count else Double.NaN
     }
 
 @JvmName("byteMean")
-internal fun Iterable<Byte>.meanOrNull(): Double? =
+internal fun Iterable<Byte>.mean(): Double =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else null
+        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else null
+        if (count > 0) sum / count else Double.NaN
     }
 
 @JvmName("longMean")
-internal fun Iterable<Long>.meanOrNull(): Double? =
+internal fun Iterable<Long>.mean(): Double =
+    if (this is Collection) {
+        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+    } else {
+        var count = 0
+        val sum = sumOf {
+            count++
+            it.toDouble()
+        }
+        if (count > 0) sum / count else Double.NaN
+    }
+
+// TODO result is Double, but should be BigDecimal, Issue #558
+@JvmName("bigIntegerMean")
+internal fun Iterable<BigInteger>.mean(): Double =
+    if (this is Collection) {
+        if (size > 0) sumOf { it.toDouble() } / size else Double.NaN
+    } else {
+        var count = 0
+        val sum = sumOf {
+            count++
+            it.toDouble()
+        }
+        if (count > 0) sum / count else Double.NaN
+    }
+
+// TODO result is Double, but should be BigDecimal, Issue #558
+@JvmName("bigDecimalMean")
+internal fun Iterable<BigDecimal>.mean(): Double =
     if (this is Collection) {
-        if (size > 0) sumOf { it.toDouble() } / size else null
+        if (size > 0) sum().toDouble() / size else Double.NaN
     } else {
         var count = 0
         val sum = sumOf {
             count++
             it.toDouble()
         }
-        if (count > 0) sum / count else null
+        if (count > 0) sum / count else Double.NaN
     }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
index 188c6673f9..71f049f5ca 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import io.kotest.matchers.doubles.shouldBeNaN
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.alsoDebug
 import org.junit.Test
@@ -39,7 +40,7 @@ class DescribeTests {
             nulls shouldBe 1
             top shouldBe 1
             freq shouldBe 1
-            mean shouldBe 4.5.toBigDecimal()
+            mean shouldBe 4.5
             std shouldBe 2.449489742783178
             min shouldBe 1.toBigDecimal()
             (p25 as BigDecimal).setScale(2) shouldBe 2.75.toBigDecimal()
@@ -64,8 +65,8 @@ class DescribeTests {
             nulls shouldBe 0
             top shouldBe 1
             freq shouldBe 1
-            mean shouldBe null
-            std.isNaN shouldBe true
+            mean.shouldBeNaN()
+            std.shouldBeNaN()
             min shouldBe 1.0 // TODO should be NaN too?
             p25 shouldBe 1.75
             median shouldBe 3.0
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
index 71e55b21d6..4241f7570a 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/BasicTests.kt
@@ -142,9 +142,9 @@ class BasicTests {
     @Test
     fun `calculate mean age for each animal`() {
         val expected = dataFrameOf("animal", "age")(
-            "cat", null,
+            "cat", Double.NaN,
             "snake", 2.5,
-            "dog", null,
+            "dog", Double.NaN,
         )
 
         df.groupBy { animal }.mean { age } shouldBe expected
@@ -213,7 +213,7 @@ class BasicTests {
         val expected = dataFrameOf("animal", "1", "3", "2")(
             "cat", 2.5, 2.5, null,
             "snake", 4.5, null, 0.5,
-            "dog", 3.0, null, 6.0,
+            "dog", 3.0, Double.NaN, 6.0,
         )
 
         val actualDfAcc = df.pivot(inward = false) { visits }.groupBy { animal }.mean(skipNA = true) { age }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
index 338be0fd59..72fd1dc4dc 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/puzzles/MediumTests.kt
@@ -66,7 +66,7 @@ class MediumTests {
             -1, 0, 1,
         )
 
-        df.convert { colsOf<Double>() }.with { (it - rowMean().toDouble()).roundToInt() } shouldBe expected
+        df.convert { colsOf<Double>() }.with { (it - rowMean()).roundToInt() } shouldBe expected
     }
 
     @Test
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
index 80c571e95a..0c0c156f26 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/BasicMathTests.kt
@@ -2,11 +2,9 @@ package org.jetbrains.kotlinx.dataframe.statistics
 
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.columnOf
 import org.jetbrains.kotlinx.dataframe.api.mean
-import org.jetbrains.kotlinx.dataframe.api.meanOrNull
-import org.jetbrains.kotlinx.dataframe.impl.nullableNothingType
+import org.jetbrains.kotlinx.dataframe.impl.nothingType
 import org.junit.Test
 import kotlin.reflect.typeOf
 
@@ -20,17 +18,10 @@ class BasicMathTests {
 
     @Test
     fun `mean with nans and nulls`() {
-        columnOf(10, 20, Double.NaN, null).meanOrNull() shouldBe null
+        columnOf(10, 20, Double.NaN, null).mean() shouldBe Double.NaN
         columnOf(10, 20, Double.NaN, null).mean(skipNA = true) shouldBe 15
 
-        DataColumn.createValueColumn("", emptyList<Nothing>(), nullableNothingType)
-            .cast<Number?>()
-            .meanOrNull() shouldBe null
-        DataColumn.createValueColumn("", emptyList<Nothing>(), typeOf<Double?>())
-            .cast<Double?>()
-            .meanOrNull() shouldBe null
-        DataColumn.createValueColumn("", listOf(null), typeOf<Double?>())
-            .cast<Double?>()
-            .meanOrNull() shouldBe null
+        DataColumn.createValueColumn("", emptyList<Nothing>(), nothingType(false)).mean() shouldBe Double.NaN
+        DataColumn.createValueColumn("", listOf(null), nothingType(true)).mean() shouldBe Double.NaN
     }
 }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
index 3283d4e9d3..f8ee4b37d5 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/animals/AnimalsTests.kt
@@ -33,7 +33,7 @@ class AnimalsTests {
         mean.columnsCount() shouldBe 2
         mean.rowsCount() shouldBe 2
         mean.name.values() shouldBe listOf("age", "visits")
-        mean.value.type() shouldBe typeOf<Double?>()
+        mean.value.type() shouldBe typeOf<Double>()
     }
 
     @Test
@@ -42,7 +42,7 @@ class AnimalsTests {
             .update { age }.with { Double.NaN }
             .update { visits }.withNull()
         val mean = cleared.mean()
-        mean[age] shouldBe null
-        mean[visits.name()] shouldBe null
+        mean[age] shouldBe Double.NaN
+        (mean[visits.name()] as Double).isNaN() shouldBe true
     }
 }
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
index 831d826d71..5bc715e078 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt
@@ -1,7 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.testSets.person
 
 import io.kotest.assertions.throwables.shouldThrow
-import io.kotest.matchers.collections.shouldBeIn
 import io.kotest.matchers.doubles.ToleranceMatcher
 import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
@@ -123,7 +122,6 @@ import org.jetbrains.kotlinx.dataframe.api.rename
 import org.jetbrains.kotlinx.dataframe.api.reorderColumnsByName
 import org.jetbrains.kotlinx.dataframe.api.replace
 import org.jetbrains.kotlinx.dataframe.api.rows
-import org.jetbrains.kotlinx.dataframe.api.schema
 import org.jetbrains.kotlinx.dataframe.api.select
 import org.jetbrains.kotlinx.dataframe.api.single
 import org.jetbrains.kotlinx.dataframe.api.sortBy
@@ -179,12 +177,11 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.isMissingColumn
 import org.jetbrains.kotlinx.dataframe.impl.emptyPath
 import org.jetbrains.kotlinx.dataframe.impl.getColumnsImpl
 import org.jetbrains.kotlinx.dataframe.impl.nothingType
-import org.jetbrains.kotlinx.dataframe.impl.nullableNothingType
 import org.jetbrains.kotlinx.dataframe.impl.trackColumnAccess
 import org.jetbrains.kotlinx.dataframe.index
 import org.jetbrains.kotlinx.dataframe.io.renderValueForStdout
 import org.jetbrains.kotlinx.dataframe.kind
-import org.jetbrains.kotlinx.dataframe.math.meanOrNull
+import org.jetbrains.kotlinx.dataframe.math.mean
 import org.jetbrains.kotlinx.dataframe.name
 import org.jetbrains.kotlinx.dataframe.ncol
 import org.jetbrains.kotlinx.dataframe.nrow
@@ -820,8 +817,8 @@ class DataFrameTests : BaseTest() {
 
     @Test
     fun `groupBy meanOf`() {
-        typed.groupBy { name }.meanOf { age * 2 }.schema() shouldBe typed
-            .groupBy { name }.aggregate { mean { age } * 2 into "mean" }.schema()
+        typed.groupBy { name }.meanOf { age * 2 } shouldBe typed
+            .groupBy { name }.aggregate { mean { age } * 2 into "mean" }
     }
 
     @Test
@@ -1488,7 +1485,7 @@ class DataFrameTests : BaseTest() {
 
     @Test
     fun `column stats`() {
-        typed.age.mean() shouldBe typed.age.toList().meanOrNull()
+        typed.age.mean() shouldBe typed.age.toList().mean()
         typed.age.min() shouldBe typed.age.toList().minOrNull()
         typed.age.max() shouldBe typed.age.toList().maxOrNull()
         typed.age.sum() shouldBe typed.age.toList().sum()
@@ -2113,8 +2110,8 @@ class DataFrameTests : BaseTest() {
             it.kind() shouldBe ColumnKind.Group
             val group = it.asColumnGroup()
             group.columnNames() shouldBe listOf("age", "weight")
-            group.columnTypes().forEach {
-                it shouldBeIn setOf(typeOf<Double?>(), nullableNothingType)
+            group.columns().forEach {
+                it.type() shouldBe typeOf<Double?>()
             }
         }
     }
diff --git a/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt b/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
index 64d0e9423c..b29fd1625a 100644
--- a/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
+++ b/examples/idea-examples/titanic/src/main/kotlin/org/jetbrains/kotlinx/dataframe/examples/titanic/ml/titanic.kt
@@ -24,10 +24,11 @@ private val model = Sequential.of(
     Input(9),
     Dense(50, Activations.Relu, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
     Dense(50, Activations.Relu, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
-    Dense(2, Activations.Linear, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros()),
+    Dense(2, Activations.Linear, kernelInitializer = HeNormal(SEED), biasInitializer = Zeros())
 )
 
 fun main() {
+
     // Set Locale for correct number parsing
     Locale.setDefault(Locale.FRANCE)
 
@@ -36,7 +37,7 @@ fun main() {
     // Calculating imputing values
     val (train, test) = df
         // imputing
-        .fillNulls { sibsp and parch and age and fare }.perCol { it.mean()?.toDouble() }
+        .fillNulls { sibsp and parch and age and fare }.perCol { it.mean() }
         .fillNulls { sex }.with { "female" }
         // one hot encoding
         .pivotMatches { pclass and sex }
@@ -49,7 +50,7 @@ fun main() {
         it.compile(
             optimizer = Adam(),
             loss = Losses.SOFT_MAX_CROSS_ENTROPY_WITH_LOGITS,
-            metric = Metrics.ACCURACY,
+            metric = Metrics.ACCURACY
         )
 
         it.summary()

From ce57e66594a97415bc5af1ec20c5a1419126c5c6 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Thu, 6 Mar 2025 16:44:10 +0100
Subject: [PATCH 07/18] extracting some common lambdas to type aliases, adding
 some docs

---
 .../aggregation/aggregators/Aggregator.kt     | 27 +++++++++---
 .../aggregation/aggregators/AggregatorBase.kt | 43 +++++++++++++------
 .../aggregators/AggregatorOptionSwitch.kt     |  4 +-
 .../aggregation/aggregators/Aggregators.kt    | 34 +++++++--------
 .../aggregators/FlatteningAggregator.kt       | 14 +++---
 .../aggregators/TwoStepAggregator.kt          | 18 ++++----
 .../aggregators/TwoStepNumbersAggregator.kt   | 41 +++++++++++++++---
 .../jetbrains/kotlinx/dataframe/math/mean.kt  |  7 +++
 .../jetbrains/kotlinx/dataframe/math/std.kt   |  7 +++
 .../jetbrains/kotlinx/dataframe/math/sum.kt   |  7 +++
 10 files changed, 142 insertions(+), 60 deletions(-)

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
index 3693c24af9..20f6a27b93 100644
--- 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
@@ -33,12 +33,14 @@ internal interface Aggregator<in Value, out Return> {
      * Aggregates the given values, taking [type] into account, 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?
 
     /**
      * Aggregates the data in the given column and computes a single resulting value.
-     * Nulls are filtered out by default, then the aggregation function (with [Iterable] and [KType]) is called.
+     * Nulls are filtered out by default, then [aggregate] (with [Iterable] and [KType]) is called.
      *
      * See [AggregatorBase.aggregate].
      */
@@ -46,20 +48,28 @@ internal interface Aggregator<in Value, out Return> {
 
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
-     *
-     * Must be overridden when using [AggregatorBase].
      */
     fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
      * This is a heavy operation and should be avoided when possible.
-     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     *
+     * @param values The values to be aggregated.
+     * @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!).
      */
     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.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param type The type of the input values.
+     * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
      */
     fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
 }
@@ -70,6 +80,13 @@ internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Ag
 @PublishedApi
 internal fun <Value, Return> Aggregator<*, *>.cast2(): Aggregator<Value, Return> = this as Aggregator<Value, Return>
 
-internal val preserveReturnTypeNullIfEmpty: (KType, Boolean) -> KType = { type, emptyInput ->
+/** 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?
+
+/** Common case for [CalculateReturnTypeOrNull], preserves return type, but makes it nullable for empty inputs. */
+internal val preserveReturnTypeNullIfEmpty: CalculateReturnTypeOrNull = { 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
index f791e50111..a2dac82aed 100644
--- 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
@@ -8,7 +8,7 @@ import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
 /**
- * Base class for [aggregators][Aggregator].
+ * Abstract base class for [aggregators][Aggregator].
  *
  * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
  * or multiple [DataColumns][DataColumn].
@@ -18,37 +18,52 @@ import kotlin.reflect.full.withNullability
  */
 internal abstract class AggregatorBase<in Value, out Return>(
     override val name: String,
-    protected val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-    protected val aggregator: (values: Iterable<Value>, type: KType) -> Return?,
+    protected val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    protected val aggregator: Aggregate<Value, Return>,
 ) : Aggregator<Value, Return> {
 
     /**
      * Base function of [Aggregator].
      *
      * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
      * Uses [aggregator] to compute the result.
+     *
+     * When the exact [type] is unknown, use [aggregateCalculatingType].
      */
     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.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * Uses [getReturnTypeOrNull] to calculate the return type.
+     *
+     * @param type The type of the input values.
+     * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
     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].
+     *
+     * Nulls are filtered out by default, then [aggregate] (with [Iterable] and [KType]) is called.
      */
+    @Suppress("UNCHECKED_CAST")
     override fun aggregate(column: DataColumn<Value?>): Return? =
-        if (column.hasNulls()) {
-            aggregate(column.asSequence().filterNotNull().asIterable(), column.type().withNullability(false))
-        } else {
-            aggregate(column.asIterable() as Iterable<Value>, column.type().withNullability(false))
-        }
+        aggregate(
+            values =
+                if (column.hasNulls()) {
+                    column.asSequence().filterNotNull().asIterable()
+                } else {
+                    column.asIterable() as Iterable<Value>
+                },
+            type = column.type().withNullability(false),
+        )
 
-    /**
-     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
-     * This is a heavy operation and should be avoided when possible.
-     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
-     */
+    /** @include [Aggregator.aggregateCalculatingType] */
     override fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>?): Return? {
         val commonType = if (valueTypes != null) {
             valueTypes.commonType(false)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index 11a7054a28..7a2b2d4496 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -29,7 +29,7 @@ internal class AggregatorOptionSwitch1<in Param1, out AggregatorType : Aggregato
      * }
      */
     class Factory<in Param1, out AggregatorType : Aggregator<*, *>>(
-        val getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
+        val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
     ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
             AggregatorOptionSwitch1(name, getAggregator)
         })
@@ -65,7 +65,7 @@ internal class AggregatorOptionSwitch2<in Param1, in Param2, out AggregatorType
      * ```
      */
     class Factory<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
-        val getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
+        val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
     ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
             AggregatorOptionSwitch2(name, getAggregator)
         })
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
index 9c6e441139..8c677a0990 100644
--- 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,13 +1,13 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.math.mean
+import org.jetbrains.kotlinx.dataframe.math.meanTypeConversion
 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.stdTypeConversion
 import org.jetbrains.kotlinx.dataframe.math.sum
-import kotlin.reflect.KType
-import kotlin.reflect.full.withNullability
-import kotlin.reflect.typeOf
+import org.jetbrains.kotlinx.dataframe.math.sumTypeConversion
 
 @PublishedApi
 internal object Aggregators {
@@ -17,7 +17,7 @@ internal object Aggregators {
      *
      * @include [TwoStepAggregator]
      */
-    private fun <Type> twoStepPreservingType(aggregator: Iterable<Type>.(type: KType) -> Type?) =
+    private fun <Type> twoStepPreservingType(aggregator: Aggregate<Type, Type>) =
         TwoStepAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             stepOneAggregator = aggregator,
@@ -31,9 +31,9 @@ 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,
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        stepOneAggregator: Aggregate<Value, Return>,
+        stepTwoAggregator: Aggregate<Return, Return>,
     ) = TwoStepAggregator.Factory(
         getReturnTypeOrNull = getReturnTypeOrNull,
         stepOneAggregator = stepOneAggregator,
@@ -46,7 +46,7 @@ internal object Aggregators {
      *
      * @include [FlatteningAggregator]
      */
-    private fun <Type> flatteningPreservingTypes(aggregate: Iterable<Type?>.(type: KType) -> Type?) =
+    private fun <Type> flatteningPreservingTypes(aggregate: Aggregate<Type, Type>) =
         FlatteningAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             aggregator = aggregate,
@@ -59,8 +59,8 @@ internal object Aggregators {
      * @include [FlatteningAggregator]
      */
     private fun <Value, Return> flatteningChangingTypes(
-        getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-        aggregate: Iterable<Value?>.(type: KType) -> Return?,
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        aggregate: Aggregate<Value, Return>,
     ) = FlatteningAggregator.Factory(
         getReturnTypeOrNull = getReturnTypeOrNull,
         aggregator = aggregate,
@@ -73,8 +73,8 @@ internal object Aggregators {
      * @include [TwoStepNumbersAggregator]
      */
     private fun <Return : Number> twoStepForNumbers(
-        getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-        aggregate: Iterable<Number>.(numberType: KType) -> Return?,
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        aggregate: Aggregate<Number, Return>,
     ) = TwoStepNumbersAggregator.Factory(
         getReturnTypeOrNull = getReturnTypeOrNull,
         aggregate = aggregate,
@@ -102,9 +102,7 @@ internal object Aggregators {
 
     // T: Number? -> Double
     val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
-        flatteningChangingTypes<Number, Double>(
-            getReturnTypeOrNull = { _, emptyInput -> typeOf<Double>().withNullability(emptyInput) },
-        ) { type ->
+        flatteningChangingTypes<Number, Double>(stdTypeConversion) { type ->
             std(type, skipNA, ddof)
         }
     }
@@ -113,7 +111,7 @@ internal object Aggregators {
     // step two: Double -> Double
     val mean by withOneOption { skipNA: Boolean ->
         twoStepChangingType(
-            getReturnTypeOrNull = { _, _ -> typeOf<Double>() },
+            getReturnTypeOrNull = meanTypeConversion,
             stepOneAggregator = { type -> mean(type, skipNA) },
             stepTwoAggregator = { mean(skipNA) },
         )
@@ -132,9 +130,7 @@ internal object Aggregators {
     }
 
     // T: Number -> T
-    val sum by twoStepForNumbers(
-        getReturnTypeOrNull = { type, _ -> type.withNullability(false) },
-    ) { type ->
+    val sum by twoStepForNumbers(sumTypeConversion) { type ->
         sum(type)
     }
 }
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
index bd4a03a9ea..b259339a69 100644
--- 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
@@ -2,7 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.impl.commonType
-import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
 /**
@@ -21,19 +20,21 @@ import kotlin.reflect.full.withNullability
  *     -> Return?
  * ```
  *
- * This is essential for aggregators that depend on the distribution of all values across the dataframe.
+ * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
+ * the median, percentile, and standard deviation.
  *
  * See [TwoStepAggregator] for different behavior for multiple columns.
  *
  * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
  * @param aggregator Functional argument for the [aggregate] function.
  *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
  * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
 internal class FlatteningAggregator<in Value, out Return>(
     name: String,
-    getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-    aggregator: (values: Iterable<Value>, type: KType) -> Return?,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    aggregator: Aggregate<Value, Return>,
     override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, aggregator) {
 
@@ -51,12 +52,13 @@ internal class FlatteningAggregator<in Value, out Return>(
     /**
      * Creates [FlatteningAggregator].
      *
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
      * @param aggregator Functional argument for the [aggregate] function.
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<in Value, out Return>(
-        private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-        private val aggregator: (Iterable<Value>, KType) -> Return?,
+        private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        private val aggregator: Aggregate<Value, Return>,
         private val preservesType: Boolean,
     ) : AggregatorProvider<FlatteningAggregator<Value, Return>> by AggregatorProvider({ name ->
             FlatteningAggregator(
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
index 8cb4ee81a1..b095708a2a 100644
--- 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
@@ -2,8 +2,6 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 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
 
@@ -32,6 +30,7 @@ import kotlin.reflect.full.withNullability
  * See [FlatteningAggregator] for different behavior for multiple columns.
  *
  * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
  * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
  * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
  *   It is run on the results of [stepOneAggregator].
@@ -39,9 +38,9 @@ import kotlin.reflect.full.withNullability
  */
 internal class TwoStepAggregator<in Value, out 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?,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    stepOneAggregator: Aggregate<Value, Return>,
+    private val stepTwoAggregator: Aggregate<Return, Return>,
     override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
@@ -49,6 +48,8 @@ internal class TwoStepAggregator<in Value, out Return>(
      * 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.
+     *
+     * Post-step-one types are calculated by [calculateReturnTypeOrNull].
      */
     override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return? {
         val (values, types) = columns.mapNotNull { col ->
@@ -68,15 +69,16 @@ internal class TwoStepAggregator<in Value, out Return>(
     /**
      * Creates [TwoStepAggregator].
      *
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
      * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
      * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<in Value, out 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 getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        private val stepOneAggregator: Aggregate<Value, Return>,
+        private val stepTwoAggregator: Aggregate<Return, Return>,
         private val preservesType: Boolean,
     ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
             TwoStepAggregator(
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
index d7bd0a8c8c..4515d14bc9 100644
--- 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
@@ -33,16 +33,26 @@ import kotlin.reflect.typeOf
  * ```
  *
  * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
  * @param aggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
  *   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>(
     name: String,
-    getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-    aggregator: (values: Iterable<Number>, numberType: KType) -> Return?,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    aggregator: Aggregate<Number, Return>,
 ) : AggregatorBase<Number, Return>(name, getReturnTypeOrNull, aggregator) {
 
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * Uses [aggregator] to compute the result.
+     *
+     * When the exact [type] is unknown, use [aggregateCalculatingType].
+     */
     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?"
@@ -50,6 +60,14 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         return super.aggregate(values, type)
     }
 
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     *
+     * This function calls [aggregator] on each column and then again on the results.
+     *
+     * After the first aggregation, the number types are found by [calculateReturnTypeOrNull] and then
+     * unified using [aggregateCalculatingType].
+     */
     override fun aggregate(columns: Iterable<DataColumn<Number?>>): Return? {
         val (values, types) = columns.mapNotNull { col ->
             val value = aggregate(col) ?: return@mapNotNull null
@@ -69,9 +87,14 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
-     * of the values at runtime.
+     * of the values at runtime and converts all numbers to this type before aggregating.
      * This is a heavy operation and should be avoided when possible.
-     * If provided, [valueTypes] can be used to avoid calculating the types of [values] at runtime.
+     *
+     * @param values The numbers to be aggregated.
+     * @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!).
      */
     @Suppress("UNCHECKED_CAST")
     override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
@@ -84,9 +107,15 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
 
     override val preservesType = false
 
+    /**
+     * Creates [TwoStepNumbersAggregator].
+     *
+     * @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>(
-        private val getReturnTypeOrNull: (type: KType, emptyInput: Boolean) -> KType?,
-        private val aggregate: Iterable<Number>.(numberType: KType) -> Return?,
+        private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        private val aggregate: Aggregate<Number, Return>,
     ) : AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->
             TwoStepNumbersAggregator(
                 name = name,
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
index bb4c1e2f21..fc5fb70b70 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -1,11 +1,13 @@
 package org.jetbrains.kotlinx.dataframe.math
 
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
 @PublishedApi
 internal fun <T : Number> Iterable<T?>.mean(type: KType, skipNA: Boolean = skipNA_default): Double =
@@ -43,6 +45,11 @@ internal fun <T : Number> Sequence<T?>.mean(type: KType, skipNA: Boolean = skipN
     }
 }
 
+// T: Number? -> Double
+internal val meanTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
+    typeOf<Double>()
+}
+
 internal fun Sequence<Double>.mean(skipNA: Boolean = skipNA_default): Double {
     var count = 0
     var sum: Double = 0.toDouble()
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
index 052556ba59..148c9ece23 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
@@ -2,11 +2,13 @@ package org.jetbrains.kotlinx.dataframe.math
 
 import org.jetbrains.kotlinx.dataframe.api.ddof_default
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
 @Suppress("UNCHECKED_CAST")
 @PublishedApi
@@ -35,6 +37,11 @@ internal fun <T : Number> Iterable<T?>.std(
     }
 }
 
+// T: Number? -> Double
+internal val stdTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
+    typeOf<Double>()
+}
+
 @JvmName("doubleStd")
 internal fun Iterable<Double>.std(skipNA: Boolean = skipNA_default, ddof: Int = ddof_default): Double =
     varianceAndMean(skipNA)?.std(ddof) ?: Double.NaN
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
index 08dae78937..1b03221988 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
@@ -1,8 +1,10 @@
 package org.jetbrains.kotlinx.dataframe.math
 
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
 @PublishedApi
 internal fun <T, R : Number> Iterable<T>.sumOf(type: KType, selector: (T) -> R?): R {
@@ -95,6 +97,11 @@ internal fun <T : Number> Iterable<T?>.sum(type: KType): T =
         else -> throw IllegalArgumentException("sum is not supported for $type")
     }
 
+// T: Number? -> T
+internal val sumTypeConversion: CalculateReturnTypeOrNull = { type, _ ->
+    type.withNullability(false)
+}
+
 @PublishedApi
 internal fun Iterable<BigDecimal>.sum(): BigDecimal {
     var sum: BigDecimal = BigDecimal.ZERO

From 65ffe9bc9dd5ae08662921924b017a0d1575540d Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Thu, 6 Mar 2025 20:15:39 +0100
Subject: [PATCH 08/18] updating from master

---
 .../jetbrains/kotlinx/dataframe/api/max.kt    |   8 +-
 .../jetbrains/kotlinx/dataframe/api/median.kt |  15 +-
 .../jetbrains/kotlinx/dataframe/api/min.kt    |   8 +-
 .../kotlinx/dataframe/api/percentile.kt       |  10 +-
 .../jetbrains/kotlinx/dataframe/api/sum.kt    |   4 +-
 .../kotlinx/dataframe/impl/NumberTypeUtils.kt |  20 ++
 .../aggregation/aggregators/Aggregator.kt     |  82 +++++-
 .../aggregation/aggregators/AggregatorBase.kt |  92 +++++-
 .../aggregators/AggregatorOptionSwitch.kt     |  73 +++--
 .../aggregators/AggregatorProvider.kt         |  26 +-
 .../aggregation/aggregators/Aggregators.kt    | 266 ++++++++++++++++--
 .../aggregators/FlatteningAggregator.kt       |  71 +++++
 .../aggregators/MergedValuesAggregator.kt     |  42 ---
 .../aggregators/NumbersAggregator.kt          |  37 ---
 .../aggregators/TwoStepAggregator.kt          |  95 ++++++-
 .../aggregators/TwoStepNumbersAggregator.kt   | 126 +++++++++
 .../dataframe/impl/aggregation/getColumns.kt  |   2 +-
 .../impl/aggregation/modes/ofRowExpression.kt |   7 +-
 .../jetbrains/kotlinx/dataframe/math/mean.kt  |  11 +-
 .../jetbrains/kotlinx/dataframe/math/std.kt   |   7 +
 .../jetbrains/kotlinx/dataframe/math/sum.kt   |   7 +
 .../kotlinx/dataframe/api/describe.kt         |   5 +-
 .../dataframe/documentation/DelimParams.kt    |  20 +-
 .../jetbrains/kotlinx/dataframe/io/readCsv.kt |  12 +-
 .../kotlinx/dataframe/io/readDelim.kt         |  12 +-
 .../jetbrains/kotlinx/dataframe/io/readTsv.kt |  12 +-
 .../kotlinx/dataframe/io/DelimCsvTsvTests.kt  |  47 ++--
 27 files changed, 895 insertions(+), 222 deletions(-)
 create mode 100644 core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
 delete mode 100644 core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
 delete mode 100644 core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
 create mode 100644 core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt

diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
index a265a20f9e..fc58dc626c 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMaxOf(): T = rowMaxOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(interComparableColumns())
+public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.max.aggregateFor(this, columns)
@@ -134,7 +134,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.maxByOrNull(column: KProperty<C?>
 
 // region GroupBy
 
-public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(interComparableColumns())
+public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.max.aggregateFor(this, columns)
@@ -246,7 +246,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.maxBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.maxFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index f2cdbb390e..b69f76b18c 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -12,7 +12,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
@@ -39,8 +39,9 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOf(noinline
 // region DataRow
 
 public fun AnyRow.rowMedianOrNull(): Any? =
-    Aggregators.median.aggregateMixed(
-        values().filterIsInstance<Comparable<Any?>>().asIterable(),
+    Aggregators.median.aggregateCalculatingType(
+        values = values().filterIsInstance<Comparable<Any?>>().asIterable(),
+        valueTypes = df().columns().filter { it.valuesAreComparable() }.map { it.type() }.toSet(),
     )
 
 public fun AnyRow.rowMedian(): Any = rowMedianOrNull().suggestIfNull("rowMedian")
@@ -54,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMedianOf(): T =
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(interComparableColumns())
+public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.median.aggregateFor(this, columns)
@@ -104,7 +105,7 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.medianOf(
 
 // region GroupBy
 
-public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(interComparableColumns())
+public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.median.aggregateFor(this, columns)
@@ -146,7 +147,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.medianOf(
 
 // region Pivot
 
-public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, interComparableColumns())
+public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Pivot<T>.medianFor(
     separate: Boolean = false,
@@ -190,7 +191,7 @@ public inline fun <T, reified R : Comparable<R>> Pivot<T>.medianOf(
 // region PivotGroupBy
 
 public fun <T> PivotGroupBy<T>.median(separate: Boolean = false): DataFrame<T> =
-    medianFor(separate, interComparableColumns())
+    medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> PivotGroupBy<T>.medianFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
index d1cae852aa..01f74ce41e 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMinOf(): T = rowMinOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(interComparableColumns())
+public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.min.aggregateFor(this, columns)
@@ -134,7 +134,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.minByOrNull(column: KProperty<C?>
 
 // region GroupBy
 
-public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(interComparableColumns())
+public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
     Aggregators.min.aggregateFor(this, columns)
@@ -247,7 +247,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.minBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.minFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
index 34c482612d..1ced2969e5 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
@@ -12,7 +12,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
@@ -63,7 +63,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowPercentileOf(percentile:
 // region DataFrame
 
 public fun <T> DataFrame<T>.percentile(percentile: Double): DataRow<T> =
-    percentileFor(percentile, interComparableColumns())
+    percentileFor(percentile, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.percentileFor(
     percentile: Double,
@@ -128,7 +128,7 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.percentileOf(
 // region GroupBy
 
 public fun <T> Grouped<T>.percentile(percentile: Double): DataFrame<T> =
-    percentileFor(percentile, interComparableColumns())
+    percentileFor(percentile, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Grouped<T>.percentileFor(
     percentile: Double,
@@ -184,7 +184,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.percentileOf(
 // region Pivot
 
 public fun <T> Pivot<T>.percentile(percentile: Double, separate: Boolean = false): DataRow<T> =
-    percentileFor(percentile, separate, interComparableColumns())
+    percentileFor(percentile, separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Pivot<T>.percentileFor(
     percentile: Double,
@@ -238,7 +238,7 @@ public inline fun <T, reified R : Comparable<R>> Pivot<T>.percentileOf(
 // region PivotGroupBy
 
 public fun <T> PivotGroupBy<T>.percentile(percentile: Double, separate: Boolean = false): DataFrame<T> =
-    percentileFor(percentile, separate, interComparableColumns())
+    percentileFor(percentile, separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> PivotGroupBy<T>.percentileFor(
     percentile: Double,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
index af9bea3657..9e84d74622 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt
@@ -44,9 +44,9 @@ public inline fun <T, reified R : Number> DataColumn<T>.sumOf(crossinline expres
 // region DataRow
 
 public fun AnyRow.rowSum(): Number =
-    Aggregators.sum.aggregateMixed(
+    Aggregators.sum.aggregateCalculatingType(
         values = values().filterIsInstance<Number>(),
-        types = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
+        valueTypes = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
     ) ?: 0
 
 public inline fun <reified T : Number> AnyRow.rowSumOf(): T = values().filterIsInstance<T>().sum(typeOf<T>())
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
index 1166742813..881f1e4741 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
@@ -194,3 +194,23 @@ internal fun Iterable<Number>.convertToUnifiedNumberType(
         converter(it) ?: error("Can not convert $it to $commonNumberType")
     }
 }
+
+/** Converts the elements of the given iterable of numbers into a common numeric type based on complexity.
+ * The common numeric type is determined using the provided [commonNumberType] parameter
+ * or calculated with [Iterable.unifiedNumberType][kotlin.collections.Iterable.unifiedNumberType] from the iterable's elements if not explicitly specified.
+ *
+ * @param commonNumberType The desired common numeric type to convert the elements to.
+ *   This is determined by default using the types of the elements in the iterable.
+ * @return A new iterable of numbers where each element is converted to the specified or inferred common number type.
+ * @throws IllegalStateException if an element cannot be converted to the common number type.
+ * @see UnifyingNumbers */
+@JvmName("convertToUnifiedNumberTypeSequence")
+@Suppress("UNCHECKED_CAST")
+internal fun Sequence<Number>.convertToUnifiedNumberType(
+    commonNumberType: KType = asIterable().types().unifiedNumberType(),
+): Sequence<Number> {
+    val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
+    return map {
+        converter(it) ?: error("Can not convert $it to $commonNumberType")
+    }
+}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
index dcd88a15a7..20f6a27b93 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
@@ -2,23 +2,91 @@ 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.
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * 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.
+ * @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]`?`.
+ */
 @PublishedApi
-internal interface Aggregator<C, R> {
+internal interface Aggregator<in Value, out Return> {
 
+    /** The name of this aggregator. */
     val name: String
 
-    fun aggregate(column: DataColumn<C?>): R?
-
+    /** If `true`, [Value][Value]`  ==  ` [Return][Return]. */
     val preservesType: Boolean
 
-    fun aggregate(columns: Iterable<DataColumn<C?>>): R?
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, 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?
+
+    /**
+     * 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?
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     */
+    fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
+
+    /**
+     * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     *
+     * @param values The values to be aggregated.
+     * @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!).
+     */
+    fun aggregateCalculatingType(values: Iterable<Value>, valueTypes: Set<KType>? = null): Return?
 
-    fun aggregate(values: Iterable<C>, type: KType): R?
+    /**
+     * Function that can give the return type of [aggregate] as [KType], given the type of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param type The type of the input values.
+     * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
 }
 
 @PublishedApi
-internal fun <T> Aggregator<*, *>.cast(): Aggregator<T, T> = this as Aggregator<T, T>
+internal fun <Type> Aggregator<*, *>.cast(): Aggregator<Type, Type> = this as Aggregator<Type, Type>
 
 @PublishedApi
-internal fun <T, P> Aggregator<*, *>.cast2(): Aggregator<T, P> = this as Aggregator<T, P>
+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?
+
+/** Common case for [CalculateReturnTypeOrNull], preserves return type, but makes it nullable for empty inputs. */
+internal val preserveReturnTypeNullIfEmpty: CalculateReturnTypeOrNull = { type, emptyInput ->
+    type.withNullability(emptyInput)
+}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
index 1deb052b2f..2f210e12b8 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
@@ -3,19 +3,95 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 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 kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
-internal abstract class AggregatorBase<C, R>(
+/**
+ * Abstract base class for [aggregators][Aggregator].
+ *
+ * Aggregators are used to compute a single value from an [Iterable] of values, a single [DataColumn],
+ * or multiple [DataColumns][DataColumn].
+ *
+ * @param name The name of this aggregator.
+ * @param aggregator Functional argument for the [aggregate] function.
+ */
+internal abstract class AggregatorBase<in Value, out Return>(
     override val name: String,
-    protected val aggregator: (Iterable<C>, KType) -> R?,
-) : Aggregator<C, R> {
+    protected val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    protected val aggregator: Aggregate<Value, Return>,
+) : Aggregator<Value, Return> {
 
-    override fun aggregate(column: DataColumn<C?>): R? =
-        if (column.hasNulls()) {
-            aggregate(column.asSequence().filterNotNull().asIterable(), column.type())
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * Uses [aggregator] to compute the result.
+     *
+     * When the exact [type] is unknown, use [aggregateCalculatingType].
+     */
+    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.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * Uses [getReturnTypeOrNull] to calculate the return type.
+     *
+     * @param type The type of the input values.
+     * @param emptyInput If `true`, the input values are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    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 by default, then [aggregate] (with [Iterable] and [KType]) is called.
+     */
+    @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),
+        )
+
+    /** Special case of [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator.aggregate] with [Iterable] that calculates the common type of the values at runtime.
+     * This is a heavy operation and should be avoided when possible.
+     *
+     * @param values The values to be aggregated.
+     * @param valueTypes The types of the values.
+     *   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? {
+        val commonType = if (valueTypes != null) {
+            valueTypes.commonType(false)
         } else {
-            aggregate(column.asIterable() as Iterable<C>, column.type())
+            var hasNulls = false
+            val classes = values.mapNotNull {
+                if (it == null) {
+                    hasNulls = true
+                    null
+                } else {
+                    it.javaClass.kotlin
+                }
+            }
+            classes.commonType(hasNulls)
         }
+        return aggregate(values, commonType)
+    }
 
-    override fun aggregate(values: Iterable<C>, type: KType): R? = aggregator(values, type)
+    /**
+     * 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?
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index 45cb01be19..7a2b2d4496 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -1,33 +1,72 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
-import kotlin.reflect.KProperty
-
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require a single parameter.
+ *
+ * Aggregators are cached by their parameter value.
+ * @see AggregatorOptionSwitch2
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch<P, C, R>(val name: String, val getAggregator: (P) -> AggregatorProvider<C, R>) {
+internal class AggregatorOptionSwitch1<in Param1, out AggregatorType : Aggregator<*, *>>(
+    val name: String,
+    val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
+) {
 
-    private val cache = mutableMapOf<P, Aggregator<C, R>>()
+    private val cache: MutableMap<Param1, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option: P) = cache.getOrPut(option) { getAggregator(option).create(name) }
+    operator fun invoke(param1: Param1): AggregatorType =
+        cache.getOrPut(param1) {
+            getAggregator(param1).create(name)
+        }
 
-    class Factory<P, C, R>(val getAggregator: (P) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch1].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch1.Factory { param1: Param1 ->
+     *   MyAggregator.Factory(param1)
+     * }
+     */
+    class Factory<in Param1, out AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch1(name, getAggregator)
+        })
 }
 
+/**
+ * Wrapper around an [aggregator factory][AggregatorProvider] for aggregators that require two parameters.
+ *
+ * Aggregators are cached by their parameter values.
+ * @see AggregatorOptionSwitch1
+ */
 @PublishedApi
-internal class AggregatorOptionSwitch2<P1, P2, C, R>(
+internal class AggregatorOptionSwitch2<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
     val name: String,
-    val getAggregator: (P1, P2) -> AggregatorProvider<C, R>,
+    val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
 ) {
 
-    private val cache = mutableMapOf<Pair<P1, P2>, Aggregator<C, R>>()
+    private val cache: MutableMap<Pair<Param1, Param2>, AggregatorType> = mutableMapOf()
 
-    operator fun invoke(option1: P1, option2: P2) =
-        cache.getOrPut(option1 to option2) {
-            getAggregator(option1, option2).create(name)
+    operator fun invoke(param1: Param1, param2: Param2): AggregatorType =
+        cache.getOrPut(param1 to param2) {
+            getAggregator(param1, param2).create(name)
         }
 
-    class Factory<P1, P2, C, R>(val getAggregator: (P1, P2) -> AggregatorProvider<C, R>) {
-        operator fun getValue(obj: Any?, property: KProperty<*>) = AggregatorOptionSwitch2(property.name, getAggregator)
-    }
+    /**
+     * Creates [AggregatorOptionSwitch2].
+     *
+     * Used like:
+     * ```kt
+     * val myAggregator by AggregatorOptionSwitch2.Factory { param1: Param1, param2: Param2 ->
+     *   MyAggregator.Factory(param1, param2)
+     * }
+     * ```
+     */
+    class Factory<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
+        val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
+    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
+            AggregatorOptionSwitch2(name, getAggregator)
+        })
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
index a8265a8175..9c16fcdb59 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorProvider.kt
@@ -2,9 +2,27 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import kotlin.reflect.KProperty
 
-internal interface AggregatorProvider<C, R> {
+/**
+ * Common interface for providers or "factory" objects that create anything of type [T].
+ *
+ * When implemented, this allows the object to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myNamedValue by MyFactory
+ * ```
+ */
+internal fun interface Provider<out T> {
 
-    operator fun getValue(obj: Any?, property: KProperty<*>): Aggregator<C, R> = create(property.name)
-
-    fun create(name: String): Aggregator<C, R>
+    fun create(name: String): T
 }
+
+internal operator fun <T> Provider<T>.getValue(obj: Any?, property: KProperty<*>): T = create(property.name)
+
+/**
+ * Common interface for providers of [Aggregators][Aggregator] or "factory" objects that create aggregators.
+ *
+ * When implemented, this allows an aggregator to be created using the `by` delegate, to give it a name, like:
+ * ```kt
+ * val myAggregator by MyAggregator.Factory
+ * ```
+ */
+internal fun interface AggregatorProvider<out AggregatorType : Aggregator<*, *>> : Provider<AggregatorType>
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
index 4c90f286d8..7ff4da9ee5 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
@@ -1,52 +1,270 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.math.mean
+import org.jetbrains.kotlinx.dataframe.math.meanTypeConversion
 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.stdTypeConversion
 import org.jetbrains.kotlinx.dataframe.math.sum
-import kotlin.reflect.KType
+import org.jetbrains.kotlinx.dataframe.math.sumTypeConversion
 
 @PublishedApi
 internal object Aggregators {
 
-    private fun <C> preservesType(aggregate: Iterable<C>.(KType) -> C?) =
-        TwoStepAggregator.Factory(aggregate, aggregate, true)
+    /**
+     * Factory for a simple aggregator that preserves the type of the input values.
+     *
+     * A slightly more advanced [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] implementation.
+     *
+     * Nulls are filtered from columns.
+     *
+     * When called on multiple columns, this [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] works in two steps:
+     * First, it aggregates within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn]/[Iterable] ([stepOneAggregator]) with their (given) type,
+     * and then in between different columns ([stepTwoAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.stepTwoAggregator]) using the results of the first and the newly
+     * calculated common type of those results.
+     *
+     * ```
+     * Iterable<Column<Value?>>
+     *     -> Iterable<Iterable<Value>> // nulls filtered out
+     *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
+     *     -> Iterable<Return> // nulls filtered out
+     *     -> stepTwoAggregator(Iterable<Return>, common valueType)
+     *     -> Return?
+     * ```
+     *
+     * It can also be used as a "simple" aggregator by providing the same function for both steps,
+     * requires [preservesType][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.preservesType] be set to `true`.
+     *
+     * See [FlatteningAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator] for different behavior for multiple columns.
+     *
+     * @param name The name of this aggregator.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param stepOneAggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
+     * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+     *   It is run on the results of [stepOneAggregator].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    private fun <Type> twoStepPreservingType(aggregator: Aggregate<Type, Type>) =
+        TwoStepAggregator.Factory(
+            getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
+            stepOneAggregator = aggregator,
+            stepTwoAggregator = aggregator,
+            preservesType = true,
+        )
 
-    private fun <C : Any, R> mergedValues(aggregate: Iterable<C?>.(KType) -> R?) =
-        MergedValuesAggregator.Factory(aggregate, true)
+    /**
+     * Factory for a simple aggregator that changes the type of the input values.
+     *
+     * A slightly more advanced [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] implementation.
+     *
+     * Nulls are filtered from columns.
+     *
+     * When called on multiple columns, this [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] works in two steps:
+     * First, it aggregates within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn]/[Iterable] ([stepOneAggregator]) with their (given) type,
+     * and then in between different columns ([stepTwoAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.stepTwoAggregator]) using the results of the first and the newly
+     * calculated common type of those results.
+     *
+     * ```
+     * Iterable<Column<Value?>>
+     *     -> Iterable<Iterable<Value>> // nulls filtered out
+     *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
+     *     -> Iterable<Return> // nulls filtered out
+     *     -> stepTwoAggregator(Iterable<Return>, common valueType)
+     *     -> Return?
+     * ```
+     *
+     * It can also be used as a "simple" aggregator by providing the same function for both steps,
+     * requires [preservesType][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.preservesType] be set to `true`.
+     *
+     * See [FlatteningAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator] for different behavior for multiple columns.
+     *
+     * @param name The name of this aggregator.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param stepOneAggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
+     * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+     *   It is run on the results of [stepOneAggregator].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    private fun <Value, Return> twoStepChangingType(
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        stepOneAggregator: Aggregate<Value, Return>,
+        stepTwoAggregator: Aggregate<Return, Return>,
+    ) = TwoStepAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
+        stepOneAggregator = stepOneAggregator,
+        stepTwoAggregator = stepTwoAggregator,
+        preservesType = false,
+    )
 
-    private fun <C : Any, R> mergedValuesChangingTypes(aggregate: Iterable<C?>.(KType) -> R?) =
-        MergedValuesAggregator.Factory(aggregate, false)
+    /**
+     * Factory for a flattening aggregator that preserves the type of the input values.
+     *
+     * Simple [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] implementation with flattening behavior for multiple columns.
+     *
+     * Nulls are filtered from columns.
+     *
+     * When called on multiple columns,
+     * the columns are flattened into a single list of values, filtering nulls as usual;
+     * then the aggregation function is called with their common type.
+     *
+     * ```
+     * Iterable<Column<Value?>>
+     *     -> Iterable<Value> // flattened without nulls
+     *     -> aggregator(Iterable<Value>, common colType)
+     *     -> Return?
+     * ```
+     *
+     * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
+     * the median, percentile, and standard deviation.
+     *
+     * See [TwoStepAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator] for different behavior for multiple columns.
+     *
+     * @param name The name of this aggregator.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @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].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    private fun <Type> flatteningPreservingTypes(aggregate: Aggregate<Type, Type>) =
+        FlatteningAggregator.Factory(
+            getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
+            aggregator = aggregate,
+            preservesType = true,
+        )
 
-    private fun <C, R> changesType(aggregate1: Iterable<C>.(KType) -> R, aggregate2: Iterable<R>.(KType) -> R) =
-        TwoStepAggregator.Factory(aggregate1, aggregate2, false)
+    /**
+     * Factory for a flattening aggregator that changes the type of the input values.
+     *
+     * Simple [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] implementation with flattening behavior for multiple columns.
+     *
+     * Nulls are filtered from columns.
+     *
+     * When called on multiple columns,
+     * the columns are flattened into a single list of values, filtering nulls as usual;
+     * then the aggregation function is called with their common type.
+     *
+     * ```
+     * Iterable<Column<Value?>>
+     *     -> Iterable<Value> // flattened without nulls
+     *     -> aggregator(Iterable<Value>, common colType)
+     *     -> Return?
+     * ```
+     *
+     * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
+     * the median, percentile, and standard deviation.
+     *
+     * See [TwoStepAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator] for different behavior for multiple columns.
+     *
+     * @param name The name of this aggregator.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @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].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    private fun <Value, Return> flatteningChangingTypes(
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        aggregate: Aggregate<Value, Return>,
+    ) = FlatteningAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
+        aggregator = aggregate,
+        preservesType = false,
+    )
 
-    private fun extendsNumbers(aggregate: Iterable<Number>.(KType) -> Number?) = NumbersAggregator.Factory(aggregate)
+    /**
+     * Factory for a two-step aggregator that works only with numbers.
+     *
+     * [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] made specifically for number calculations.
+     *
+     * Nulls are filtered from columns.
+     *
+     * When called on multiple columns (with potentially different [Number] types),
+     * this [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] works in two steps:
+     *
+     * First, it aggregates within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn]/[Iterable] with their (given) [Number] type,
+     * and then between different columns
+     * using the results of the first and the newly calculated [unified number][org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers] type of those results.
+     *
+     * ```
+     * Iterable<Column<Number?>>
+     *     -> Iterable<Iterable<Number>> // nulls filtered out
+     *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+     *     -> Iterable<Return> // nulls filtered out
+     *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+     *     -> Return?
+     * ```
+     *
+     * @param name The name of this aggregator.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param aggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepNumbersAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
+     *   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(
+        getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        aggregate: Aggregate<Number, Return>,
+    ) = TwoStepNumbersAggregator.Factory(
+        getReturnTypeOrNull = getReturnTypeOrNull,
+        aggregate = aggregate,
+    )
 
-    private fun <P, C, R> withOption(getAggregator: (P) -> AggregatorProvider<C, R>) =
-        AggregatorOptionSwitch.Factory(getAggregator)
+    /** Wrapper around an [aggregator factory][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorProvider] for aggregators that require a single parameter.
+     *
+     * Aggregators are cached by their parameter value.
+     * @see AggregatorOptionSwitch2 */
+    private fun <Param1, AggregatorType : Aggregator<*, *>> withOneOption(
+        getAggregator: (Param1) -> AggregatorProvider<AggregatorType>,
+    ) = AggregatorOptionSwitch1.Factory(getAggregator)
 
-    private fun <P1, P2, C, R> withOption2(getAggregator: (P1, P2) -> AggregatorProvider<C, R>) =
-        AggregatorOptionSwitch2.Factory(getAggregator)
+    /** Wrapper around an [aggregator factory][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorProvider] for aggregators that require two parameters.
+     *
+     * Aggregators are cached by their parameter values.
+     * @see AggregatorOptionSwitch1 */
+    private fun <Param1, Param2, AggregatorType : Aggregator<*, *>> withTwoOptions(
+        getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
+    ) = AggregatorOptionSwitch2.Factory(getAggregator)
 
-    val min by preservesType<Comparable<Any?>> { minOrNull() }
+    // T: Comparable<T> -> T?
+    val min by twoStepPreservingType<Comparable<Any?>> {
+        minOrNull()
+    }
 
-    val max by preservesType<Comparable<Any?>> { maxOrNull() }
+    // T: Comparable<T> -> T?
+    val max by twoStepPreservingType<Comparable<Any?>> {
+        maxOrNull()
+    }
 
-    val std by withOption2<Boolean, Int, Number, Double> { skipNA, ddof ->
-        mergedValuesChangingTypes { std(it, skipNA, ddof) }
+    // T: Number? -> Double
+    val std by withTwoOptions { skipNA: Boolean, ddof: Int ->
+        flatteningChangingTypes<Number, Double>(stdTypeConversion) { type ->
+            std(type, skipNA, ddof)
+        }
     }
 
-    val mean by withOption<Boolean, Number, Double> { skipNA ->
-        changesType({ mean(it, skipNA) }) { mean(skipNA) }
+    // step one: T: Number? -> Double
+    // step two: Double -> Double
+    val mean by withOneOption { skipNA: Boolean ->
+        twoStepChangingType(
+            getReturnTypeOrNull = meanTypeConversion,
+            stepOneAggregator = { type -> mean(type, skipNA) },
+            stepTwoAggregator = { mean(skipNA) },
+        )
     }
 
-    val percentile by withOption<Double, Comparable<Any?>, Comparable<Any?>> { percentile ->
-        mergedValuesChangingTypes { type -> percentile(percentile, type) }
+    // T: Comparable<T>? -> T
+    val percentile by withOneOption { percentile: Double ->
+        flatteningPreservingTypes<Comparable<Any?>> { type ->
+            percentile(percentile, type)
+        }
     }
 
-    val median by mergedValues<Comparable<Any?>, Comparable<Any?>> { median(it) }
+    // T: Comparable<T>? -> T
+    val median by flatteningPreservingTypes<Comparable<Any?>> { type ->
+        median(type)
+    }
 
-    val sum by extendsNumbers { sum(it) }
+    // T: Number -> T
+    val sum by twoStepForNumbers(sumTypeConversion) { type ->
+        sum(type)
+    }
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
new file mode 100644
index 0000000000..b259339a69
--- /dev/null
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
@@ -0,0 +1,71 @@
+package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
+
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.full.withNullability
+
+/**
+ * Simple [Aggregator] implementation with flattening behavior for multiple columns.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns,
+ * the columns are flattened into a single list of values, filtering nulls as usual;
+ * then the aggregation function is called with their common type.
+ *
+ * ```
+ * Iterable<Column<Value?>>
+ *     -> Iterable<Value> // flattened without nulls
+ *     -> aggregator(Iterable<Value>, common colType)
+ *     -> Return?
+ * ```
+ *
+ * This is essential for aggregators that depend on the distribution of all values across the dataframe, like
+ * the median, percentile, and standard deviation.
+ *
+ * See [TwoStepAggregator] for different behavior for multiple columns.
+ *
+ * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
+ * @param aggregator Functional argument for the [aggregate] function.
+ *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
+ * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+ */
+internal class FlatteningAggregator<in Value, out Return>(
+    name: String,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    aggregator: Aggregate<Value, Return>,
+    override val preservesType: Boolean,
+) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, aggregator) {
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     * 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? {
+        val commonType = columns.map { it.type() }.commonType().withNullability(false)
+        val allValues = columns.asSequence().flatMap { it.values() }.filterNotNull()
+        return aggregate(allValues.asIterable(), commonType)
+    }
+
+    /**
+     * Creates [FlatteningAggregator].
+     *
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
+     * @param aggregator Functional argument for the [aggregate] function.
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+     */
+    class Factory<in Value, out Return>(
+        private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        private val aggregator: Aggregate<Value, Return>,
+        private val preservesType: Boolean,
+    ) : AggregatorProvider<FlatteningAggregator<Value, Return>> by AggregatorProvider({ name ->
+            FlatteningAggregator(
+                name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
+                aggregator = aggregator,
+                preservesType = preservesType,
+            )
+        })
+}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
deleted file mode 100644
index 135ba0a5ec..0000000000
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/MergedValuesAggregator.kt
+++ /dev/null
@@ -1,42 +0,0 @@
-package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
-
-import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.impl.commonType
-import kotlin.reflect.KProperty
-import kotlin.reflect.KType
-
-internal class MergedValuesAggregator<C : Any, R>(
-    name: String,
-    val aggregateWithType: (Iterable<C?>, KType) -> R?,
-    override val preservesType: Boolean,
-) : AggregatorBase<C, R>(name, aggregateWithType) {
-
-    override fun aggregate(columns: Iterable<DataColumn<C?>>): R? {
-        val commonType = columns.map { it.type() }.commonType()
-        val allValues = columns.flatMap { it.values() }
-        return aggregateWithType(allValues, commonType)
-    }
-
-    fun aggregateMixed(values: Iterable<C?>): R? {
-        var hasNulls = false
-        val classes = values.mapNotNull {
-            if (it == null) {
-                hasNulls = true
-                null
-            } else {
-                it.javaClass.kotlin
-            }
-        }
-        return aggregateWithType(values, classes.commonType(hasNulls))
-    }
-
-    class Factory<C : Any, R>(
-        private val aggregateWithType: (Iterable<C?>, KType) -> R?,
-        private val preservesType: Boolean,
-    ) : AggregatorProvider<C, R> {
-        override fun create(name: String) = MergedValuesAggregator(name, aggregateWithType, preservesType)
-
-        override operator fun getValue(obj: Any?, property: KProperty<*>): MergedValuesAggregator<C, R> =
-            create(property.name)
-    }
-}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
deleted file mode 100644
index 00ef22febe..0000000000
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt
+++ /dev/null
@@ -1,37 +0,0 @@
-package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
-
-import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
-import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
-import kotlin.reflect.KProperty
-import kotlin.reflect.KType
-
-internal class NumbersAggregator(name: String, aggregate: (Iterable<Number>, KType) -> Number?) :
-    AggregatorBase<Number, Number>(name, aggregate) {
-
-    override fun aggregate(columns: Iterable<DataColumn<Number?>>): Number? =
-        aggregateMixed(
-            values = columns.mapNotNull { aggregate(it) },
-            types = columns.map { it.type() }.toSet(),
-        )
-
-    class Factory(private val aggregate: Iterable<Number>.(KType) -> Number?) : AggregatorProvider<Number, Number> {
-        override fun create(name: String) = NumbersAggregator(name, aggregate)
-
-        override operator fun getValue(obj: Any?, property: KProperty<*>): NumbersAggregator = create(property.name)
-    }
-
-    /**
-     * Can aggregate numbers with different types by first converting them to a compatible type.
-     */
-    @Suppress("UNCHECKED_CAST")
-    fun aggregateMixed(values: Iterable<Number>, types: Set<KType>): Number? {
-        val commonType = types.unifiedNumberType()
-        return aggregate(
-            values = values.convertToUnifiedNumberType(commonType),
-            type = commonType,
-        )
-    }
-
-    override val preservesType = false
-}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
index 9d01169d02..b095708a2a 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
@@ -2,26 +2,91 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.impl.commonType
-import kotlin.reflect.KType
+import kotlin.reflect.full.starProjectedType
+import kotlin.reflect.full.withNullability
 
-internal class TwoStepAggregator<C, R>(
+/**
+ * A slightly more advanced [Aggregator] implementation.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns, this [Aggregator] works in two steps:
+ * First, it aggregates within a [DataColumn]/[Iterable] ([stepOneAggregator]) with their (given) type,
+ * and then in between different columns ([stepTwoAggregator]) using the results of the first and the newly
+ * calculated common type of those results.
+ *
+ * ```
+ * Iterable<Column<Value?>>
+ *     -> Iterable<Iterable<Value>> // nulls filtered out
+ *     -> stepOneAggregator(Iterable<Value>, colType) // called on each iterable
+ *     -> Iterable<Return> // nulls filtered out
+ *     -> stepTwoAggregator(Iterable<Return>, common valueType)
+ *     -> Return?
+ * ```
+ *
+ * It can also be used as a "simple" aggregator by providing the same function for both steps,
+ * requires [preservesType] be set to `true`.
+ *
+ * See [FlatteningAggregator] for different behavior for multiple columns.
+ *
+ * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
+ * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+ * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+ *   It is run on the results of [stepOneAggregator].
+ * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
+ */
+internal class TwoStepAggregator<in Value, out Return>(
     name: String,
-    aggregateWithType: (Iterable<C>, KType) -> R?,
-    private val aggregateValues: (Iterable<R>, KType) -> R?,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    stepOneAggregator: Aggregate<Value, Return>,
+    private val stepTwoAggregator: Aggregate<Return, Return>,
     override val preservesType: Boolean,
-) : AggregatorBase<C, R>(name, aggregateWithType) {
+) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
-    override fun aggregate(columns: Iterable<DataColumn<C?>>): R? {
-        val columnValues = columns.mapNotNull { aggregate(it) }
-        val commonType = columnValues.map { it.javaClass.kotlin }.commonType(false)
-        return aggregateValues(columnValues, commonType)
+    /**
+     * 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.
+     *
+     * Post-step-one types are calculated by [calculateReturnTypeOrNull].
+     */
+    override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return? {
+        val (values, types) = columns.mapNotNull { col ->
+            // uses stepOneAggregator
+            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)
     }
 
-    class Factory<C, R>(
-        private val aggregateWithType: (Iterable<C>, KType) -> R?,
-        private val aggregateValues: (Iterable<R>, KType) -> R?,
+    /**
+     * Creates [TwoStepAggregator].
+     *
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
+     * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+     * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
+     *   It is run on the results of [stepOneAggregator].
+     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][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 preservesType: Boolean,
-    ) : AggregatorProvider<C, R> {
-        override fun create(name: String) = TwoStepAggregator(name, aggregateWithType, aggregateValues, preservesType)
-    }
+    ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
+            TwoStepAggregator(
+                name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
+                stepOneAggregator = stepOneAggregator,
+                stepTwoAggregator = stepTwoAggregator,
+                preservesType = preservesType,
+            )
+        })
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
new file mode 100644
index 0000000000..4515d14bc9
--- /dev/null
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
@@ -0,0 +1,126 @@
+package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
+
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
+import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+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
+
+/**
+ * [Aggregator] made specifically for number calculations.
+ *
+ * Nulls are filtered from columns.
+ *
+ * When called on multiple columns (with potentially different [Number] types),
+ * this [Aggregator] works in two steps:
+ *
+ * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type,
+ * and then between different columns
+ * using the results of the first and the newly calculated [unified number][UnifyingNumbers] type of those results.
+ *
+ * ```
+ * Iterable<Column<Number?>>
+ *     -> Iterable<Iterable<Number>> // nulls filtered out
+ *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+ *     -> Iterable<Return> // nulls filtered out
+ *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+ *     -> Return?
+ * ```
+ *
+ * @param name The name of this aggregator.
+ * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
+ * @param aggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
+ *   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>(
+    name: String,
+    getReturnTypeOrNull: CalculateReturnTypeOrNull,
+    aggregator: Aggregate<Number, Return>,
+) : AggregatorBase<Number, Return>(name, getReturnTypeOrNull, aggregator) {
+
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * Uses [aggregator] to compute the result.
+     *
+     * When the exact [type] is unknown, use [aggregateCalculatingType].
+     */
+    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 super.aggregate(values, type)
+    }
+
+    /**
+     * Aggregates the data in the multiple given columns and computes a single resulting value.
+     *
+     * This function calls [aggregator] on each column and then again on the results.
+     *
+     * After the first aggregation, the number types are found by [calculateReturnTypeOrNull] and then
+     * unified using [aggregateCalculatingType].
+     */
+    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]
+     * of the values at runtime and converts all numbers to this type before aggregating.
+     * This is a heavy operation and should be avoided when possible.
+     *
+     * @param values The numbers to be aggregated.
+     * @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!).
+     */
+    @Suppress("UNCHECKED_CAST")
+    override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
+        val commonType = (valueTypes ?: values.types()).unifiedNumberType().withNullability(false)
+        return aggregate(
+            values = values.convertToUnifiedNumberType(commonType),
+            type = commonType,
+        )
+    }
+
+    override val preservesType = false
+
+    /**
+     * Creates [TwoStepNumbersAggregator].
+     *
+     * @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>(
+        private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
+        private val aggregate: Aggregate<Number, Return>,
+    ) : AggregatorProvider<TwoStepNumbersAggregator<Return>> by AggregatorProvider({ name ->
+            TwoStepNumbersAggregator(
+                name = name,
+                getReturnTypeOrNull = getReturnTypeOrNull,
+                aggregator = aggregate,
+            )
+        })
+}
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
index 6f514d95eb..b7b2c1052d 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/getColumns.kt
@@ -14,7 +14,7 @@ internal inline fun <T> Aggregatable<T>.remainingColumns(
     crossinline predicate: (AnyCol) -> Boolean,
 ): ColumnsSelector<T, Any?> = remainingColumnsSelector().filter { predicate(it.data) }
 
-internal fun <T> Aggregatable<T>.interComparableColumns() =
+internal fun <T> Aggregatable<T>.intraComparableColumns() =
     remainingColumns { it.valuesAreComparable() } as ColumnsSelector<T, Comparable<Any?>>
 
 internal fun <T> Aggregatable<T>.numberColumns() = remainingColumns { it.isNumber() } as ColumnsSelector<T, Number?>
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
index 4d43fb6128..299587a3b9 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
@@ -30,7 +30,7 @@ internal inline fun <C, reified V, R> Aggregator<V, R>.aggregateOf(
 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) } // TODO: inline
+): R? = (this as Aggregator<C, R>).aggregateOf(frame.rows()) { expression(it, it) }
 
 @PublishedApi
 internal fun <T, C, R> Aggregator<*, R>.aggregateOfDelegated(
@@ -50,7 +50,7 @@ internal inline fun <T, reified C, R> Aggregator<*, R>.of(
 
 @PublishedApi
 internal inline fun <C, reified V, R> Aggregator<V, R>.of(data: DataColumn<C>, crossinline expression: (C) -> V): R? =
-    aggregateOf(data.values()) { expression(it) } // TODO: inline
+    aggregateOf(data.values()) { expression(it) }
 
 @PublishedApi
 internal inline fun <T, reified C, reified R> Aggregator<*, R>.aggregateOf(
@@ -75,7 +75,8 @@ internal inline fun <T, reified C, reified R> Grouped<T>.aggregateOf(
     val type = typeOf<R>()
     return aggregateInternal {
         val value = aggregator.aggregateOf(df, expression)
-        yield(path, value, type, null, false)
+        val inferType = !aggregator.preservesType
+        yield(path, value, type, null, inferType)
     }
 }
 
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
index a1ab845624..fc5fb70b70 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -1,18 +1,20 @@
 package org.jetbrains.kotlinx.dataframe.math
 
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
 @PublishedApi
-internal fun <T : Number> Iterable<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Double =
+internal fun <T : Number> Iterable<T?>.mean(type: KType, skipNA: Boolean = skipNA_default): Double =
     asSequence().mean(type, skipNA)
 
 @Suppress("UNCHECKED_CAST")
-internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA_default): Double {
+internal fun <T : Number> Sequence<T?>.mean(type: KType, skipNA: Boolean = skipNA_default): Double {
     if (type.isMarkedNullable) {
         return filterNotNull().mean(type.withNullability(false), skipNA)
     }
@@ -43,6 +45,11 @@ internal fun <T : Number> Sequence<T>.mean(type: KType, skipNA: Boolean = skipNA
     }
 }
 
+// T: Number? -> Double
+internal val meanTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
+    typeOf<Double>()
+}
+
 internal fun Sequence<Double>.mean(skipNA: Boolean = skipNA_default): Double {
     var count = 0
     var sum: Double = 0.toDouble()
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
index 052556ba59..148c9ece23 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
@@ -2,11 +2,13 @@ package org.jetbrains.kotlinx.dataframe.math
 
 import org.jetbrains.kotlinx.dataframe.api.ddof_default
 import org.jetbrains.kotlinx.dataframe.api.skipNA_default
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
+import kotlin.reflect.typeOf
 
 @Suppress("UNCHECKED_CAST")
 @PublishedApi
@@ -35,6 +37,11 @@ internal fun <T : Number> Iterable<T?>.std(
     }
 }
 
+// T: Number? -> Double
+internal val stdTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
+    typeOf<Double>()
+}
+
 @JvmName("doubleStd")
 internal fun Iterable<Double>.std(skipNA: Boolean = skipNA_default, ddof: Int = ddof_default): Double =
     varianceAndMean(skipNA)?.std(ddof) ?: Double.NaN
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
index 08dae78937..1b03221988 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
@@ -1,8 +1,10 @@
 package org.jetbrains.kotlinx.dataframe.math
 
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.CalculateReturnTypeOrNull
 import java.math.BigDecimal
 import java.math.BigInteger
 import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 
 @PublishedApi
 internal fun <T, R : Number> Iterable<T>.sumOf(type: KType, selector: (T) -> R?): R {
@@ -95,6 +97,11 @@ internal fun <T : Number> Iterable<T?>.sum(type: KType): T =
         else -> throw IllegalArgumentException("sum is not supported for $type")
     }
 
+// T: Number? -> T
+internal val sumTypeConversion: CalculateReturnTypeOrNull = { type, _ ->
+    type.withNullability(false)
+}
+
 @PublishedApi
 internal fun Iterable<BigDecimal>.sum(): BigDecimal {
     var sum: BigDecimal = BigDecimal.ZERO
diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
index 04694ad901..71f049f5ca 100644
--- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
+++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/describe.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import io.kotest.matchers.doubles.shouldBeNaN
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.alsoDebug
 import org.junit.Test
@@ -64,8 +65,8 @@ class DescribeTests {
             nulls shouldBe 0
             top shouldBe 1
             freq shouldBe 1
-            mean.isNaN() shouldBe true
-            std.isNaN() shouldBe true
+            mean.shouldBeNaN()
+            std.shouldBeNaN()
             min shouldBe 1.0 // TODO should be NaN too?
             p25 shouldBe 1.75
             median shouldBe 3.0
diff --git a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/DelimParams.kt b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/DelimParams.kt
index 9a069fd714..c15dbc8b74 100644
--- a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/DelimParams.kt
+++ b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/DelimParams.kt
@@ -17,16 +17,28 @@ import org.jetbrains.kotlinx.dataframe.io.QuoteMode
 @Suppress("ktlint:standard:class-naming", "ClassName", "KDocUnresolvedReference")
 internal object DelimParams {
 
-    /** @param path The file path to read. Can also be compressed as `.gz` or `.zip`, see [Compression]. */
+    /**
+     * @param path The file path to read.
+     *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
+     */
     interface PATH_READ
 
-    /** @param file The file to read. Can also be compressed as `.gz` or `.zip`, see [Compression]. */
+    /**
+     * @param file The file to read.
+     *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
+     */
     interface FILE_READ
 
-    /** @param url The URL from which to fetch the data. Can also be compressed as `.gz` or `.zip`, see [Compression]. */
+    /**
+     * @param url The URL from which to fetch the data.
+     *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
+     */
     interface URL_READ
 
-    /** @param fileOrUrl The file path or URL to read the data from. Can also be compressed as `.gz` or `.zip`, see [Compression]. */
+    /**
+     * @param fileOrUrl The file path or URL to read the data from.
+     *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
+     */
     interface FILE_OR_URL_READ
 
     /** @param inputStream Represents the file to read. */
diff --git a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readCsv.kt b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readCsv.kt
index 814baa5718..0e6274c3b9 100644
--- a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readCsv.kt
+++ b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readCsv.kt
@@ -65,7 +65,8 @@ import kotlin.io.path.inputStream
  *
  * [DataFrame.readCsvStr][readCsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param path The file path to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param path The file path to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -216,7 +217,8 @@ public fun DataFrame.Companion.readCsv(
  *
  * [DataFrame.readCsvStr][readCsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param file The file to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param file The file to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -367,7 +369,8 @@ public fun DataFrame.Companion.readCsv(
  *
  * [DataFrame.readCsvStr][readCsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param url The URL from which to fetch the data. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param url The URL from which to fetch the data.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -518,7 +521,8 @@ public fun DataFrame.Companion.readCsv(
  *
  * [DataFrame.readCsvStr][readCsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param fileOrUrl The file path or URL to read the data from. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param fileOrUrl The file path or URL to read the data from.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
diff --git a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readDelim.kt b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readDelim.kt
index 329ef00cb5..65f46899bf 100644
--- a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readDelim.kt
+++ b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readDelim.kt
@@ -71,7 +71,8 @@ import kotlin.io.path.inputStream
  *
  * [DataFrame.readDelimStr][readDelimStr]`("a,b,c", delimiter = ",")`
  *
- * @param path The file path to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param path The file path to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -222,7 +223,8 @@ public fun DataFrame.Companion.readDelim(
  *
  * [DataFrame.readDelimStr][readDelimStr]`("a,b,c", delimiter = ",")`
  *
- * @param file The file to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param file The file to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -373,7 +375,8 @@ public fun DataFrame.Companion.readDelim(
  *
  * [DataFrame.readDelimStr][readDelimStr]`("a,b,c", delimiter = ",")`
  *
- * @param url The URL from which to fetch the data. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param url The URL from which to fetch the data.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -524,7 +527,8 @@ public fun DataFrame.Companion.readDelim(
  *
  * [DataFrame.readDelimStr][readDelimStr]`("a,b,c", delimiter = ",")`
  *
- * @param fileOrUrl The file path or URL to read the data from. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param fileOrUrl The file path or URL to read the data from.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: ','.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
diff --git a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readTsv.kt b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readTsv.kt
index 0acbede3e1..7e6ad6c7d0 100644
--- a/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readTsv.kt
+++ b/dataframe-csv/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readTsv.kt
@@ -65,7 +65,8 @@ import kotlin.io.path.inputStream
  *
  * [DataFrame.readTsvStr][readTsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param path The file path to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param path The file path to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: '\t'.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -216,7 +217,8 @@ public fun DataFrame.Companion.readTsv(
  *
  * [DataFrame.readTsvStr][readTsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param file The file to read. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param file The file to read.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: '\t'.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -367,7 +369,8 @@ public fun DataFrame.Companion.readTsv(
  *
  * [DataFrame.readTsvStr][readTsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param url The URL from which to fetch the data. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param url The URL from which to fetch the data.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: '\t'.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
@@ -518,7 +521,8 @@ public fun DataFrame.Companion.readTsv(
  *
  * [DataFrame.readTsvStr][readTsvStr]`("a,b,c", delimiter = ",")`
  *
- * @param fileOrUrl The file path or URL to read the data from. Can also be compressed as `.gz` or `.zip`, see [Compression].
+ * @param fileOrUrl The file path or URL to read the data from.
+ *   Can also be compressed as `.gz` or `.zip`, see [Compression][org.jetbrains.kotlinx.dataframe.io.Compression].
  * @param delimiter The field delimiter character. Default: '\t'.
  *
  *   Ignored if [hasFixedWidthColumns] is `true`.
diff --git a/dataframe-csv/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/DelimCsvTsvTests.kt b/dataframe-csv/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/DelimCsvTsvTests.kt
index a5903eab10..31abcfd041 100644
--- a/dataframe-csv/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/DelimCsvTsvTests.kt
+++ b/dataframe-csv/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/DelimCsvTsvTests.kt
@@ -518,29 +518,32 @@ class DelimCsvTsvTests {
 
         dutchDf["price"].type() shouldBe typeOf<Double?>()
 
-        // while negative numbers in RTL languages cannot be parsed, thanks to Java, others work
-        @Language("csv")
-        val arabicCsv =
-            """
-            الاسم; السعر;
-            أ;١٢٫٤٥;
-            ب;١٣٫٣٥;
-            ج;١٠٠٫١٢٣;
-            د;٢٠٤٫٢٣٥;
-            هـ;ليس رقم;
-            و;null;
-            """.trimIndent()
-
-        val easternArabicDf = DataFrame.readCsvStr(
-            arabicCsv,
-            delimiter = ';',
-            parserOptions = ParserOptions(
-                locale = Locale.forLanguageTag("ar-001"),
-            ),
-        )
+        // skipping this test on windows due to lack of support for Arabic locales
+        if (!System.getProperty("os.name").startsWith("Windows")) {
+            // while negative numbers in RTL languages cannot be parsed thanks to Java, others work
+            @Language("csv")
+            val arabicCsv =
+                """
+                الاسم; السعر;
+                أ;١٢٫٤٥;
+                ب;١٣٫٣٥;
+                ج;١٠٠٫١٢٣;
+                د;٢٠٤٫٢٣٥;
+                هـ;ليس رقم;
+                و;null;
+                """.trimIndent()
+
+            val easternArabicDf = DataFrame.readCsvStr(
+                arabicCsv,
+                delimiter = ';',
+                parserOptions = ParserOptions(
+                    locale = Locale.forLanguageTag("ar-001"),
+                ),
+            )
 
-        easternArabicDf["السعر"].type() shouldBe typeOf<Double?>()
-        easternArabicDf["الاسم"].type() shouldBe typeOf<String>() // apparently not a char
+            easternArabicDf["السعر"].type() shouldBe typeOf<Double?>()
+            easternArabicDf["الاسم"].type() shouldBe typeOf<String>() // apparently not a char
+        }
     }
 
     @Test

From 8611254f78cb7553f8fecee634dd119828757988 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Fri, 7 Mar 2025 15:02:12 +0100
Subject: [PATCH 09/18] added missed casts to median/percentile. Could result
 in Comparable<Any?> columns

---
 .../kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt     | 5 +++--
 .../kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt | 4 ++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index b69f76b18c..81698f8b1c 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -16,6 +16,7 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.columns.toComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull
@@ -141,7 +142,7 @@ public fun <T, C : Comparable<C>> Grouped<T>.median(vararg columns: KProperty<C?
 public inline fun <T, reified R : Comparable<R>> Grouped<T>.medianOf(
     name: String? = null,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.median.aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.median.cast<R?>().aggregateOf(this, name, expression)
 
 // endregion
 
@@ -228,6 +229,6 @@ public fun <T, C : Comparable<C>> PivotGroupBy<T>.median(vararg columns: KProper
 
 public inline fun <T, reified R : Comparable<R>> PivotGroupBy<T>.medianOf(
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.median.aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.median.cast<R?>().aggregateOf(this, expression)
 
 // endregion
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
index 1ced2969e5..b0a08bef6d 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/percentile.kt
@@ -177,7 +177,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.percentileOf(
     percentile: Double,
     name: String? = null,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.percentile(percentile).aggregateOf(this, name, expression)
+): DataFrame<T> = Aggregators.percentile(percentile).cast<R?>().aggregateOf(this, name, expression)
 
 // endregion
 
@@ -289,6 +289,6 @@ public fun <T, C : Comparable<C>> PivotGroupBy<T>.percentile(
 public inline fun <T, reified R : Comparable<R>> PivotGroupBy<T>.percentileOf(
     percentile: Double,
     crossinline expression: RowExpression<T, R?>,
-): DataFrame<T> = Aggregators.percentile(percentile).aggregateOf(this, expression)
+): DataFrame<T> = Aggregators.percentile(percentile).cast<R?>().aggregateOf(this, expression)
 
 // endregion

From c8e4d21d550b74bb5ddcc765723d6de45b012ce7 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 12:49:55 +0100
Subject: [PATCH 10/18] linting

---
 core/api/core.api                                            | 5 +----
 .../kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt     | 1 -
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/core/api/core.api b/core/api/core.api
index a01cbed2a3..5489bca2d4 100644
--- a/core/api/core.api
+++ b/core/api/core.api
@@ -5303,15 +5303,12 @@ 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 static synthetic fun aggregateCalculatingType$default (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;Ljava/lang/Iterable;Ljava/util/Set;ILjava/lang/Object;)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
 }
 
-public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator$DefaultImpls {
-	public static synthetic fun aggregateCalculatingType$default (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;Ljava/lang/Iterable;Ljava/util/Set;ILjava/lang/Object;)Ljava/lang/Object;
-}
-
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorKt {
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
 	public static final fun cast2 (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;)Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index 81698f8b1c..450c0137dc 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -16,7 +16,6 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.columns.toComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull

From 0b5988b57600fc29e8358fb44f92d27e7be38172 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 13:43:57 +0100
Subject: [PATCH 11/18] TwoStepNumbersAggregator now always unifies numbers

---
 .../aggregators/TwoStepNumbersAggregator.kt   | 56 +++++++++++--------
 1 file changed, 34 insertions(+), 22 deletions(-)

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
index 4515d14bc9..0b296baad1 100644
--- 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
@@ -13,22 +13,23 @@ import kotlin.reflect.typeOf
 
 /**
  * [Aggregator] made specifically for number calculations.
+ * Mixed number types are [unified][UnifyingNumbers].
  *
  * Nulls are filtered from columns.
  *
- * When called on multiple columns (with potentially different [Number] types),
+ * When called on multiple columns (with potentially mixed [Number] types),
  * this [Aggregator] works in two steps:
  *
- * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type,
- * and then between different columns
+ * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type
+ * (potentially unifying the types), and then between different columns
  * using the results of the first and the newly calculated [unified number][UnifyingNumbers] type of those results.
  *
  * ```
  * Iterable<Column<Number?>>
  *     -> Iterable<Iterable<Number>> // nulls filtered out
- *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+ *     -> aggregator(Iterable<specific Number>, unified number type of common colType) // called on each iterable
  *     -> Iterable<Return> // nulls filtered out
- *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+ *     -> aggregator(Iterable<specific Return>, unified number type of common valueType)
  *     -> Return?
  * ```
  *
@@ -44,22 +45,6 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
     aggregator: Aggregate<Number, Return>,
 ) : AggregatorBase<Number, Return>(name, getReturnTypeOrNull, aggregator) {
 
-    /**
-     * Base function of [Aggregator].
-     *
-     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
-     *
-     * Uses [aggregator] to compute the result.
-     *
-     * When the exact [type] is unknown, use [aggregateCalculatingType].
-     */
-    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 super.aggregate(values, type)
-    }
-
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
      *
@@ -85,6 +70,33 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * 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? {
+        require(type.isSubtypeOf(typeOf<Number?>())) {
+            "${TwoStepNumbersAggregator::class.simpleName}: Type $type is not a subtype of Number?"
+        }
+
+        // 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
+        // or calling aggregateCalculatingType with all known number types
+        return if (type.withNullability(false) == typeOf<Number>()) {
+            aggregateCalculatingType(values)
+        } else {
+            super.aggregate(values, type)
+        }
+    }
+
     /**
      * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
      * of the values at runtime and converts all numbers to this type before aggregating.
@@ -99,7 +111,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
     @Suppress("UNCHECKED_CAST")
     override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
         val commonType = (valueTypes ?: values.types()).unifiedNumberType().withNullability(false)
-        return aggregate(
+        return super.aggregate(
             values = values.convertToUnifiedNumberType(commonType),
             type = commonType,
         )

From a53588f074f9440f28f3996a298f9d7eecfcf69d Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 14:26:46 +0100
Subject: [PATCH 12/18] added UnifiedNumberTypeOptions such that the number
 aggregator can run on primitives only

---
 .../documentation/UnifyingNumbers.kt          |  15 ++-
 .../kotlinx/dataframe/impl/NumberTypeUtils.kt | 122 +++++++++++++-----
 .../aggregators/TwoStepNumbersAggregator.kt   |  19 ++-
 3 files changed, 115 insertions(+), 41 deletions(-)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
index 42db06463c..2d1a0125c5 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
@@ -1,5 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.documentation
 
+import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions
+
 /**
  * ## Unifying Numbers
  *
@@ -12,16 +14,23 @@ package org.jetbrains.kotlinx.dataframe.documentation
  * For each number type in the graph, it holds that a number of that type can be expressed lossless by
  * a number of a more complex type (any of its parents).
  * This is either because the more complex type has a larger range or higher precision (in terms of bits).
+ *
+ * There are variants of this graph that exclude some types, such as `BigDecimal` and `BigInteger`.
+ * In these cases `Double` could be considered the most complex type.
+ * `Long`/`ULong` and `Double` could be joined to `Double`,
+ * potentially losing a little precision, but a warning will be given.
+ *
+ * See [UnifiedNumberTypeOptions] for these settings.
  */
 internal interface UnifyingNumbers {
 
     /**
      * ```
-     *            BigDecimal
+     *           (BigDecimal)
      *            /      \\
-     *      BigInteger    \\
+     *     (BigInteger)   \\
      *        /   \\        \\
-     *    ULong   Long    Double
+     * <~ ULong   Long ~> Double ..
      * ..   |    /   |   /   |  \\..
      *   \\  |   /    |  /    |
      *     UInt     Int    Float
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
index 1745147f47..73f8544997 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
@@ -9,6 +9,27 @@ import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+/**
+ * @param useBigNumbers Whether to include [BigDecimal] and [BigInteger] in the graph.
+ *   If set to `false`, consider setting [allowLongToDoubleConversion] to `true` to have a single "most complex" number type.
+ * @param allowLongToDoubleConversion Whether to allow [Long]/[ULong] -> [Double] conversion.
+ *   If set to `true`, [Long] and [ULong] will be joined to [Double] in the graph.
+ */
+internal data class UnifiedNumberTypeOptions(val useBigNumbers: Boolean, val allowLongToDoubleConversion: Boolean) {
+    companion object {
+        val DEFAULT = UnifiedNumberTypeOptions(
+            useBigNumbers = true,
+            allowLongToDoubleConversion = false,
+        )
+        val PRIMITIVES_ONLY = UnifiedNumberTypeOptions(
+            useBigNumbers = false,
+            allowLongToDoubleConversion = true,
+        )
+    }
+}
+
+private val unifiedNumberTypeGraphs = mutableMapOf<UnifiedNumberTypeOptions, DirectedAcyclicGraph<KType>>()
+
 /**
  * Number type graph, structured in terms of number complexity.
  * A number can always be expressed lossless by a number of a more complex type (any of its parents).
@@ -17,46 +38,57 @@ import kotlin.reflect.typeOf
  *
  * For any two numbers, we can find the nearest common ancestor in this graph
  * by calling [DirectedAcyclicGraph.findNearestCommonVertex].
+ *
+ * @param options See [UnifiedNumberTypeOptions]
  * @see getUnifiedNumberClass
  * @see unifiedNumberClass
  * @see UnifyingNumbers
  */
-internal val unifiedNumberTypeGraph: DirectedAcyclicGraph<KType> by lazy {
-    buildDag {
-        addEdge(typeOf<BigDecimal>(), typeOf<BigInteger>())
-        addEdge(typeOf<BigDecimal>(), typeOf<Double>())
+internal fun getUnifiedNumberTypeGraph(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): DirectedAcyclicGraph<KType> =
+    unifiedNumberTypeGraphs.getOrPut(options) {
+        buildDag {
+            if (options.useBigNumbers) {
+                addEdge(typeOf<BigDecimal>(), typeOf<BigInteger>())
+                addEdge(typeOf<BigDecimal>(), typeOf<Double>())
 
-        addEdge(typeOf<BigInteger>(), typeOf<ULong>())
-        addEdge(typeOf<BigInteger>(), typeOf<Long>())
+                addEdge(typeOf<BigInteger>(), typeOf<ULong>())
+                addEdge(typeOf<BigInteger>(), typeOf<Long>())
+            }
+            if (options.allowLongToDoubleConversion) {
+                addEdge(typeOf<Double>(), typeOf<Long>())
+                addEdge(typeOf<Double>(), typeOf<ULong>())
+            }
 
-        addEdge(typeOf<ULong>(), typeOf<UInt>())
+            addEdge(typeOf<ULong>(), typeOf<UInt>())
 
-        addEdge(typeOf<Long>(), typeOf<UInt>())
-        addEdge(typeOf<Long>(), typeOf<Int>())
+            addEdge(typeOf<Long>(), typeOf<UInt>())
+            addEdge(typeOf<Long>(), typeOf<Int>())
 
-        addEdge(typeOf<Double>(), typeOf<Int>())
-        addEdge(typeOf<Double>(), typeOf<Float>())
-        addEdge(typeOf<Double>(), typeOf<UInt>())
+            addEdge(typeOf<Double>(), typeOf<Int>())
+            addEdge(typeOf<Double>(), typeOf<Float>())
+            addEdge(typeOf<Double>(), typeOf<UInt>())
 
-        addEdge(typeOf<UInt>(), typeOf<UShort>())
+            addEdge(typeOf<UInt>(), typeOf<UShort>())
 
-        addEdge(typeOf<Int>(), typeOf<UShort>())
-        addEdge(typeOf<Int>(), typeOf<Short>())
+            addEdge(typeOf<Int>(), typeOf<UShort>())
+            addEdge(typeOf<Int>(), typeOf<Short>())
 
-        addEdge(typeOf<Float>(), typeOf<Short>())
-        addEdge(typeOf<Float>(), typeOf<UShort>())
+            addEdge(typeOf<Float>(), typeOf<Short>())
+            addEdge(typeOf<Float>(), typeOf<UShort>())
 
-        addEdge(typeOf<UShort>(), typeOf<UByte>())
+            addEdge(typeOf<UShort>(), typeOf<UByte>())
 
-        addEdge(typeOf<Short>(), typeOf<UByte>())
-        addEdge(typeOf<Short>(), typeOf<Byte>())
+            addEdge(typeOf<Short>(), typeOf<UByte>())
+            addEdge(typeOf<Short>(), typeOf<Byte>())
+        }
     }
-}
 
-/** @include [unifiedNumberTypeGraph] */
-internal val unifiedNumberClassGraph: DirectedAcyclicGraph<KClass<*>> by lazy {
-    unifiedNumberTypeGraph.map { it.classifier as KClass<*> }
-}
+/** @include [getUnifiedNumberTypeGraph] */
+internal fun getUnifiedNumberClassGraph(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): DirectedAcyclicGraph<KClass<*>> = getUnifiedNumberTypeGraph(options).map { it.classifier as KClass<*> }
 
 /**
  * Determines the nearest common numeric type, in terms of complexity, between two given classes/types.
@@ -67,11 +99,16 @@ internal val unifiedNumberClassGraph: DirectedAcyclicGraph<KClass<*>> by lazy {
  *
  * @param first The first numeric type to compare. Can be null, in which case the second to is returned.
  * @param second The second numeric to compare. Cannot be null.
+ * @param options See [UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the two input classes.
  *   If no common class is found, [IllegalStateException] is thrown.
  * @see UnifyingNumbers
  */
-internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
+internal fun getUnifiedNumberType(
+    first: KType?,
+    second: KType,
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KType {
     if (first == null) return second
 
     val firstWithoutNullability = first.withNullability(false)
@@ -80,7 +117,7 @@ internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
     val result = if (firstWithoutNullability == secondWithoutNullability) {
         firstWithoutNullability
     } else {
-        unifiedNumberTypeGraph.findNearestCommonVertex(firstWithoutNullability, secondWithoutNullability)
+        getUnifiedNumberTypeGraph(options).findNearestCommonVertex(firstWithoutNullability, secondWithoutNullability)
             ?: error("Can not find common number type for $first and $second")
     }
 
@@ -89,13 +126,17 @@ internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
 
 /** @include [getUnifiedNumberType] */
 @Suppress("IntroduceWhenSubject")
-internal fun getUnifiedNumberClass(first: KClass<*>?, second: KClass<*>): KClass<*> =
+internal fun getUnifiedNumberClass(
+    first: KClass<*>?,
+    second: KClass<*>,
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KClass<*> =
     when {
         first == null -> second
 
         first == second -> first
 
-        else -> unifiedNumberClassGraph.findNearestCommonVertex(first, second)
+        else -> getUnifiedNumberClassGraph(options).findNearestCommonVertex(first, second)
             ?: error("Can not find common number type for $first and $second")
     }
 
@@ -106,16 +147,25 @@ internal fun getUnifiedNumberClass(first: KClass<*>?, second: KClass<*>): KClass
  * but unless the input solely exists of unsigned numbers, it will never be returned.
  * Meaning, given a [Number] in the input, the output will always be a [Number].
  *
+ * @param options See [UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the input types.
  *   If no common type is found, it returns [Number].
  * @see UnifyingNumbers
  */
-internal fun Iterable<KType>.unifiedNumberType(): KType =
-    fold(null as KType?, ::getUnifiedNumberType) ?: typeOf<Number>()
+internal fun Iterable<KType>.unifiedNumberType(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KType =
+    fold(null as KType?) { a, b ->
+        getUnifiedNumberType(a, b, options)
+    } ?: typeOf<Number>()
 
 /** @include [unifiedNumberType] */
-internal fun Iterable<KClass<*>>.unifiedNumberClass(): KClass<*> =
-    fold(null as KClass<*>?, ::getUnifiedNumberClass) ?: Number::class
+internal fun Iterable<KClass<*>>.unifiedNumberClass(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KClass<*> =
+    fold(null as KClass<*>?) { a, b ->
+        getUnifiedNumberClass(a, b, options)
+    } ?: Number::class
 
 /**
  * Converts the elements of the given iterable of numbers into a common numeric type based on complexity.
@@ -130,7 +180,8 @@ internal fun Iterable<KClass<*>>.unifiedNumberClass(): KClass<*> =
  */
 @Suppress("UNCHECKED_CAST")
 internal fun Iterable<Number>.convertToUnifiedNumberType(
-    commonNumberType: KType = this.types().unifiedNumberType(),
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+    commonNumberType: KType = this.types().unifiedNumberType(options),
 ): Iterable<Number> {
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
@@ -142,7 +193,8 @@ internal fun Iterable<Number>.convertToUnifiedNumberType(
 @JvmName("convertToUnifiedNumberTypeSequence")
 @Suppress("UNCHECKED_CAST")
 internal fun Sequence<Number>.convertToUnifiedNumberType(
-    commonNumberType: KType = asIterable().types().unifiedNumberType(),
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+    commonNumberType: KType = asIterable().types().unifiedNumberType(options),
 ): Sequence<Number> {
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
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
index 0b296baad1..a310150711 100644
--- 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
@@ -1,7 +1,9 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
+import io.github.oshai.kotlinlogging.KotlinLogging
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
+import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
 import org.jetbrains.kotlinx.dataframe.impl.types
 import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
@@ -11,9 +13,11 @@ import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+private val logger = KotlinLogging.logger { }
+
 /**
  * [Aggregator] made specifically for number calculations.
- * Mixed number types are [unified][UnifyingNumbers].
+ * Mixed number types are [unified][UnifyingNumbers] to [primitives][PRIMITIVES_ONLY].
  *
  * Nulls are filtered from columns.
  *
@@ -110,9 +114,18 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      */
     @Suppress("UNCHECKED_CAST")
     override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
-        val commonType = (valueTypes ?: values.types()).unifiedNumberType().withNullability(false)
+        val valueTypes = valueTypes ?: values.types()
+        val commonType = valueTypes
+            .unifiedNumberType(PRIMITIVES_ONLY)
+            .withNullability(false)
+
+        if (commonType == 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."
+            }
+        }
         return super.aggregate(
-            values = values.convertToUnifiedNumberType(commonType),
+            values = values.convertToUnifiedNumberType(commonNumberType = commonType),
             type = commonType,
         )
     }

From f922d9f7f764d8bd938131d76004435653819fb6 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 15:19:27 +0100
Subject: [PATCH 13/18] better exceptions for unsupported/mixed number types

---
 .../kotlinx/dataframe/impl/NumberTypeUtils.kt | 10 ++++++++++
 .../aggregators/TwoStepNumbersAggregator.kt   | 20 +++++++++++++++----
 .../kotlinx/dataframe/statistics/sum.kt       | 19 +++++++++++++-----
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
index 73f8544997..fca8cb23e7 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
@@ -201,3 +201,13 @@ internal fun Sequence<Number>.convertToUnifiedNumberType(
         converter(it) ?: error("Can not convert $it to $commonNumberType")
     }
 }
+
+internal val primitiveNumberTypes =
+    setOf(
+        typeOf<Byte>(),
+        typeOf<Short>(),
+        typeOf<Int>(),
+        typeOf<Long>(),
+        typeOf<Float>(),
+        typeOf<Double>(),
+    )
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
index a310150711..6f87b04dd3 100644
--- 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
@@ -5,6 +5,8 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
 import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
+import org.jetbrains.kotlinx.dataframe.impl.renderType
 import org.jetbrains.kotlinx.dataframe.impl.types
 import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
 import kotlin.reflect.KType
@@ -94,10 +96,14 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         // 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
         // or calling aggregateCalculatingType with all known number types
-        return if (type.withNullability(false) == typeOf<Number>()) {
-            aggregateCalculatingType(values)
-        } else {
-            super.aggregate(values, type)
+        return when (type.withNullability(false)) {
+            typeOf<Number>() -> aggregateCalculatingType(values)
+
+            !in primitiveNumberTypes -> throw IllegalArgumentException(
+                "Cannot calculate $name of ${renderType(type)}, only primitive numbers are supported.",
+            )
+
+            else -> super.aggregate(values, type)
         }
     }
 
@@ -124,6 +130,12 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
                 "Number unification of Long -> Double happened during aggregation. Loss of precision may have occurred."
             }
         }
+        if (commonType !in primitiveNumberTypes) {
+            throw IllegalArgumentException(
+                "Cannot calculate $name of ${renderType(commonType)}, only primitive numbers are supported.",
+            )
+        }
+
         return super.aggregate(
             values = values.convertToUnifiedNumberType(commonNumberType = commonType),
             type = commonType,
diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
index 513d7f4d19..1980e5da70 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.statistics
 
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.columnOf
@@ -7,8 +8,8 @@ import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.rowSum
 import org.jetbrains.kotlinx.dataframe.api.sum
 import org.jetbrains.kotlinx.dataframe.api.sumOf
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.junit.Test
-import java.math.BigDecimal
 
 class SumTests {
 
@@ -58,10 +59,10 @@ class SumTests {
         df.sumOf { value3() } shouldBe expected3
         df.sum(value1) shouldBe expected1
         df.sum(value2) shouldBe expected2
-        df.sum(value3) shouldBe expected3
+        // TODO sum rework, has Number in results df.sum(value3) shouldBe expected3
         df.sum { value1 } shouldBe expected1
         df.sum { value2 } shouldBe expected2
-        df.sum { value3 } shouldBe expected3
+        // TODO sum rework, has Number in results df.sum { value3 } shouldBe expected3
     }
 
     /** [Issue #1068](https://github.com/Kotlin/dataframe/issues/1068) */
@@ -78,9 +79,17 @@ class SumTests {
             it::class shouldBe Int::class
         }
 
+        // NOTE! lossy conversion from long -> double happens
         dataFrameOf("a", "b")(1.0, 2L)[0].rowSum().let {
-            it shouldBe (3.0.toBigDecimal())
-            it::class shouldBe BigDecimal::class
+            it shouldBe 3.0
+            it::class shouldBe Double::class
+        }
+    }
+
+    @Test
+    fun `unknown number type`() {
+        shouldThrow<IllegalArgumentException> {
+            columnOf(1.toBigDecimal(), 2.toBigDecimal()).toDataFrame().sum()
         }
     }
 }

From 87165bad6baa66b90504525d3f7b89302665d656 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 16:11:45 +0100
Subject: [PATCH 14/18] added back support for Nothing in
 TwoStepNumbersAggregator

---
 .../aggregators/TwoStepNumbersAggregator.kt          | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

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
index 6f87b04dd3..088bf9b352 100644
--- 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
@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
 import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+import org.jetbrains.kotlinx.dataframe.impl.nothingType
 import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
 import org.jetbrains.kotlinx.dataframe.impl.renderType
 import org.jetbrains.kotlinx.dataframe.impl.types
@@ -93,12 +94,15 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
             "${TwoStepNumbersAggregator::class.simpleName}: Type $type is not a subtype of Number?"
         }
 
-        // 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
-        // or calling aggregateCalculatingType with all known number types
         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
+            // or calling aggregateCalculatingType with all known number types
             typeOf<Number>() -> aggregateCalculatingType(values)
 
+            // Nothing can occur when values are empty
+            nothingType -> super.aggregate(values, type)
+
             !in primitiveNumberTypes -> throw IllegalArgumentException(
                 "Cannot calculate $name of ${renderType(type)}, only primitive numbers are supported.",
             )
@@ -130,7 +134,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
                 "Number unification of Long -> Double happened during aggregation. Loss of precision may have occurred."
             }
         }
-        if (commonType !in primitiveNumberTypes) {
+        if (commonType !in primitiveNumberTypes && commonType != nothingType) {
             throw IllegalArgumentException(
                 "Cannot calculate $name of ${renderType(commonType)}, only primitive numbers are supported.",
             )

From 851b1a6e2950f7a74de5c7758439cdde5fccbb11 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Mon, 10 Mar 2025 20:50:53 +0100
Subject: [PATCH 15/18] marked aggregateBy for removal

---
 .../dataframe/impl/aggregation/aggregators/Aggregator.kt        | 2 +-
 .../impl/aggregation/aggregators/TwoStepNumbersAggregator.kt    | 2 +-
 .../kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt     | 2 ++
 3 files changed, 4 insertions(+), 2 deletions(-)

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
index 20f6a27b93..3ed36e4363 100644
--- 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
@@ -53,7 +53,7 @@ internal interface Aggregator<in Value, out Return> {
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
-     * This is a heavy operation and should be avoided when possible.
+     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
      *
      * @param values The values to be aggregated.
      * @param valueTypes The types of the values.
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
index 088bf9b352..2e01bb6480 100644
--- 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
@@ -114,7 +114,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
     /**
      * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
      * of the values at runtime and converts all numbers to this type before aggregating.
-     * This is a heavy operation and should be avoided when possible.
+     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
      *
      * @param values The numbers to be aggregated.
      * @param valueTypes The types of the numbers.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
index 239ed236de..659180a3aa 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
@@ -3,12 +3,14 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.modes
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataFrameExpression
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.CandidateForRemoval
 import org.jetbrains.kotlinx.dataframe.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.namedValues
 
+@CandidateForRemoval
 internal fun <T> Grouped<T>.aggregateBy(body: DataFrameExpression<T, DataRow<T>?>): DataFrame<T> {
     require(this is GroupBy<*, T>)
     val keyColumns = keys.columnNames().toSet()

From 868b8a185e07f1b9bebc0d16851071686b21acd3 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Tue, 11 Mar 2025 12:23:00 +0100
Subject: [PATCH 16/18] update from master

---
 .../jetbrains/kotlinx/dataframe/api/max.kt    |   8 +-
 .../jetbrains/kotlinx/dataframe/api/median.kt |  15 +-
 .../jetbrains/kotlinx/dataframe/api/min.kt    |   8 +-
 .../documentation/UnifyingNumbers.kt          |  21 ++-
 .../kotlinx/dataframe/impl/NumberTypeUtils.kt | 146 +++++++++++++-----
 .../aggregation/aggregators/Aggregator.kt     |   2 +-
 .../aggregation/aggregators/AggregatorBase.kt |   2 +-
 .../aggregation/aggregators/Aggregators.kt    |  11 +-
 .../aggregators/TwoStepNumbersAggregator.kt   |  91 ++++++++---
 .../impl/aggregation/modes/aggregateBy.kt     |   2 +
 .../kotlinx/dataframe/statistics/sum.kt       |  19 ++-
 .../kotlinx/dataframe/types/UtilTests.kt      |   6 +-
 .../jetbrains/kotlinx/dataframe/api/median.kt |   1 -
 13 files changed, 230 insertions(+), 102 deletions(-)

diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
index 8d7d6b3b47..6de276a7e4 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/max.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMaxOf(): T = rowMaxOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(interComparableColumns())
+public fun <T> DataFrame<T>.max(): DataRow<T> = maxFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.max.aggregateFor(this, columns)
@@ -135,7 +135,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.maxByOrNull(column: KProperty<C?>
 // region GroupBy
 @Refine
 @Interpretable("GroupByMax1")
-public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(interComparableColumns())
+public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(intraComparableColumns())
 
 @Refine
 @Interpretable("GroupByMax0")
@@ -251,7 +251,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.maxBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = maxFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.maxFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index ac8d8a92f8..8da5194a7d 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
@@ -41,8 +41,9 @@ public inline fun <T, reified R : Comparable<R>> DataColumn<T>.medianOf(noinline
 // region DataRow
 
 public fun AnyRow.rowMedianOrNull(): Any? =
-    Aggregators.median.aggregateMixed(
-        values().filterIsInstance<Comparable<Any?>>().asIterable(),
+    Aggregators.median.aggregateCalculatingType(
+        values = values().filterIsInstance<Comparable<Any?>>().asIterable(),
+        valueTypes = df().columns().filter { it.valuesAreComparable() }.map { it.type() }.toSet(),
     )
 
 public fun AnyRow.rowMedian(): Any = rowMedianOrNull().suggestIfNull("rowMedian")
@@ -56,7 +57,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMedianOf(): T =
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(interComparableColumns())
+public fun <T> DataFrame<T>.median(): DataRow<T> = medianFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.medianFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.median.aggregateFor(this, columns)
@@ -107,7 +108,7 @@ public inline fun <T, reified R : Comparable<R>> DataFrame<T>.medianOf(
 // region GroupBy
 @Refine
 @Interpretable("GroupByMedian1")
-public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(interComparableColumns())
+public fun <T> Grouped<T>.median(): DataFrame<T> = medianFor(intraComparableColumns())
 
 @Refine
 @Interpretable("GroupByMedian0")
@@ -155,7 +156,7 @@ public inline fun <T, reified R : Comparable<R>> Grouped<T>.medianOf(
 
 // region Pivot
 
-public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, interComparableColumns())
+public fun <T> Pivot<T>.median(separate: Boolean = false): DataRow<T> = medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> Pivot<T>.medianFor(
     separate: Boolean = false,
@@ -199,7 +200,7 @@ public inline fun <T, reified R : Comparable<R>> Pivot<T>.medianOf(
 // region PivotGroupBy
 
 public fun <T> PivotGroupBy<T>.median(separate: Boolean = false): DataFrame<T> =
-    medianFor(separate, interComparableColumns())
+    medianFor(separate, intraComparableColumns())
 
 public fun <T, C : Comparable<C>> PivotGroupBy<T>.medianFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
index 0a9c79b5a1..c843cc871f 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/min.kt
@@ -14,7 +14,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.interComparableColumns
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
@@ -55,7 +55,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMinOf(): T = rowMinOfOrN
 
 // region DataFrame
 
-public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(interComparableColumns())
+public fun <T> DataFrame<T>.min(): DataRow<T> = minFor(intraComparableColumns())
 
 public fun <T, C : Comparable<C>> DataFrame<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataRow<T> =
     Aggregators.min.aggregateFor(this, columns)
@@ -135,7 +135,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.minByOrNull(column: KProperty<C?>
 // region GroupBy
 @Refine
 @Interpretable("GroupByMin1")
-public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(interComparableColumns())
+public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(intraComparableColumns())
 
 @Refine
 @Interpretable("GroupByMin0")
@@ -252,7 +252,7 @@ public fun <T, C : Comparable<C>> Pivot<T>.minBy(column: KProperty<C?>): Reduced
 
 // region PivotGroupBy
 
-public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, interComparableColumns())
+public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = minFor(separate, intraComparableColumns())
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.minFor(
     separate: Boolean = false,
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
index 6b1646828d..55bc3b7599 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/documentation/UnifyingNumbers.kt
@@ -1,5 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.documentation
 
+import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions
+
 /**
  * ## Unifying Numbers
  *
@@ -9,11 +11,11 @@ package org.jetbrains.kotlinx.dataframe.documentation
  * The order is top-down from the most complex type to the simplest one.
  *
  * ```
- *            BigDecimal
+ *           (BigDecimal)
  *            /      \
- *      BigInteger    \
+ *     (BigInteger)   \
  *        /   \        \
- *    ULong   Long    Double
+ * <~ ULong   Long ~> Double ..
  * ..   |    /   |   /   |  \..
  *   \  |   /    |  /    |
  *     UInt     Int    Float
@@ -27,16 +29,23 @@ package org.jetbrains.kotlinx.dataframe.documentation
  * For each number type in the graph, it holds that a number of that type can be expressed lossless by
  * a number of a more complex type (any of its parents).
  * This is either because the more complex type has a larger range or higher precision (in terms of bits).
+ *
+ * There are variants of this graph that exclude some types, such as `BigDecimal` and `BigInteger`.
+ * In these cases `Double` could be considered the most complex type.
+ * `Long`/`ULong` and `Double` could be joined to `Double`,
+ * potentially losing a little precision, but a warning will be given.
+ *
+ * See [UnifiedNumberTypeOptions] for these settings.
  */
 internal interface UnifyingNumbers {
 
     /**
      * ```
-     *            BigDecimal
+     *           (BigDecimal)
      *            /      \
-     *      BigInteger    \
+     *     (BigInteger)   \
      *        /   \        \
-     *    ULong   Long    Double
+     * <~ ULong   Long ~> Double ..
      * ..   |    /   |   /   |  \..
      *   \  |   /    |  /    |
      *     UInt     Int    Float
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
index 881f1e4741..c4e1a9679a 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/NumberTypeUtils.kt
@@ -9,16 +9,37 @@ import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+/**
+ * @param useBigNumbers Whether to include [BigDecimal] and [BigInteger] in the graph.
+ *   If set to `false`, consider setting [allowLongToDoubleConversion] to `true` to have a single "most complex" number type.
+ * @param allowLongToDoubleConversion Whether to allow [Long]/[ULong] -> [Double] conversion.
+ *   If set to `true`, [Long] and [ULong] will be joined to [Double] in the graph.
+ */
+internal data class UnifiedNumberTypeOptions(val useBigNumbers: Boolean, val allowLongToDoubleConversion: Boolean) {
+    companion object {
+        val DEFAULT = UnifiedNumberTypeOptions(
+            useBigNumbers = true,
+            allowLongToDoubleConversion = false,
+        )
+        val PRIMITIVES_ONLY = UnifiedNumberTypeOptions(
+            useBigNumbers = false,
+            allowLongToDoubleConversion = true,
+        )
+    }
+}
+
+private val unifiedNumberTypeGraphs = mutableMapOf<UnifiedNumberTypeOptions, DirectedAcyclicGraph<KType>>()
+
 /**
  * Number type graph, structured in terms of number complexity.
  * A number can always be expressed lossless by a number of a more complex type (any of its parents).
  *
  * ```
- *            BigDecimal
+ *           (BigDecimal)
  *            /      \
- *      BigInteger    \
+ *     (BigInteger)   \
  *        /   \        \
- *    ULong   Long    Double
+ * <~ ULong   Long ~> Double ..
  * ..   |    /   |   /   |  \..
  *   \  |   /    |  /    |
  *     UInt     Int    Float
@@ -32,51 +53,62 @@ import kotlin.reflect.typeOf
  *
  * For any two numbers, we can find the nearest common ancestor in this graph
  * by calling [DirectedAcyclicGraph.findNearestCommonVertex].
+ *
+ * @param options See [UnifiedNumberTypeOptions]
  * @see getUnifiedNumberClass
  * @see unifiedNumberClass
  * @see UnifyingNumbers
  */
-internal val unifiedNumberTypeGraph: DirectedAcyclicGraph<KType> by lazy {
-    buildDag {
-        addEdge(typeOf<BigDecimal>(), typeOf<BigInteger>())
-        addEdge(typeOf<BigDecimal>(), typeOf<Double>())
+internal fun getUnifiedNumberTypeGraph(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): DirectedAcyclicGraph<KType> =
+    unifiedNumberTypeGraphs.getOrPut(options) {
+        buildDag {
+            if (options.useBigNumbers) {
+                addEdge(typeOf<BigDecimal>(), typeOf<BigInteger>())
+                addEdge(typeOf<BigDecimal>(), typeOf<Double>())
 
-        addEdge(typeOf<BigInteger>(), typeOf<ULong>())
-        addEdge(typeOf<BigInteger>(), typeOf<Long>())
+                addEdge(typeOf<BigInteger>(), typeOf<ULong>())
+                addEdge(typeOf<BigInteger>(), typeOf<Long>())
+            }
+            if (options.allowLongToDoubleConversion) {
+                addEdge(typeOf<Double>(), typeOf<Long>())
+                addEdge(typeOf<Double>(), typeOf<ULong>())
+            }
 
-        addEdge(typeOf<ULong>(), typeOf<UInt>())
+            addEdge(typeOf<ULong>(), typeOf<UInt>())
 
-        addEdge(typeOf<Long>(), typeOf<UInt>())
-        addEdge(typeOf<Long>(), typeOf<Int>())
+            addEdge(typeOf<Long>(), typeOf<UInt>())
+            addEdge(typeOf<Long>(), typeOf<Int>())
 
-        addEdge(typeOf<Double>(), typeOf<Int>())
-        addEdge(typeOf<Double>(), typeOf<Float>())
-        addEdge(typeOf<Double>(), typeOf<UInt>())
+            addEdge(typeOf<Double>(), typeOf<Int>())
+            addEdge(typeOf<Double>(), typeOf<Float>())
+            addEdge(typeOf<Double>(), typeOf<UInt>())
 
-        addEdge(typeOf<UInt>(), typeOf<UShort>())
+            addEdge(typeOf<UInt>(), typeOf<UShort>())
 
-        addEdge(typeOf<Int>(), typeOf<UShort>())
-        addEdge(typeOf<Int>(), typeOf<Short>())
+            addEdge(typeOf<Int>(), typeOf<UShort>())
+            addEdge(typeOf<Int>(), typeOf<Short>())
 
-        addEdge(typeOf<Float>(), typeOf<Short>())
-        addEdge(typeOf<Float>(), typeOf<UShort>())
+            addEdge(typeOf<Float>(), typeOf<Short>())
+            addEdge(typeOf<Float>(), typeOf<UShort>())
 
-        addEdge(typeOf<UShort>(), typeOf<UByte>())
+            addEdge(typeOf<UShort>(), typeOf<UByte>())
 
-        addEdge(typeOf<Short>(), typeOf<UByte>())
-        addEdge(typeOf<Short>(), typeOf<Byte>())
+            addEdge(typeOf<Short>(), typeOf<UByte>())
+            addEdge(typeOf<Short>(), typeOf<Byte>())
+        }
     }
-}
 
 /** Number type graph, structured in terms of number complexity.
  * A number can always be expressed lossless by a number of a more complex type (any of its parents).
  *
  * ```
- *            BigDecimal
+ *           (BigDecimal)
  *            /      \
- *      BigInteger    \
+ *     (BigInteger)   \
  *        /   \        \
- *    ULong   Long    Double
+ * <~ ULong   Long ~> Double ..
  * ..   |    /   |   /   |  \..
  *   \  |   /    |  /    |
  *     UInt     Int    Float
@@ -90,12 +122,14 @@ internal val unifiedNumberTypeGraph: DirectedAcyclicGraph<KType> by lazy {
  *
  * For any two numbers, we can find the nearest common ancestor in this graph
  * by calling [DirectedAcyclicGraph.findNearestCommonVertex][org.jetbrains.kotlinx.dataframe.impl.DirectedAcyclicGraph.findNearestCommonVertex].
+ *
+ * @param options See [UnifiedNumberTypeOptions][org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions]
  * @see getUnifiedNumberClass
  * @see unifiedNumberClass
  * @see UnifyingNumbers */
-internal val unifiedNumberClassGraph: DirectedAcyclicGraph<KClass<*>> by lazy {
-    unifiedNumberTypeGraph.map { it.classifier as KClass<*> }
-}
+internal fun getUnifiedNumberClassGraph(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): DirectedAcyclicGraph<KClass<*>> = getUnifiedNumberTypeGraph(options).map { it.classifier as KClass<*> }
 
 /**
  * Determines the nearest common numeric type, in terms of complexity, between two given classes/types.
@@ -106,11 +140,16 @@ internal val unifiedNumberClassGraph: DirectedAcyclicGraph<KClass<*>> by lazy {
  *
  * @param first The first numeric type to compare. Can be null, in which case the second to is returned.
  * @param second The second numeric to compare. Cannot be null.
+ * @param options See [UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the two input classes.
  *   If no common class is found, [IllegalStateException] is thrown.
  * @see UnifyingNumbers
  */
-internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
+internal fun getUnifiedNumberType(
+    first: KType?,
+    second: KType,
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KType {
     if (first == null) return second
 
     val firstWithoutNullability = first.withNullability(false)
@@ -119,7 +158,7 @@ internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
     val result = if (firstWithoutNullability == secondWithoutNullability) {
         firstWithoutNullability
     } else {
-        unifiedNumberTypeGraph.findNearestCommonVertex(firstWithoutNullability, secondWithoutNullability)
+        getUnifiedNumberTypeGraph(options).findNearestCommonVertex(firstWithoutNullability, secondWithoutNullability)
             ?: error("Can not find common number type for $first and $second")
     }
 
@@ -134,17 +173,22 @@ internal fun getUnifiedNumberType(first: KType?, second: KType): KType {
  *
  * @param first The first numeric type to compare. Can be null, in which case the second to is returned.
  * @param second The second numeric to compare. Cannot be null.
+ * @param options See [UnifiedNumberTypeOptions][org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the two input classes.
  *   If no common class is found, [IllegalStateException] is thrown.
  * @see UnifyingNumbers */
 @Suppress("IntroduceWhenSubject")
-internal fun getUnifiedNumberClass(first: KClass<*>?, second: KClass<*>): KClass<*> =
+internal fun getUnifiedNumberClass(
+    first: KClass<*>?,
+    second: KClass<*>,
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KClass<*> =
     when {
         first == null -> second
 
         first == second -> first
 
-        else -> unifiedNumberClassGraph.findNearestCommonVertex(first, second)
+        else -> getUnifiedNumberClassGraph(options).findNearestCommonVertex(first, second)
             ?: error("Can not find common number type for $first and $second")
     }
 
@@ -155,12 +199,17 @@ internal fun getUnifiedNumberClass(first: KClass<*>?, second: KClass<*>): KClass
  * but unless the input solely exists of unsigned numbers, it will never be returned.
  * Meaning, given a [Number] in the input, the output will always be a [Number].
  *
+ * @param options See [UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the input types.
  *   If no common type is found, it returns [Number].
  * @see UnifyingNumbers
  */
-internal fun Iterable<KType>.unifiedNumberType(): KType =
-    fold(null as KType?, ::getUnifiedNumberType) ?: typeOf<Number>()
+internal fun Iterable<KType>.unifiedNumberType(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KType =
+    fold(null as KType?) { a, b ->
+        getUnifiedNumberType(a, b, options)
+    } ?: typeOf<Number>()
 
 /** Determines the nearest common numeric type, in terms of complexity, all types in [this].
  *
@@ -168,11 +217,16 @@ internal fun Iterable<KType>.unifiedNumberType(): KType =
  * but unless the input solely exists of unsigned numbers, it will never be returned.
  * Meaning, given a [Number] in the input, the output will always be a [Number].
  *
+ * @param options See [UnifiedNumberTypeOptions][org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions]
  * @return The nearest common numeric type between the input types.
  *   If no common type is found, it returns [Number].
  * @see UnifyingNumbers */
-internal fun Iterable<KClass<*>>.unifiedNumberClass(): KClass<*> =
-    fold(null as KClass<*>?, ::getUnifiedNumberClass) ?: Number::class
+internal fun Iterable<KClass<*>>.unifiedNumberClass(
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+): KClass<*> =
+    fold(null as KClass<*>?) { a, b ->
+        getUnifiedNumberClass(a, b, options)
+    } ?: Number::class
 
 /**
  * Converts the elements of the given iterable of numbers into a common numeric type based on complexity.
@@ -187,7 +241,8 @@ internal fun Iterable<KClass<*>>.unifiedNumberClass(): KClass<*> =
  */
 @Suppress("UNCHECKED_CAST")
 internal fun Iterable<Number>.convertToUnifiedNumberType(
-    commonNumberType: KType = this.types().unifiedNumberType(),
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+    commonNumberType: KType = this.types().unifiedNumberType(options),
 ): Iterable<Number> {
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
@@ -207,10 +262,21 @@ internal fun Iterable<Number>.convertToUnifiedNumberType(
 @JvmName("convertToUnifiedNumberTypeSequence")
 @Suppress("UNCHECKED_CAST")
 internal fun Sequence<Number>.convertToUnifiedNumberType(
-    commonNumberType: KType = asIterable().types().unifiedNumberType(),
+    options: UnifiedNumberTypeOptions = UnifiedNumberTypeOptions.DEFAULT,
+    commonNumberType: KType = asIterable().types().unifiedNumberType(options),
 ): Sequence<Number> {
     val converter = createConverter(typeOf<Number>(), commonNumberType)!! as (Number) -> Number?
     return map {
         converter(it) ?: error("Can not convert $it to $commonNumberType")
     }
 }
+
+internal val primitiveNumberTypes =
+    setOf(
+        typeOf<Byte>(),
+        typeOf<Short>(),
+        typeOf<Int>(),
+        typeOf<Long>(),
+        typeOf<Float>(),
+        typeOf<Double>(),
+    )
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
index 20f6a27b93..3ed36e4363 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
@@ -53,7 +53,7 @@ internal interface Aggregator<in Value, out Return> {
 
     /**
      * Special case of [aggregate] with [Iterable] that calculates the common type of the values at runtime.
-     * This is a heavy operation and should be avoided when possible.
+     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
      *
      * @param values The values to be aggregated.
      * @param valueTypes The types of the values.
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
index 2f210e12b8..755e6bf582 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
@@ -64,7 +64,7 @@ internal abstract class AggregatorBase<in Value, out Return>(
         )
 
     /** Special case of [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator.aggregate] with [Iterable] that calculates the common type of the values at runtime.
-     * This is a heavy operation and should be avoided when possible.
+     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
      *
      * @param values The values to be aggregated.
      * @param valueTypes The types of the values.
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
index 7ff4da9ee5..abbdffc575 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
@@ -175,22 +175,23 @@ internal object Aggregators {
      * Factory for a two-step aggregator that works only with numbers.
      *
      * [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] made specifically for number calculations.
+     * Mixed number types are [unified][org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers] to [primitives][org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY].
      *
      * Nulls are filtered from columns.
      *
-     * When called on multiple columns (with potentially different [Number] types),
+     * When called on multiple columns (with potentially mixed [Number] types),
      * this [Aggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator] works in two steps:
      *
-     * First, it aggregates within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn]/[Iterable] with their (given) [Number] type,
-     * and then between different columns
+     * First, it aggregates within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn]/[Iterable] with their (given) [Number] type
+     * (potentially unifying the types), and then between different columns
      * using the results of the first and the newly calculated [unified number][org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers] type of those results.
      *
      * ```
      * Iterable<Column<Number?>>
      *     -> Iterable<Iterable<Number>> // nulls filtered out
-     *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+     *     -> aggregator(Iterable<specific Number>, unified number type of common colType) // called on each iterable
      *     -> Iterable<Return> // nulls filtered out
-     *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+     *     -> aggregator(Iterable<specific Return>, unified number type of common valueType)
      *     -> Return?
      * ```
      *
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
index 4515d14bc9..2e01bb6480 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
@@ -1,8 +1,13 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
+import io.github.oshai.kotlinlogging.KotlinLogging
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
+import org.jetbrains.kotlinx.dataframe.impl.UnifiedNumberTypeOptions.Companion.PRIMITIVES_ONLY
 import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
+import org.jetbrains.kotlinx.dataframe.impl.nothingType
+import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
+import org.jetbrains.kotlinx.dataframe.impl.renderType
 import org.jetbrains.kotlinx.dataframe.impl.types
 import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
 import kotlin.reflect.KType
@@ -11,24 +16,27 @@ import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+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.
  *
- * When called on multiple columns (with potentially different [Number] types),
+ * When called on multiple columns (with potentially mixed [Number] types),
  * this [Aggregator] works in two steps:
  *
- * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type,
- * and then between different columns
+ * First, it aggregates within a [DataColumn]/[Iterable] with their (given) [Number] type
+ * (potentially unifying the types), and then between different columns
  * using the results of the first and the newly calculated [unified number][UnifyingNumbers] type of those results.
  *
  * ```
  * Iterable<Column<Number?>>
  *     -> Iterable<Iterable<Number>> // nulls filtered out
- *     -> aggregator(Iterable<Number>, colType) // called on each iterable
+ *     -> aggregator(Iterable<specific Number>, unified number type of common colType) // called on each iterable
  *     -> Iterable<Return> // nulls filtered out
- *     -> aggregator(Iterable<Return>, unified number type of common valueType)
+ *     -> aggregator(Iterable<specific Return>, unified number type of common valueType)
  *     -> Return?
  * ```
  *
@@ -44,22 +52,6 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
     aggregator: Aggregate<Number, Return>,
 ) : AggregatorBase<Number, Return>(name, getReturnTypeOrNull, aggregator) {
 
-    /**
-     * Base function of [Aggregator].
-     *
-     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
-     *
-     * Uses [aggregator] to compute the result.
-     *
-     * When the exact [type] is unknown, use [aggregateCalculatingType].
-     */
-    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 super.aggregate(values, type)
-    }
-
     /**
      * Aggregates the data in the multiple given columns and computes a single resulting value.
      *
@@ -85,10 +77,44 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
+    /**
+     * Base function of [Aggregator].
+     *
+     * Aggregates the given values, taking [type] into account, and computes a single resulting value.
+     *
+     * 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? {
+        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
+            // or calling aggregateCalculatingType with all known number types
+            typeOf<Number>() -> aggregateCalculatingType(values)
+
+            // Nothing can occur when values are empty
+            nothingType -> super.aggregate(values, type)
+
+            !in primitiveNumberTypes -> throw IllegalArgumentException(
+                "Cannot calculate $name of ${renderType(type)}, only primitive numbers are supported.",
+            )
+
+            else -> super.aggregate(values, type)
+        }
+    }
+
     /**
      * Special case of [aggregate] with [Iterable] that calculates the [unified number type][UnifyingNumbers]
      * of the values at runtime and converts all numbers to this type before aggregating.
-     * This is a heavy operation and should be avoided when possible.
+     * Without [valueTypes], this is a heavy operation and should be avoided when possible.
      *
      * @param values The numbers to be aggregated.
      * @param valueTypes The types of the numbers.
@@ -98,9 +124,24 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
      */
     @Suppress("UNCHECKED_CAST")
     override fun aggregateCalculatingType(values: Iterable<Number>, valueTypes: Set<KType>?): Return? {
-        val commonType = (valueTypes ?: values.types()).unifiedNumberType().withNullability(false)
-        return aggregate(
-            values = values.convertToUnifiedNumberType(commonType),
+        val valueTypes = valueTypes ?: values.types()
+        val commonType = valueTypes
+            .unifiedNumberType(PRIMITIVES_ONLY)
+            .withNullability(false)
+
+        if (commonType == 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) {
+            throw IllegalArgumentException(
+                "Cannot calculate $name of ${renderType(commonType)}, only primitive numbers are supported.",
+            )
+        }
+
+        return super.aggregate(
+            values = values.convertToUnifiedNumberType(commonNumberType = commonType),
             type = commonType,
         )
     }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
index 239ed236de..659180a3aa 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/aggregateBy.kt
@@ -3,12 +3,14 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.modes
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataFrameExpression
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.CandidateForRemoval
 import org.jetbrains.kotlinx.dataframe.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.namedValues
 
+@CandidateForRemoval
 internal fun <T> Grouped<T>.aggregateBy(body: DataFrameExpression<T, DataRow<T>?>): DataFrame<T> {
     require(this is GroupBy<*, T>)
     val keyColumns = keys.columnNames().toSet()
diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
index 513d7f4d19..1980e5da70 100644
--- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
+++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.statistics
 
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.columnOf
@@ -7,8 +8,8 @@ import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.rowSum
 import org.jetbrains.kotlinx.dataframe.api.sum
 import org.jetbrains.kotlinx.dataframe.api.sumOf
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.junit.Test
-import java.math.BigDecimal
 
 class SumTests {
 
@@ -58,10 +59,10 @@ class SumTests {
         df.sumOf { value3() } shouldBe expected3
         df.sum(value1) shouldBe expected1
         df.sum(value2) shouldBe expected2
-        df.sum(value3) shouldBe expected3
+        // TODO sum rework, has Number in results df.sum(value3) shouldBe expected3
         df.sum { value1 } shouldBe expected1
         df.sum { value2 } shouldBe expected2
-        df.sum { value3 } shouldBe expected3
+        // TODO sum rework, has Number in results df.sum { value3 } shouldBe expected3
     }
 
     /** [Issue #1068](https://github.com/Kotlin/dataframe/issues/1068) */
@@ -78,9 +79,17 @@ class SumTests {
             it::class shouldBe Int::class
         }
 
+        // NOTE! lossy conversion from long -> double happens
         dataFrameOf("a", "b")(1.0, 2L)[0].rowSum().let {
-            it shouldBe (3.0.toBigDecimal())
-            it::class shouldBe BigDecimal::class
+            it shouldBe 3.0
+            it::class shouldBe Double::class
+        }
+    }
+
+    @Test
+    fun `unknown number type`() {
+        shouldThrow<IllegalArgumentException> {
+            columnOf(1.toBigDecimal(), 2.toBigDecimal()).toDataFrame().sum()
         }
     }
 }
diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/types/UtilTests.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/types/UtilTests.kt
index 7027d7e194..5a95606733 100644
--- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/types/UtilTests.kt
+++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/types/UtilTests.kt
@@ -422,11 +422,11 @@ class UtilTests {
     /**
      * See [UnifyingNumbers] for more information.
      * ```
-     *            BigDecimal
+     *           (BigDecimal)
      *            /      \
-     *      BigInteger    \
+     *     (BigInteger)   \
      *        /   \        \
-     *    ULong   Long    Double
+     * <~ ULong   Long ~> Double ..
      * ..   |    /   |   /   |  \..
      *   \  |   /    |  /    |
      *     UInt     Int    Float
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
index 366a7de6d7..8da5194a7d 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/median.kt
@@ -18,7 +18,6 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.intraComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateAll
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateFor
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOf
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.aggregateOfDelegated
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.modes.of
 import org.jetbrains.kotlinx.dataframe.impl.columns.toComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.suggestIfNull

From e584e9741863a2ee034e702971298665259f5f27 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Tue, 11 Mar 2025 12:38:24 +0100
Subject: [PATCH 17/18] fixed :core statistics tests

---
 .../kotlinx/dataframe/api/statistics.kt       | 29 ++++++++++---------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
index 006b8048b2..1d4d76b637 100644
--- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
+++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
@@ -148,14 +148,15 @@ class StatisticsTests {
             "city",
             "name",
             "age",
+            "weight",
             "height",
-            "yearsToRetirement"
-        ) // TODO: why double values from weight are not in the list? are they not Comparable?
+            "yearsToRetirement",
+        )
 
         val median01 = res0["age"][0] as Int
         median01 shouldBe 22
-        //val median02 = res0["weight"][0] as Double
-        //median02 shouldBe 66.0
+        // val median02 = res0["weight"][0] as Double
+        // median02 shouldBe 66.0
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").medianFor("age")
@@ -276,14 +277,15 @@ class StatisticsTests {
             "city",
             "name",
             "age",
+            "weight",
             "height",
-            "yearsToRetirement"
-        ) // TODO: why it's working for height and doesn't work for Double column weight
+            "yearsToRetirement",
+        )
 
         val min01 = res0["age"][0] as Int
         min01 shouldBe 15
-        //val min02 = res0["weight"][0] as Double
-        //min02 shouldBe 38.85756039691633
+        // val min02 = res0["weight"][0] as Double
+        // min02 shouldBe 38.85756039691633
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").minFor("age")
@@ -342,7 +344,7 @@ class StatisticsTests {
             "age",
             "weight",
             "height",
-            "yearsToRetirement"
+            "yearsToRetirement",
         ) // TODO: why is here weight presented? looks like inconsitency
 
         val min41 = res4["age"][0] as Int
@@ -362,12 +364,12 @@ class StatisticsTests {
     fun `max on GroupBy`() {
         // scenario #0: all numerical columns
         val res0 = personsDf.groupBy("city").max()
-        res0.columnNames() shouldBe listOf("city", "name", "age", "height", "yearsToRetirement") // TODO: DOUBLE weight?
+        res0.columnNames() shouldBe listOf("city", "name", "age", "weight", "height", "yearsToRetirement")
 
         val max01 = res0["age"][0] as Int
         max01 shouldBe 35
-        //val max02 = res0["weight"][0] as Double
-        //max02 shouldBe 140.0
+        // val max02 = res0["weight"][0] as Double
+        // max02 shouldBe 140.0
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").maxFor("age")
@@ -426,7 +428,7 @@ class StatisticsTests {
             "age",
             "weight",
             "height",
-            "yearsToRetirement"
+            "yearsToRetirement",
         ) // TODO: weight is here?
 
         val max41 = res4["age"][0] as Int
@@ -442,4 +444,3 @@ class StatisticsTests {
         max51 shouldBe 35
     }
 }
-

From 502042a0c1d54487a6c89a48346ebc38f5582394 Mon Sep 17 00:00:00 2001
From: Jolan Rensen <jolan.rensen@jetbrains.com>
Date: Wed, 12 Mar 2025 15:28:57 +0100
Subject: [PATCH 18/18] fixed aggregators based on feedback. Removed
 `preservesType` property. It's unneeded as we can calculate return types at
 runtime quickly. Added overload for calculateReturnTypeOrNull for multiple
 columns. Aggregator callers now use this function instead of `preservesType`

---
 core/api/core.api                             |  2 +-
 .../kotlinx/dataframe/columns/BaseColumn.kt   |  2 ++
 .../aggregation/aggregators/Aggregator.kt     | 14 +++++++--
 .../aggregation/aggregators/AggregatorBase.kt | 11 +++++++
 .../aggregators/AggregatorOptionSwitch.kt     | 10 +++----
 .../aggregation/aggregators/Aggregators.kt    | 24 +++++----------
 .../aggregators/FlatteningAggregator.kt       | 20 +++++++++----
 .../aggregators/TwoStepAggregator.kt          | 29 ++++++++++++++-----
 .../aggregators/TwoStepNumbersAggregator.kt   | 28 ++++++++++++++++--
 .../impl/aggregation/modes/forEveryColumn.kt  | 11 +++++--
 .../impl/aggregation/modes/ofRowExpression.kt | 16 ++++++++--
 .../aggregation/modes/withinAllColumns.kt     | 25 ++++++++++++++--
 .../jetbrains/kotlinx/dataframe/math/mean.kt  |  2 +-
 .../jetbrains/kotlinx/dataframe/math/std.kt   |  2 +-
 .../jetbrains/kotlinx/dataframe/math/sum.kt   |  2 +-
 .../kotlinx/dataframe/api/statistics.kt       | 29 ++++++++++---------
 .../kotlinx/dataframe/columns/BaseColumn.kt   |  2 ++
 .../aggregation/aggregators/Aggregator.kt     | 14 +++++++--
 .../aggregation/aggregators/AggregatorBase.kt | 11 +++++++
 .../aggregators/AggregatorOptionSwitch.kt     | 10 +++----
 .../aggregation/aggregators/Aggregators.kt    |  4 ---
 .../aggregators/FlatteningAggregator.kt       | 20 +++++++++----
 .../aggregators/TwoStepAggregator.kt          | 29 ++++++++++++++-----
 .../aggregators/TwoStepNumbersAggregator.kt   | 28 ++++++++++++++++--
 .../impl/aggregation/modes/forEveryColumn.kt  | 11 +++++--
 .../impl/aggregation/modes/ofRowExpression.kt | 16 ++++++++--
 .../aggregation/modes/withinAllColumns.kt     | 25 ++++++++++++++--
 .../jetbrains/kotlinx/dataframe/math/mean.kt  |  2 +-
 .../jetbrains/kotlinx/dataframe/math/std.kt   |  2 +-
 .../jetbrains/kotlinx/dataframe/math/sum.kt   |  2 +-
 30 files changed, 297 insertions(+), 106 deletions(-)

diff --git a/core/api/core.api b/core/api/core.api
index c146483135..2acda50a99 100644
--- a/core/api/core.api
+++ b/core/api/core.api
@@ -5304,9 +5304,9 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/impl/aggregation
 	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 static synthetic fun aggregateCalculatingType$default (Lorg/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator;Ljava/lang/Iterable;Ljava/util/Set;ILjava/lang/Object;)Ljava/lang/Object;
+	public abstract fun calculateReturnTypeOrNull (Ljava/util/Set;Z)Lkotlin/reflect/KType;
 	public abstract fun calculateReturnTypeOrNull (Lkotlin/reflect/KType;Z)Lkotlin/reflect/KType;
 	public abstract fun getName ()Ljava/lang/String;
-	public abstract fun getPreservesType ()Z
 }
 
 public final class org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorKt {
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
index cfed2a1de4..9117e01bf8 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
@@ -100,3 +100,5 @@ public interface BaseColumn<out T> : ColumnReference<T> {
 internal val <T> BaseColumn<T>.values: Iterable<T> get() = values()
 
 internal val AnyBaseCol.size: Int get() = size()
+
+internal val AnyBaseCol.isEmpty: Boolean get() = size() == 0
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
index 3ed36e4363..0050145715 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregator.kt
@@ -24,9 +24,6 @@ internal interface Aggregator<in Value, out Return> {
     /** The name of this aggregator. */
     val name: String
 
-    /** If `true`, [Value][Value]`  ==  ` [Return][Return]. */
-    val preservesType: Boolean
-
     /**
      * Base function of [Aggregator].
      *
@@ -72,6 +69,17 @@ internal interface Aggregator<in Value, out Return> {
      * @return The return type of [aggregate] as [KType].
      */
     fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
+
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
 }
 
 @PublishedApi
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
index 755e6bf582..99c7b33c61 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorBase.kt
@@ -94,4 +94,15 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * Must be overridden to use.
      */
     abstract override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
+
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    abstract override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index 7a2b2d4496..a21b06c401 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -30,9 +30,8 @@ internal class AggregatorOptionSwitch1<in Param1, out AggregatorType : Aggregato
      */
     class Factory<in Param1, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
-    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
-            AggregatorOptionSwitch1(name, getAggregator)
-        })
+    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by
+        Provider({ name -> AggregatorOptionSwitch1(name, getAggregator) })
 }
 
 /**
@@ -66,7 +65,6 @@ internal class AggregatorOptionSwitch2<in Param1, in Param2, out AggregatorType
      */
     class Factory<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
-    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
-            AggregatorOptionSwitch2(name, getAggregator)
-        })
+    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by
+        Provider({ name -> AggregatorOptionSwitch2(name, getAggregator) })
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
index abbdffc575..ac05bc6e79 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt
@@ -33,24 +33,21 @@ internal object Aggregators {
      *     -> Return?
      * ```
      *
-     * It can also be used as a "simple" aggregator by providing the same function for both steps,
-     * requires [preservesType][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.preservesType] be set to `true`.
+     * It can also be used as a "simple" aggregator by providing the same function for both steps.
      *
      * See [FlatteningAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator] for different behavior for multiple columns.
      *
      * @param name The name of this aggregator.
-     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.calculateReturnTypeOrNull] function.
      * @param stepOneAggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     private fun <Type> twoStepPreservingType(aggregator: Aggregate<Type, Type>) =
         TwoStepAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             stepOneAggregator = aggregator,
             stepTwoAggregator = aggregator,
-            preservesType = true,
         )
 
     /**
@@ -74,17 +71,15 @@ internal object Aggregators {
      *     -> Return?
      * ```
      *
-     * It can also be used as a "simple" aggregator by providing the same function for both steps,
-     * requires [preservesType][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.preservesType] be set to `true`.
+     * It can also be used as a "simple" aggregator by providing the same function for both steps.
      *
      * See [FlatteningAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator] for different behavior for multiple columns.
      *
      * @param name The name of this aggregator.
-     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.calculateReturnTypeOrNull] function.
      * @param stepOneAggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     private fun <Value, Return> twoStepChangingType(
         getReturnTypeOrNull: CalculateReturnTypeOrNull,
@@ -94,7 +89,6 @@ internal object Aggregators {
         getReturnTypeOrNull = getReturnTypeOrNull,
         stepOneAggregator = stepOneAggregator,
         stepTwoAggregator = stepTwoAggregator,
-        preservesType = false,
     )
 
     /**
@@ -121,16 +115,14 @@ internal object Aggregators {
      * See [TwoStepAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator] for different behavior for multiple columns.
      *
      * @param name The name of this aggregator.
-     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator.calculateReturnTypeOrNull] function.
      * @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].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     private fun <Type> flatteningPreservingTypes(aggregate: Aggregate<Type, Type>) =
         FlatteningAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             aggregator = aggregate,
-            preservesType = true,
         )
 
     /**
@@ -157,10 +149,9 @@ internal object Aggregators {
      * See [TwoStepAggregator][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepAggregator] for different behavior for multiple columns.
      *
      * @param name The name of this aggregator.
-     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.FlatteningAggregator.calculateReturnTypeOrNull] function.
      * @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].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     private fun <Value, Return> flatteningChangingTypes(
         getReturnTypeOrNull: CalculateReturnTypeOrNull,
@@ -168,7 +159,6 @@ internal object Aggregators {
     ) = FlatteningAggregator.Factory(
         getReturnTypeOrNull = getReturnTypeOrNull,
         aggregator = aggregate,
-        preservesType = false,
     )
 
     /**
@@ -196,7 +186,7 @@ internal object Aggregators {
      * ```
      *
      * @param name The name of this aggregator.
-     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.AggregatorBase.calculateReturnTypeOrNull] function.
+     * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepNumbersAggregator.calculateReturnTypeOrNull] function.
      * @param aggregator Functional argument for the [aggregate][org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.TwoStepNumbersAggregator.aggregate] function, used within a [DataColumn][org.jetbrains.kotlinx.dataframe.DataColumn] or [Iterable].
      *   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].
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
index b259339a69..270777e7a2 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/FlatteningAggregator.kt
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
 /**
@@ -29,13 +30,11 @@ import kotlin.reflect.full.withNullability
  * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
  * @param aggregator Functional argument for the [aggregate] function.
  *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
- * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
 internal class FlatteningAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     aggregator: Aggregate<Value, Return>,
-    override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, aggregator) {
 
     /**
@@ -49,23 +48,34 @@ internal class FlatteningAggregator<in Value, out Return>(
         return aggregate(allValues.asIterable(), commonType)
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val commonType = colTypes.commonType().withNullability(false)
+        return calculateReturnTypeOrNull(commonType, colsEmpty)
+    }
+
     /**
      * Creates [FlatteningAggregator].
      *
      * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
      * @param aggregator Functional argument for the [aggregate] function.
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<in Value, out Return>(
         private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
         private val aggregator: Aggregate<Value, Return>,
-        private val preservesType: Boolean,
     ) : AggregatorProvider<FlatteningAggregator<Value, Return>> by AggregatorProvider({ name ->
             FlatteningAggregator(
                 name = name,
                 getReturnTypeOrNull = getReturnTypeOrNull,
                 aggregator = aggregator,
-                preservesType = preservesType,
             )
         })
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
index b095708a2a..11738fbf5e 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepAggregator.kt
@@ -1,7 +1,9 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.KType
 import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 
@@ -24,8 +26,7 @@ import kotlin.reflect.full.withNullability
  *     -> Return?
  * ```
  *
- * It can also be used as a "simple" aggregator by providing the same function for both steps,
- * requires [preservesType] be set to `true`.
+ * It can also be used as a "simple" aggregator by providing the same function for both steps.
  *
  * See [FlatteningAggregator] for different behavior for multiple columns.
  *
@@ -34,14 +35,12 @@ import kotlin.reflect.full.withNullability
  * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
  * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
  *   It is run on the results of [stepOneAggregator].
- * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
 internal class TwoStepAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     stepOneAggregator: Aggregate<Value, Return>,
     private val stepTwoAggregator: Aggregate<Return, Return>,
-    override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
     /**
@@ -57,7 +56,7 @@ internal class TwoStepAggregator<in Value, out Return>(
             val value = aggregate(col) ?: return@mapNotNull null
             val type = calculateReturnTypeOrNull(
                 type = col.type().withNullability(false),
-                emptyInput = col.size() == 0,
+                emptyInput = col.isEmpty,
             ) ?: value::class.starProjectedType // heavy fallback type calculation
 
             value to type
@@ -66,6 +65,23 @@ internal class TwoStepAggregator<in Value, out Return>(
         return stepTwoAggregator(values, commonType)
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val typesAfterStepOne = colTypes.map { type ->
+            calculateReturnTypeOrNull(type = type.withNullability(false), emptyInput = colsEmpty)
+        }
+        if (typesAfterStepOne.any { it == null }) return null
+        return typesAfterStepOne.commonType()
+    }
+
     /**
      * Creates [TwoStepAggregator].
      *
@@ -73,20 +89,17 @@ internal class TwoStepAggregator<in Value, out Return>(
      * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][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 preservesType: Boolean,
     ) : AggregatorProvider<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
             TwoStepAggregator(
                 name = name,
                 getReturnTypeOrNull = getReturnTypeOrNull,
                 stepOneAggregator = stepOneAggregator,
                 stepTwoAggregator = stepTwoAggregator,
-                preservesType = preservesType,
             )
         })
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
index 2e01bb6480..bb229720d3 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/TwoStepNumbersAggregator.kt
@@ -2,8 +2,10 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import io.github.oshai.kotlinlogging.KotlinLogging
 import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 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.nothingType
 import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
@@ -65,7 +67,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
             val value = aggregate(col) ?: return@mapNotNull null
             val type = calculateReturnTypeOrNull(
                 type = col.type().withNullability(false),
-                emptyInput = col.size() == 0,
+                emptyInput = col.isEmpty,
             ) ?: value::class.starProjectedType // heavy fallback type calculation
 
             value to type
@@ -77,6 +79,28 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    @Suppress("UNCHECKED_CAST")
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val typesAfterStepOne = colTypes.map { type ->
+            calculateReturnTypeOrNull(type = type.withNullability(false), emptyInput = colsEmpty)
+        }
+        if (typesAfterStepOne.anyNull()) return null
+        val commonType = (typesAfterStepOne as List<KType>)
+            .toSet()
+            .unifiedNumberType(PRIMITIVES_ONLY)
+            .withNullability(false)
+        return commonType
+    }
+
     /**
      * Base function of [Aggregator].
      *
@@ -146,8 +170,6 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
-    override val preservesType = false
-
     /**
      * Creates [TwoStepNumbersAggregator].
      *
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
index fbb932e03c..6ec6459ff0 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
@@ -48,7 +49,13 @@ internal fun <T, C, R> AggregateInternalDsl<T>.aggregateFor(
     cols.forEach { col ->
         val path = getPath(col, isSingle)
         val value = aggregator.aggregate(col.data)
-        val inferType = !aggregator.preservesType
-        yield(path, value, col.type, col.default, inferType)
+        val returnType = aggregator.calculateReturnTypeOrNull(col.data.type, col.data.isEmpty)
+        yield(
+            path = path,
+            value = value,
+            type = returnType,
+            default = col.default,
+            guessType = returnType == null,
+        )
     }
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
index 299587a3b9..80bdc5bc33 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateBody
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
+import org.jetbrains.kotlinx.dataframe.api.isEmpty
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.rows
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
@@ -72,11 +73,20 @@ internal inline fun <T, reified C, reified R> Grouped<T>.aggregateOf(
     aggregator: Aggregator<C, R>,
 ): DataFrame<T> {
     val path = pathOf(resultName ?: aggregator.name)
-    val type = typeOf<R>()
+    val expressionResultType = typeOf<C>()
     return aggregateInternal {
         val value = aggregator.aggregateOf(df, expression)
-        val inferType = !aggregator.preservesType
-        yield(path, value, type, null, inferType)
+        val returnType = aggregator.calculateReturnTypeOrNull(
+            type = expressionResultType,
+            emptyInput = df.isEmpty(),
+        )
+        yield(
+            path = path,
+            value = value,
+            type = returnType,
+            default = null,
+            guessType = returnType == null,
+        )
     }
 }
 
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
index c6edf1400e..c46481bf65 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
 import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.get
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
@@ -52,8 +53,28 @@ internal fun <T, C, R> PivotGroupBy<T>.aggregateAll(
     aggregate {
         val cols = get(columns)
         if (cols.size == 1) {
-            internal().yield(emptyPath(), aggregator.aggregate(cols[0]))
+            val returnType = aggregator.calculateReturnTypeOrNull(
+                type = cols[0].type(),
+                emptyInput = cols[0].isEmpty,
+            )
+            internal().yield(
+                path = emptyPath(),
+                value = aggregator.aggregate(cols[0]),
+                type = returnType,
+                default = null,
+                guessType = returnType == null,
+            )
         } else {
-            internal().yield(emptyPath(), aggregator.aggregate(cols))
+            val returnType = aggregator.calculateReturnTypeOrNull(
+                colTypes = cols.map { it.type() }.toSet(),
+                colsEmpty = cols.any { it.isEmpty },
+            )
+            internal().yield(
+                path = emptyPath(),
+                value = aggregator.aggregate(cols),
+                type = returnType,
+                default = null,
+                guessType = returnType == null,
+            )
         }
     }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
index fc5fb70b70..d2d5bb4004 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -45,7 +45,7 @@ internal fun <T : Number> Sequence<T?>.mean(type: KType, skipNA: Boolean = skipN
     }
 }
 
-// T: Number? -> Double
+/** T: Number? -> Double */
 internal val meanTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
     typeOf<Double>()
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
index 148c9ece23..a91184985c 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
@@ -37,7 +37,7 @@ internal fun <T : Number> Iterable<T?>.std(
     }
 }
 
-// T: Number? -> Double
+/** T: Number? -> Double */
 internal val stdTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
     typeOf<Double>()
 }
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
index 1b03221988..07de30db44 100644
--- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
+++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
@@ -97,7 +97,7 @@ internal fun <T : Number> Iterable<T?>.sum(type: KType): T =
         else -> throw IllegalArgumentException("sum is not supported for $type")
     }
 
-// T: Number? -> T
+/** T: Number? -> T */
 internal val sumTypeConversion: CalculateReturnTypeOrNull = { type, _ ->
     type.withNullability(false)
 }
diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
index 006b8048b2..1d4d76b637 100644
--- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
+++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt
@@ -148,14 +148,15 @@ class StatisticsTests {
             "city",
             "name",
             "age",
+            "weight",
             "height",
-            "yearsToRetirement"
-        ) // TODO: why double values from weight are not in the list? are they not Comparable?
+            "yearsToRetirement",
+        )
 
         val median01 = res0["age"][0] as Int
         median01 shouldBe 22
-        //val median02 = res0["weight"][0] as Double
-        //median02 shouldBe 66.0
+        // val median02 = res0["weight"][0] as Double
+        // median02 shouldBe 66.0
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").medianFor("age")
@@ -276,14 +277,15 @@ class StatisticsTests {
             "city",
             "name",
             "age",
+            "weight",
             "height",
-            "yearsToRetirement"
-        ) // TODO: why it's working for height and doesn't work for Double column weight
+            "yearsToRetirement",
+        )
 
         val min01 = res0["age"][0] as Int
         min01 shouldBe 15
-        //val min02 = res0["weight"][0] as Double
-        //min02 shouldBe 38.85756039691633
+        // val min02 = res0["weight"][0] as Double
+        // min02 shouldBe 38.85756039691633
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").minFor("age")
@@ -342,7 +344,7 @@ class StatisticsTests {
             "age",
             "weight",
             "height",
-            "yearsToRetirement"
+            "yearsToRetirement",
         ) // TODO: why is here weight presented? looks like inconsitency
 
         val min41 = res4["age"][0] as Int
@@ -362,12 +364,12 @@ class StatisticsTests {
     fun `max on GroupBy`() {
         // scenario #0: all numerical columns
         val res0 = personsDf.groupBy("city").max()
-        res0.columnNames() shouldBe listOf("city", "name", "age", "height", "yearsToRetirement") // TODO: DOUBLE weight?
+        res0.columnNames() shouldBe listOf("city", "name", "age", "weight", "height", "yearsToRetirement")
 
         val max01 = res0["age"][0] as Int
         max01 shouldBe 35
-        //val max02 = res0["weight"][0] as Double
-        //max02 shouldBe 140.0
+        // val max02 = res0["weight"][0] as Double
+        // max02 shouldBe 140.0
 
         // scenario #1: particular column
         val res1 = personsDf.groupBy("city").maxFor("age")
@@ -426,7 +428,7 @@ class StatisticsTests {
             "age",
             "weight",
             "height",
-            "yearsToRetirement"
+            "yearsToRetirement",
         ) // TODO: weight is here?
 
         val max41 = res4["age"][0] as Int
@@ -442,4 +444,3 @@ class StatisticsTests {
         max51 shouldBe 35
     }
 }
-
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
index cfed2a1de4..9117e01bf8 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/BaseColumn.kt
@@ -100,3 +100,5 @@ public interface BaseColumn<out T> : ColumnReference<T> {
 internal val <T> BaseColumn<T>.values: Iterable<T> get() = values()
 
 internal val AnyBaseCol.size: Int get() = size()
+
+internal val AnyBaseCol.isEmpty: Boolean get() = size() == 0
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
index 3ed36e4363..0050145715 100644
--- 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
@@ -24,9 +24,6 @@ internal interface Aggregator<in Value, out Return> {
     /** The name of this aggregator. */
     val name: String
 
-    /** If `true`, [Value][Value]`  ==  ` [Return][Return]. */
-    val preservesType: Boolean
-
     /**
      * Base function of [Aggregator].
      *
@@ -72,6 +69,17 @@ internal interface Aggregator<in Value, out Return> {
      * @return The return type of [aggregate] as [KType].
      */
     fun calculateReturnTypeOrNull(type: KType, emptyInput: Boolean): KType?
+
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
 }
 
 @PublishedApi
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
index a2dac82aed..906b40dc83 100644
--- 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
@@ -87,4 +87,15 @@ internal abstract class AggregatorBase<in Value, out Return>(
      * Must be overridden to use.
      */
     abstract override fun aggregate(columns: Iterable<DataColumn<Value?>>): Return?
+
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    abstract override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType?
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
index 7a2b2d4496..a21b06c401 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt
@@ -30,9 +30,8 @@ internal class AggregatorOptionSwitch1<in Param1, out AggregatorType : Aggregato
      */
     class Factory<in Param1, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (param1: Param1) -> AggregatorProvider<AggregatorType>,
-    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by Provider({ name ->
-            AggregatorOptionSwitch1(name, getAggregator)
-        })
+    ) : Provider<AggregatorOptionSwitch1<Param1, AggregatorType>> by
+        Provider({ name -> AggregatorOptionSwitch1(name, getAggregator) })
 }
 
 /**
@@ -66,7 +65,6 @@ internal class AggregatorOptionSwitch2<in Param1, in Param2, out AggregatorType
      */
     class Factory<in Param1, in Param2, out AggregatorType : Aggregator<*, *>>(
         val getAggregator: (param1: Param1, param2: Param2) -> AggregatorProvider<AggregatorType>,
-    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by Provider({ name ->
-            AggregatorOptionSwitch2(name, getAggregator)
-        })
+    ) : Provider<AggregatorOptionSwitch2<Param1, Param2, AggregatorType>> by
+        Provider({ name -> AggregatorOptionSwitch2(name, getAggregator) })
 }
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
index 8c677a0990..5017288c2e 100644
--- 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
@@ -22,7 +22,6 @@ internal object Aggregators {
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             stepOneAggregator = aggregator,
             stepTwoAggregator = aggregator,
-            preservesType = true,
         )
 
     /**
@@ -38,7 +37,6 @@ internal object Aggregators {
         getReturnTypeOrNull = getReturnTypeOrNull,
         stepOneAggregator = stepOneAggregator,
         stepTwoAggregator = stepTwoAggregator,
-        preservesType = false,
     )
 
     /**
@@ -50,7 +48,6 @@ internal object Aggregators {
         FlatteningAggregator.Factory(
             getReturnTypeOrNull = preserveReturnTypeNullIfEmpty,
             aggregator = aggregate,
-            preservesType = true,
         )
 
     /**
@@ -64,7 +61,6 @@ internal object Aggregators {
     ) = FlatteningAggregator.Factory(
         getReturnTypeOrNull = getReturnTypeOrNull,
         aggregator = aggregate,
-        preservesType = false,
     )
 
     /**
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
index b259339a69..270777e7a2 100644
--- 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
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.KType
 import kotlin.reflect.full.withNullability
 
 /**
@@ -29,13 +30,11 @@ import kotlin.reflect.full.withNullability
  * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
  * @param aggregator Functional argument for the [aggregate] function.
  *   Note that it must be able to handle `null` values for the [Iterable] overload of [aggregate].
- * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
 internal class FlatteningAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     aggregator: Aggregate<Value, Return>,
-    override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, aggregator) {
 
     /**
@@ -49,23 +48,34 @@ internal class FlatteningAggregator<in Value, out Return>(
         return aggregate(allValues.asIterable(), commonType)
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val commonType = colTypes.commonType().withNullability(false)
+        return calculateReturnTypeOrNull(commonType, colsEmpty)
+    }
+
     /**
      * Creates [FlatteningAggregator].
      *
      * @param getReturnTypeOrNull Functional argument for the [calculateReturnTypeOrNull] function.
      * @param aggregator Functional argument for the [aggregate] function.
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
      */
     class Factory<in Value, out Return>(
         private val getReturnTypeOrNull: CalculateReturnTypeOrNull,
         private val aggregator: Aggregate<Value, Return>,
-        private val preservesType: Boolean,
     ) : AggregatorProvider<FlatteningAggregator<Value, Return>> by AggregatorProvider({ name ->
             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
index b095708a2a..11738fbf5e 100644
--- 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,7 +1,9 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.impl.commonType
+import kotlin.reflect.KType
 import kotlin.reflect.full.starProjectedType
 import kotlin.reflect.full.withNullability
 
@@ -24,8 +26,7 @@ import kotlin.reflect.full.withNullability
  *     -> Return?
  * ```
  *
- * It can also be used as a "simple" aggregator by providing the same function for both steps,
- * requires [preservesType] be set to `true`.
+ * It can also be used as a "simple" aggregator by providing the same function for both steps.
  *
  * See [FlatteningAggregator] for different behavior for multiple columns.
  *
@@ -34,14 +35,12 @@ import kotlin.reflect.full.withNullability
  * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
  * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
  *   It is run on the results of [stepOneAggregator].
- * @param preservesType If `true`, [Value][Value]`  ==  `[Return][Return].
  */
 internal class TwoStepAggregator<in Value, out Return>(
     name: String,
     getReturnTypeOrNull: CalculateReturnTypeOrNull,
     stepOneAggregator: Aggregate<Value, Return>,
     private val stepTwoAggregator: Aggregate<Return, Return>,
-    override val preservesType: Boolean,
 ) : AggregatorBase<Value, Return>(name, getReturnTypeOrNull, stepOneAggregator) {
 
     /**
@@ -57,7 +56,7 @@ internal class TwoStepAggregator<in Value, out Return>(
             val value = aggregate(col) ?: return@mapNotNull null
             val type = calculateReturnTypeOrNull(
                 type = col.type().withNullability(false),
-                emptyInput = col.size() == 0,
+                emptyInput = col.isEmpty,
             ) ?: value::class.starProjectedType // heavy fallback type calculation
 
             value to type
@@ -66,6 +65,23 @@ internal class TwoStepAggregator<in Value, out Return>(
         return stepTwoAggregator(values, commonType)
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val typesAfterStepOne = colTypes.map { type ->
+            calculateReturnTypeOrNull(type = type.withNullability(false), emptyInput = colsEmpty)
+        }
+        if (typesAfterStepOne.any { it == null }) return null
+        return typesAfterStepOne.commonType()
+    }
+
     /**
      * Creates [TwoStepAggregator].
      *
@@ -73,20 +89,17 @@ internal class TwoStepAggregator<in Value, out Return>(
      * @param stepOneAggregator Functional argument for the [aggregate] function, used within a [DataColumn] or [Iterable].
      * @param stepTwoAggregator Functional argument for the aggregation function used between different columns.
      *   It is run on the results of [stepOneAggregator].
-     * @param preservesType If `true`, [Value][Value]`  ==  `[Return][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 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
index 2e01bb6480..bb229720d3 100644
--- 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
@@ -2,8 +2,10 @@ package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
 
 import io.github.oshai.kotlinlogging.KotlinLogging
 import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 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.nothingType
 import org.jetbrains.kotlinx.dataframe.impl.primitiveNumberTypes
@@ -65,7 +67,7 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
             val value = aggregate(col) ?: return@mapNotNull null
             val type = calculateReturnTypeOrNull(
                 type = col.type().withNullability(false),
-                emptyInput = col.size() == 0,
+                emptyInput = col.isEmpty,
             ) ?: value::class.starProjectedType // heavy fallback type calculation
 
             value to type
@@ -77,6 +79,28 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
+    /**
+     * Function that can give the return type of [aggregate] with columns as [KType],
+     * given the multiple types of the input.
+     * This allows aggregators to avoid runtime type calculations.
+     *
+     * @param colTypes The types of the input columns.
+     * @param colsEmpty If `true`, all the input columns are considered empty. This often affects the return type.
+     * @return The return type of [aggregate] as [KType].
+     */
+    @Suppress("UNCHECKED_CAST")
+    override fun calculateReturnTypeOrNull(colTypes: Set<KType>, colsEmpty: Boolean): KType? {
+        val typesAfterStepOne = colTypes.map { type ->
+            calculateReturnTypeOrNull(type = type.withNullability(false), emptyInput = colsEmpty)
+        }
+        if (typesAfterStepOne.anyNull()) return null
+        val commonType = (typesAfterStepOne as List<KType>)
+            .toSet()
+            .unifiedNumberType(PRIMITIVES_ONLY)
+            .withNullability(false)
+        return commonType
+    }
+
     /**
      * Base function of [Aggregator].
      *
@@ -146,8 +170,6 @@ internal class TwoStepNumbersAggregator<out Return : Number>(
         )
     }
 
-    override val preservesType = false
-
     /**
      * Creates [TwoStepNumbersAggregator].
      *
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
index fbb932e03c..6ec6459ff0 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/forEveryColumn.kt
@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.aggregation.ColumnsForAggregateSelector
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.cast
@@ -48,7 +49,13 @@ internal fun <T, C, R> AggregateInternalDsl<T>.aggregateFor(
     cols.forEach { col ->
         val path = getPath(col, isSingle)
         val value = aggregator.aggregate(col.data)
-        val inferType = !aggregator.preservesType
-        yield(path, value, col.type, col.default, inferType)
+        val returnType = aggregator.calculateReturnTypeOrNull(col.data.type, col.data.isEmpty)
+        yield(
+            path = path,
+            value = value,
+            type = returnType,
+            default = col.default,
+            guessType = returnType == null,
+        )
     }
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
index 299587a3b9..80bdc5bc33 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/ofRowExpression.kt
@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateBody
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
+import org.jetbrains.kotlinx.dataframe.api.isEmpty
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.rows
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
@@ -72,11 +73,20 @@ internal inline fun <T, reified C, reified R> Grouped<T>.aggregateOf(
     aggregator: Aggregator<C, R>,
 ): DataFrame<T> {
     val path = pathOf(resultName ?: aggregator.name)
-    val type = typeOf<R>()
+    val expressionResultType = typeOf<C>()
     return aggregateInternal {
         val value = aggregator.aggregateOf(df, expression)
-        val inferType = !aggregator.preservesType
-        yield(path, value, type, null, inferType)
+        val returnType = aggregator.calculateReturnTypeOrNull(
+            type = expressionResultType,
+            emptyInput = df.isEmpty(),
+        )
+        yield(
+            path = path,
+            value = value,
+            type = returnType,
+            default = null,
+            guessType = returnType == null,
+        )
     }
 }
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
index c6edf1400e..c46481bf65 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/modes/withinAllColumns.kt
@@ -5,6 +5,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.Grouped
 import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
 import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.columns.isEmpty
 import org.jetbrains.kotlinx.dataframe.get
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregateInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
@@ -52,8 +53,28 @@ internal fun <T, C, R> PivotGroupBy<T>.aggregateAll(
     aggregate {
         val cols = get(columns)
         if (cols.size == 1) {
-            internal().yield(emptyPath(), aggregator.aggregate(cols[0]))
+            val returnType = aggregator.calculateReturnTypeOrNull(
+                type = cols[0].type(),
+                emptyInput = cols[0].isEmpty,
+            )
+            internal().yield(
+                path = emptyPath(),
+                value = aggregator.aggregate(cols[0]),
+                type = returnType,
+                default = null,
+                guessType = returnType == null,
+            )
         } else {
-            internal().yield(emptyPath(), aggregator.aggregate(cols))
+            val returnType = aggregator.calculateReturnTypeOrNull(
+                colTypes = cols.map { it.type() }.toSet(),
+                colsEmpty = cols.any { it.isEmpty },
+            )
+            internal().yield(
+                path = emptyPath(),
+                value = aggregator.aggregate(cols),
+                type = returnType,
+                default = null,
+                guessType = returnType == null,
+            )
         }
     }
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
index fc5fb70b70..d2d5bb4004 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/mean.kt
@@ -45,7 +45,7 @@ internal fun <T : Number> Sequence<T?>.mean(type: KType, skipNA: Boolean = skipN
     }
 }
 
-// T: Number? -> Double
+/** T: Number? -> Double */
 internal val meanTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
     typeOf<Double>()
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
index 148c9ece23..a91184985c 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/std.kt
@@ -37,7 +37,7 @@ internal fun <T : Number> Iterable<T?>.std(
     }
 }
 
-// T: Number? -> Double
+/** T: Number? -> Double */
 internal val stdTypeConversion: CalculateReturnTypeOrNull = { _, _ ->
     typeOf<Double>()
 }
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
index 1b03221988..07de30db44 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt
@@ -97,7 +97,7 @@ internal fun <T : Number> Iterable<T?>.sum(type: KType): T =
         else -> throw IllegalArgumentException("sum is not supported for $type")
     }
 
-// T: Number? -> T
+/** T: Number? -> T */
 internal val sumTypeConversion: CalculateReturnTypeOrNull = { type, _ ->
     type.withNullability(false)
 }