update from mean branch, added sumOf

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt

@@ -1,3 +1,5 @@
+@file:OptIn(ExperimentalTypeInference::class)
+
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyRow
@@ -13,19 +15,16 @@ import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.toColumnsSetOf
-import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregator
 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.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.aggregateOfRow
-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.primitiveNumberTypes
 import org.jetbrains.kotlinx.dataframe.impl.zero
-import org.jetbrains.kotlinx.dataframe.math.sumOf
+import kotlin.experimental.ExperimentalTypeInference
 import kotlin.reflect.KProperty
 import kotlin.reflect.typeOf
 
@@ -52,8 +51,37 @@ public fun DataColumn<Double?>.sum(): Double = Aggregators.sum.aggregate(this) a
 @JvmName("sumNumber")
 public fun DataColumn<Number?>.sum(): Number = Aggregators.sum.aggregate(this)
 
-public inline fun <T, reified R : Number> DataColumn<T>.sumOf(noinline expression: (T) -> R): R? =
-    (Aggregators.sum as Aggregator<*, *>).cast<R>().aggregateOf(this, expression)
+@JvmName("sumOfInt")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Int?): Int = Aggregators.sum.aggregateOf(this, expression) as Int
+
+@JvmName("sumOfShort")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Short?): Int =
+    Aggregators.sum.aggregateOf(this, expression) as Int
+
+@JvmName("sumOfByte")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Byte?): Int = Aggregators.sum.aggregateOf(this, expression) as Int
+
+@JvmName("sumOfLong")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Long?): Long =
+    Aggregators.sum.aggregateOf(this, expression) as Long
+
+@JvmName("sumOfFloat")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Float?): Float =
+    Aggregators.sum.aggregateOf(this, expression) as Float
+
+@JvmName("sumOfDouble")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Double?): Double =
+    Aggregators.sum.aggregateOf(this, expression) as Double
+
+@JvmName("sumOfNumber")
+@OverloadResolutionByLambdaReturnType
+public fun <T> DataColumn<T>.sumOf(expression: (T) -> Number?): Number = Aggregators.sum.aggregateOf(this, expression)
 
 // endregion
 

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

@@ -41,7 +41,7 @@ internal fun <T, C, R> Aggregator<*, R>.aggregateFor(
     }
 
 internal fun <T, C, R> AggregateInternalDsl<T>.aggregateFor(
-    columns: ColumnsForAggregateSelector<T, C>,
+    columns: ColumnsForAggregateSelector<T, C?>,
     aggregator: Aggregator<C, R>,
 ) {
     val cols = df.getAggregateColumns(columns)

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

@@ -15,26 +15,26 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.internal
 import org.jetbrains.kotlinx.dataframe.impl.emptyPath
 
 @PublishedApi
-internal fun <T, C, R> Aggregator<*, R>.aggregateAll(data: DataFrame<T>, columns: ColumnsSelector<T, C>): R =
+internal fun <T, C, R> Aggregator<*, R>.aggregateAll(data: DataFrame<T>, columns: ColumnsSelector<T, C?>): R =
     data.aggregateAll(cast2(), columns)
 
 internal fun <T, R, C> Aggregator<*, R>.aggregateAll(
     data: Grouped<T>,
     name: String?,
-    columns: ColumnsSelector<T, C>,
+    columns: ColumnsSelector<T, C?>,
 ): DataFrame<T> = data.aggregateAll(cast(), columns, name)
 
 internal fun <T, R, C> Aggregator<*, R>.aggregateAll(
     data: PivotGroupBy<T>,
-    columns: ColumnsSelector<T, C>,
+    columns: ColumnsSelector<T, C?>,
 ): DataFrame<T> = data.aggregateAll(cast(), columns)
 
 internal fun <T, C, R> DataFrame<T>.aggregateAll(aggregator: Aggregator<C, R>, columns: ColumnsSelector<T, C?>): R =
     aggregator.aggregate(get(columns))
 
 internal fun <T, C, R> Grouped<T>.aggregateAll(
     aggregator: Aggregator<C, R>,
-    columns: ColumnsSelector<T, C>,
+    columns: ColumnsSelector<T, C?>,
     name: String?,
 ): DataFrame<T> =
     aggregateInternal {
@@ -48,7 +48,7 @@ internal fun <T, C, R> Grouped<T>.aggregateAll(
 
 internal fun <T, C, R> PivotGroupBy<T>.aggregateAll(
     aggregator: Aggregator<C, R>,
-    columns: ColumnsSelector<T, C>,
+    columns: ColumnsSelector<T, C?>,
 ): DataFrame<T> =
     aggregate {
         val cols = get(columns)

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/sum.kt

@@ -8,40 +8,6 @@ import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 import kotlin.sequences.filterNotNull
 
-@PublishedApi
-internal fun <T, R : Number> Iterable<T>.sumOf(type: KType, selector: (T) -> R?): Number =
-    asSequence().sumOf(type, selector)
-
-@Suppress("UNCHECKED_CAST")
-@PublishedApi
-internal fun <T, R : Number> Sequence<T>.sumOf(type: KType, selector: (T) -> R?): Number {
-    if (type.isMarkedNullable) {
-        return filterNotNull().sumOf(type.withNullability(false), selector)
-    }
-    return when (type.withNullability(false)) {
-        typeOf<Double>() -> sumOf(selector as (T) -> Double)
-
-        typeOf<Float>() -> map(selector as (T) -> Float).sum()
-
-        typeOf<Int>() -> sumOf(selector as (T) -> Int)
-
-        // Note: returns Int
-        typeOf<Short>() -> map(selector as (T) -> Short).sum()
-
-        // Note: returns Int
-        typeOf<Byte>() -> map(selector as (T) -> Byte).sum()
-
-        typeOf<Long>() -> sumOf(selector as (T) -> Long)
-
-        nothingType -> 0.0
-
-        typeOf<Number>() ->
-            error("Encountered non-specific Number type in sumOf function. This should not occur.")
-
-        else -> throw IllegalArgumentException("sumOf is not supported for $type")
-    }
-}
-
 internal fun Iterable<Number?>.sum(type: KType): Number = asSequence().sum(type)
 
 @Suppress("UNCHECKED_CAST")