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

Author: Jolanrensen

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 {

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

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

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?
 }

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) })
 }

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,18 +149,16 @@ 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,
         aggregate: Aggregate<Value, Return>,
     ) = 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].

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,
             )
         })
 }

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,27 +65,41 @@ 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].
      *
      * @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<TwoStepAggregator<Value, Return>> by AggregatorProvider({ name ->
             TwoStepAggregator(
                 name = name,
                 getReturnTypeOrNull = getReturnTypeOrNull,
                 stepOneAggregator = stepOneAggregator,
                 stepTwoAggregator = stepTwoAggregator,
-                preservesType = preservesType,
             )
         })
 }