starting min/max

Author: Jolanrensen

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

@@ -58,7 +58,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMaxOf(): T = rowMaxOfOrN
 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)
+    Aggregators.max<C>().aggregateFor(this, columns)
 
 public fun <T> DataFrame<T>.maxFor(vararg columns: String): DataRow<T> = maxFor { columns.toComparableColumns() }
 
@@ -84,7 +84,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.max(vararg columns: KProperty<C?>
     maxOrNull(*columns).suggestIfNull("max")
 
 public fun <T, C : Comparable<C>> DataFrame<T>.maxOrNull(columns: ColumnsSelector<T, C?>): C? =
-    Aggregators.max.aggregateAll(this, columns) as C?
+    Aggregators.max<C>().aggregateAll(this, columns) as C?
 
 public fun <T> DataFrame<T>.maxOrNull(vararg columns: String): Comparable<Any?>? =
     maxOrNull { columns.toComparableColumns() }
@@ -140,7 +140,7 @@ public fun <T> Grouped<T>.max(): DataFrame<T> = maxFor(intraComparableColumns())
 @Refine
 @Interpretable("GroupByMax0")
 public fun <T, C : Comparable<C>> Grouped<T>.maxFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
-    Aggregators.max.aggregateFor(this, columns)
+    Aggregators.max<C>().aggregateFor(this, columns)
 
 public fun <T> Grouped<T>.maxFor(vararg columns: String): DataFrame<T> = maxFor { columns.toComparableColumns() }
 
@@ -155,7 +155,7 @@ public fun <T, C : Comparable<C>> Grouped<T>.maxFor(vararg columns: KProperty<C?
 @Refine
 @Interpretable("GroupByMax0")
 public fun <T, C : Comparable<C>> Grouped<T>.max(name: String? = null, columns: ColumnsSelector<T, C?>): DataFrame<T> =
-    Aggregators.max.aggregateAll(this, name, columns)
+    Aggregators.max<C>().aggregateAll(this, name, columns)
 
 public fun <T> Grouped<T>.max(vararg columns: String, name: String? = null): DataFrame<T> =
     max(name) { columns.toComparableColumns() }
@@ -175,7 +175,7 @@ public fun <T, C : Comparable<C>> Grouped<T>.max(vararg columns: KProperty<C?>,
 public fun <T, C : Comparable<C>> Grouped<T>.maxOf(
     name: String? = null,
     expression: RowExpression<T, C>,
-): DataFrame<T> = Aggregators.max.aggregateOfDelegated(this, name) { maxOfOrNull(expression) }
+): DataFrame<T> = Aggregators.max<C>().aggregateOfDelegated(this, name) { maxOfOrNull(expression) }
 
 @Interpretable("GroupByReduceExpression")
 public fun <T, G, R : Comparable<R>> GroupBy<T, G>.maxBy(rowExpression: RowExpression<G, R?>): ReducedGroupBy<T, G> =
@@ -256,7 +256,7 @@ public fun <T> PivotGroupBy<T>.max(separate: Boolean = false): DataFrame<T> = ma
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.maxFor(
     separate: Boolean = false,
     columns: ColumnsForAggregateSelector<T, R?>,
-): DataFrame<T> = Aggregators.max.aggregateFor(this, separate, columns)
+): DataFrame<T> = Aggregators.max<R>().aggregateFor(this, separate, columns)
 
 public fun <T> PivotGroupBy<T>.maxFor(vararg columns: String, separate: Boolean = false): DataFrame<T> =
     maxFor(separate) { columns.toComparableColumns() }
@@ -274,7 +274,7 @@ public fun <T, R : Comparable<R>> PivotGroupBy<T>.maxFor(
 ): DataFrame<T> = maxFor(separate) { columns.toColumnSet() }
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.max(columns: ColumnsSelector<T, R?>): DataFrame<T> =
-    Aggregators.max.aggregateAll(this, columns)
+    Aggregators.max<R>().aggregateAll(this, columns)
 
 public fun <T> PivotGroupBy<T>.max(vararg columns: String): DataFrame<T> = max { columns.toComparableColumns() }
 

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

@@ -12,11 +12,12 @@ import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 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.values
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 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.aggregateByOrNull
 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.columns.toComparableColumns
 import org.jetbrains.kotlinx.dataframe.impl.indexOfMin
@@ -27,17 +28,19 @@ import kotlin.reflect.KProperty
 
 public fun <T : Comparable<T>> DataColumn<T?>.min(): T = minOrNull().suggestIfNull("min")
 
-public fun <T : Comparable<T>> DataColumn<T?>.minOrNull(): T? = asSequence().filterNotNull().minOrNull()
+public fun <T : Comparable<T>> DataColumn<T?>.minOrNull(): T? = Aggregators.min<T>().aggregate(this)
 
-public fun <T, R : Comparable<R>> DataColumn<T>.minBy(selector: (T) -> R): T =
+public inline fun <T, reified R : Comparable<R>> DataColumn<T>.minBy(noinline selector: (T) -> R?): T =
     minByOrNull(selector).suggestIfNull("minBy")
 
-public fun <T, R : Comparable<R>> DataColumn<T>.minByOrNull(selector: (T) -> R): T? = values.minByOrNull(selector)
+public inline fun <T, reified R : Comparable<R>> DataColumn<T>.minByOrNull(noinline selector: (T) -> R?): T? =
+    Aggregators.min<R>().aggregateByOrNull(this, selector)
 
-public fun <T, R : Comparable<R>> DataColumn<T>.minOf(selector: (T) -> R): R =
+public inline fun <T, reified R : Comparable<R>> DataColumn<T>.minOf(crossinline selector: (T) -> R?): R =
     minOfOrNull(selector).suggestIfNull("minOf")
 
-public fun <T, R : Comparable<R>> DataColumn<T>.minOfOrNull(selector: (T) -> R): R? = values.minOfOrNull(selector)
+public inline fun <T, reified R : Comparable<R>> DataColumn<T>.minOfOrNull(crossinline selector: (T) -> R?): R? =
+    Aggregators.min<R>().aggregateOf(this, selector)
 
 // endregion
 
@@ -47,6 +50,8 @@ public fun AnyRow.rowMinOrNull(): Any? = values().filterIsInstance<Comparable<*>
 
 public fun AnyRow.rowMin(): Any = rowMinOrNull().suggestIfNull("rowMin")
 
+// todo add rowMinBy?
+
 public inline fun <reified T : Comparable<T>> AnyRow.rowMinOfOrNull(): T? = values().filterIsInstance<T>().minOrNull()
 
 public inline fun <reified T : Comparable<T>> AnyRow.rowMinOf(): T = rowMinOfOrNull<T>().suggestIfNull("rowMinOf")
@@ -58,7 +63,7 @@ public inline fun <reified T : Comparable<T>> AnyRow.rowMinOf(): T = rowMinOfOrN
 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)
+    Aggregators.min<C>().aggregateFor(this, columns)
 
 public fun <T> DataFrame<T>.minFor(vararg columns: String): DataRow<T> = minFor { columns.toComparableColumns() }
 
@@ -84,7 +89,7 @@ public fun <T, C : Comparable<C>> DataFrame<T>.min(vararg columns: KProperty<C?>
     minOrNull(*columns).suggestIfNull("min")
 
 public fun <T, C : Comparable<C>> DataFrame<T>.minOrNull(columns: ColumnsSelector<T, C?>): C? =
-    Aggregators.min.aggregateAll(this, columns) as C?
+    Aggregators.min<C>().aggregateAll(this, columns) as C?
 
 public fun <T> DataFrame<T>.minOrNull(vararg columns: String): Comparable<Any?>? =
     minOrNull { columns.toComparableColumns() }
@@ -140,7 +145,7 @@ public fun <T> Grouped<T>.min(): DataFrame<T> = minFor(intraComparableColumns())
 @Refine
 @Interpretable("GroupByMin0")
 public fun <T, C : Comparable<C>> Grouped<T>.minFor(columns: ColumnsForAggregateSelector<T, C?>): DataFrame<T> =
-    Aggregators.min.aggregateFor(this, columns)
+    Aggregators.min<C>().aggregateFor(this, columns)
 
 public fun <T> Grouped<T>.minFor(vararg columns: String): DataFrame<T> = minFor { columns.toComparableColumns() }
 
@@ -155,7 +160,7 @@ public fun <T, C : Comparable<C>> Grouped<T>.minFor(vararg columns: KProperty<C?
 @Refine
 @Interpretable("GroupByMin0")
 public fun <T, C : Comparable<C>> Grouped<T>.min(name: String? = null, columns: ColumnsSelector<T, C?>): DataFrame<T> =
-    Aggregators.min.aggregateAll(this, name, columns)
+    Aggregators.min<C>().aggregateAll(this, name, columns)
 
 public fun <T> Grouped<T>.min(vararg columns: String, name: String? = null): DataFrame<T> =
     min(name) { columns.toComparableColumns() }
@@ -175,7 +180,7 @@ public fun <T, C : Comparable<C>> Grouped<T>.min(vararg columns: KProperty<C?>,
 public fun <T, C : Comparable<C>> Grouped<T>.minOf(
     name: String? = null,
     expression: RowExpression<T, C>,
-): DataFrame<T> = Aggregators.min.aggregateOfDelegated(this, name) { minOfOrNull(expression) }
+): DataFrame<T> = Aggregators.min<C>().aggregateOfDelegated(this, name) { minOfOrNull(expression) }
 
 @Interpretable("GroupByReduceExpression")
 public fun <T, G, R : Comparable<R>> GroupBy<T, G>.minBy(rowExpression: RowExpression<G, R?>): ReducedGroupBy<T, G> =
@@ -257,7 +262,7 @@ public fun <T> PivotGroupBy<T>.min(separate: Boolean = false): DataFrame<T> = mi
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.minFor(
     separate: Boolean = false,
     columns: ColumnsForAggregateSelector<T, R?>,
-): DataFrame<T> = Aggregators.min.aggregateFor(this, separate, columns)
+): DataFrame<T> = Aggregators.min<R>().aggregateFor(this, separate, columns)
 
 public fun <T> PivotGroupBy<T>.minFor(vararg columns: String, separate: Boolean = false): DataFrame<T> =
     minFor(separate) { columns.toComparableColumns() }
@@ -275,7 +280,7 @@ public fun <T, R : Comparable<R>> PivotGroupBy<T>.minFor(
 ): DataFrame<T> = minFor(separate) { columns.toColumnSet() }
 
 public fun <T, R : Comparable<R>> PivotGroupBy<T>.min(columns: ColumnsSelector<T, R?>): DataFrame<T> =
-    Aggregators.min.aggregateAll(this, columns)
+    Aggregators.min<R>().aggregateAll(this, columns)
 
 public fun <T> PivotGroupBy<T>.min(vararg columns: String): DataFrame<T> = min { columns.toComparableColumns() }
 

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/Aggregators.kt

@@ -3,6 +3,7 @@ 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.minOrNull
 import org.jetbrains.kotlinx.dataframe.math.percentile
 import org.jetbrains.kotlinx.dataframe.math.std
 import org.jetbrains.kotlinx.dataframe.math.stdTypeConversion
@@ -86,13 +87,23 @@ internal object Aggregators {
         getAggregator: (Param1, Param2) -> AggregatorProvider<AggregatorType>,
     ) = AggregatorOptionSwitch2.Factory(getAggregator)
 
+    private fun <Value, Return> AggregatorProvider<Aggregator<Value, Return>>.asByAggregator(
+        aggregatorBy: AggregateBy<Any?, Value, Return>,
+    ) where Value : Any, Value : Comparable<Value> = ByAggregator.Factory(this, aggregatorBy)
+
     // T: Comparable<T> -> T?
-    val min by twoStepPreservingType<Comparable<Any?>> {
-        minOrNull()
+    fun <T : Comparable<T>?> min() = min.cast2<T, T?>()
+
+    private val min by twoStepPreservingType<Comparable<Any?>?> { type ->
+        minOrNull(type)
+    }.asByAggregator { sourceType, valueType, selector ->
+        minByOrNull(selector)
     }
 
     // T: Comparable<T> -> T?
-    val max by twoStepPreservingType<Comparable<Any?>> {
+    fun <T : Comparable<T>?> max() = max.cast2<T, T?>()
+
+    private val max by twoStepPreservingType<Comparable<Any?>> {
         maxOrNull()
     }
 

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

@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.impl.aggregation.modes
 
+import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataFrameExpression
 import org.jetbrains.kotlinx.dataframe.DataRow
@@ -8,10 +9,13 @@ 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.aggregation.aggregators.Aggregator
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.Aggregators
 import org.jetbrains.kotlinx.dataframe.impl.namedValues
+import kotlin.reflect.typeOf
 
 @CandidateForRemoval
-internal fun <T> Grouped<T>.aggregateBy(body: DataFrameExpression<T, DataRow<T>?>): DataFrame<T> {
+internal fun <T> Grouped<T>.aggregateByOrNull(body: DataFrameExpression<T, DataRow<T>?>): DataFrame<T> {
     require(this is GroupBy<*, T>)
     val keyColumns = keys.columnNames().toSet()
     return aggregateInternal {
@@ -21,3 +25,40 @@ internal fun <T> Grouped<T>.aggregateBy(body: DataFrameExpression<T, DataRow<T>?
         }
     }.cast()
 }
+
+/**
+ * Aggregates the values of the column using the provided [Aggregator] `by` the provided [selector].
+ * H
+ */
+@Suppress("UNCHECKED_CAST")
+@PublishedApi
+internal inline fun <C, reified V : Comparable<V>> Aggregator<V, V?>.aggregateByOrNull(
+    values: Iterable<C>,
+    noinline selector: (C) -> V?,
+): C? =
+    when (name) { // todo?
+        Aggregators.min<Comparable<Any?>>().name ->
+            values.asSequence()
+                .filterNot { selector(it) == null }
+                .minByOrNull(selector as (C) -> V)
+
+        Aggregators.max<Comparable<Any?>>().name ->
+            values.asSequence()
+                .filterNot { selector(it) == null }
+                .maxByOrNull(selector as (C) -> V)
+
+        else -> {
+            // less efficient but more generic
+            val aggregateResult = aggregate(
+                values = values.asSequence().map { selector(it) }.asIterable(),
+                type = typeOf<V>(),
+            )
+            values.first { selector(it) == aggregateResult }
+        }
+    }
+
+@PublishedApi
+internal inline fun <C, reified V : Comparable<V>?> Aggregator<V, V?>.aggregateByOrNull(
+    column: DataColumn<C>,
+    noinline selector: (C) -> V?,
+): C? = aggregateByOrNull(column.values(), selector)

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

@@ -0,0 +1,20 @@
+package org.jetbrains.kotlinx.dataframe.math
+
+import org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators.preserveReturnTypeNullIfEmpty
+import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
+
+@PublishedApi
+internal fun <T : Comparable<T>> Iterable<T?>.minOrNull(type: KType): T? = asSequence().minOrNull(type)
+
+@Suppress("UNCHECKED_CAST")
+@PublishedApi
+internal fun <T : Comparable<T>> Sequence<T?>.minOrNull(type: KType): T? {
+    if (type.isMarkedNullable) {
+        return filterNotNull().minOrNull(type.withNullability(false))
+    }
+    return (this as Sequence<T>).minOrNull()
+}
+
+/** T: Comparable<T> -> T(?) */
+internal val minTypeConversion = preserveReturnTypeNullIfEmpty