Merge branch 'master' into mean

Author: Jolanrensen

core/api/core.api

@@ -5310,9 +5310,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,
             )
         })
 }