Add support for DataFrame `sum` operation with tests #1148

zaleslaw · 2025-04-23T13:25:33Z

No description provided.

Introduced the `sum` operation for DataFrames, supporting numerical columns aggregation. Updated relevant tests and added new test cases to verify functionality. Included schema modifications for handling numerical column operations.

Converted various internal classes, interfaces, and functions related to aggregation into public entities. This change expands their visibility, enabling external usage and facilitating integration with other modules or libraries.

…re compatibility and correctness in sum calculations.

Copilot

Pull Request Overview

This PR adds support for the DataFrame sum operation along with comprehensive tests covering various summation scenarios and aggregation handlers. Key changes include:

Adding a new sum operation test in both generated tests and dedicated test data for verifying correct summation.
Extending the aggregation framework with new Sum0 and Sum1 implementations.
Changing visibility modifiers from internal to public in several core aggregator components.

Reviewed Changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java	Adds a new test method for the sum operation
plugins/kotlin-dataframe/testData/box/sum.kt	Introduces test scenarios for sum over all, selective, and expression-based columns
plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt	Updates load interpreter to support new Sum aggregators
plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/statistics.kt	Adds implementations for summation aggregators
core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt	Expands tests for groupBy operations to include new numerical columns
core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/sum.kt	Annotates and exposes the new sum APIs
Other core aggregator files	Adjusts visibility and typing to support public sum operation API

Comments suppressed due to low confidence (1)

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/statistics.kt:571

In the groupBy maxBy test, the column names are checked using 'res4' while the aggregation result is later accessed from 'res5'. This inconsistency might lead to erroneous test behavior; ensure the same result object is used throughout the test.

res4.columnNames() shouldBe listOf("city", "name", "age", "weight", "height", "yearsToRetirement", "workExperienceYears", "dependentsCount", "annualIncome")

plugins/kotlin-dataframe/testData/box/sum.kt

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

koperagen · 2025-04-24T10:04:41Z

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/statistics.kt

+    val kClass = this.classifier as? KClass<*> ?: return null
+    val classId = kClass.toClassId() ?: return null
+
+    return ConeClassLikeTypeImpl(


There's FirSession available inside Arguments scope to access primitive types

fun Arguments.test() { session.builtinTypes.doubleType.coneType }

For class id, there's StandardClassIds
And, if needed, there's ClassId.constructClassLikeType instead of calling ConeClassLikeTypeImpl directly
Please use these APIs

koperagen · 2025-04-24T10:28:15Z

...rame/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

@@ -233,6 +233,12 @@ public void testFlexibleReturnType() {
    runTest("testData/box/flexibleReturnType.kt");
  }

+  @Test
+  @TestMetadata("sum.kt")


After you create new file in testData, please use this task to update this

koperagen · 2025-04-24T10:30:16Z

plugins/kotlin-dataframe/testData/box/sum.kt

+    val res11 = personsDf.sumFor { dependentsCount }
+    val sum111: Int? = res11.dependentsCount
+
+    // scenario #2: sum of all values in two columns of Int hierarchy of types


Let's remove scenario 2 from plugin tests for DataFrame.sum? It returns Number so no schema changes

Jolanrensen · 2025-04-24T11:00:34Z

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

-    val inputHandler: AggregatorInputHandler<Value, Return>,
-    val multipleColumnsHandler: AggregatorMultipleColumnsHandler<Value, Return>,
-    val name: String,
+public class Aggregator<in Value : Any, Return : Any?>(


why'd you remove out? it should only appear in the out-position of the aggregation, at least, that's the idea

Jolanrensen · 2025-04-24T11:01:31Z

...aframe/impl/aggregation/aggregators/multipleColumnsHandlers/TwoStepMultipleColumnsHandler.kt

@@ -31,7 +31,7 @@ import kotlin.reflect.KType
 *   If not supplied, the handler of the first step is reused.
 * @see [FlatteningMultipleColumnsHandler]
 */
-internal class TwoStepMultipleColumnsHandler<in Value : Any, out Return : Any?>(
+internal class TwoStepMultipleColumnsHandler<in Value : Any, Return : Any?>(


same here with out, though, I could understand the removal here because Return is also used as input for stepTwo

Jolanrensen · 2025-04-24T11:07:21Z

...otlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/AggregatorOptionSwitch.kt

@@ -31,7 +30,7 @@ internal class AggregatorOptionSwitch1<in Param1, in Value : Any, out Return : A
         *   MyAggregator.Factory(param1)
         * }
         */
-        fun <Param1, Value : Any, Return : Any?> Factory(
+        internal fun <Param1, Value : Any, Return : Any?> Factory(


if we make everything public, the factory can also be public imo

Jolanrensen · 2025-04-24T11:13:09Z

plugins/kotlin-dataframe/testData/box/sum.kt

+    // scenario #0: all numerical columns
+    val res0 = personsDf.sum()
+
+    val sum01: Int? = res0.age


The sum of a column is never nullable, right?

no, and i think it can't be helped

alternatively, do we really return nullable int for an empty dataframe? maybe we should throw an exception instead?
but we should be careful about how it will affect GroupBy.sum

For min/max/median/percentile, the result is nullable however, but only null if the input is empty.

@koperagen, is there a compile-time check for when a DF is empty or not? Otherwise we could generate the result-type more tightly.

There's no such check

alternatively, do we really return nullable int for an empty dataframe? maybe we should throw an exception instead? but we should be careful about how it will affect GroupBy.sum

No we don't, the result of sum is never null, the sum of an empty Int column is 0. Even if we have a column of type Nothing? we have a non-null return type, defaulting to 0.0.

There's no such check

Okay then for those we can follow what I wrote here #1148 (comment)

Jolanrensen · 2025-04-24T11:41:47Z

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/statistics.kt

+): SimpleCol {
+    val originalType = column.type.type
+    val inputKType = originalType.toKType()
+    val resultKType = inputKType?.let { statisticAggregator.calculateReturnType(it, emptyInput = false) }


might cause issues for min/max/median/percentile to set emptyInput = false. I'd set it to true. The name emptyInput is also kind of a misnomer, it should be "inputPotentiallyEmpty" or something.

It changes intCol.min() from Int (when it's set to false) to Int? (when it's set to true).

For intCol.sum() the result will always be Int, regardless of what it's set to.

Feel free to rename the argument everywhere it's used :)

zaleslaw added 4 commits April 23, 2025 15:24

Add support for DataFrame sum operation with tests

7729d03

Introduced the `sum` operation for DataFrames, supporting numerical columns aggregation. Updated relevant tests and added new test cases to verify functionality. Included schema modifications for handling numerical column operations.

Merge branch 'master' into issue-1138

4b4feb8

Make aggregator-related classes and functions public

0376fec

Converted various internal classes, interfaces, and functions related to aggregation into public entities. This change expands their visibility, enabling external usage and facilitating integration with other modules or libraries.

Enhance type conversions between KType and ConeKotlinType to ensu…

644685d

…re compatibility and correctness in sum calculations.

zaleslaw requested review from Copilot, Jolanrensen and koperagen April 23, 2025 16:34

Copilot AI reviewed Apr 23, 2025

View reviewed changes

plugins/kotlin-dataframe/testData/box/sum.kt Outdated Show resolved Hide resolved

zaleslaw and others added 2 commits April 23, 2025 18:39

Update plugins/kotlin-dataframe/testData/box/sum.kt

2db823f

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Refactor type conversion and column handling logic

9844b7d

koperagen reviewed Apr 24, 2025

View reviewed changes

Jolanrensen reviewed Apr 24, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for DataFrame `sum` operation with tests #1148

Add support for DataFrame `sum` operation with tests #1148

zaleslaw commented Apr 23, 2025

Copilot AI left a comment

koperagen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025

koperagen Apr 24, 2025

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025

Jolanrensen Apr 24, 2025 •

edited

Loading

Add support for DataFrame sum operation with tests #1148

Are you sure you want to change the base?

Add support for DataFrame sum operation with tests #1148

Conversation

zaleslaw commented Apr 23, 2025

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

koperagen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

koperagen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jolanrensen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

koperagen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

koperagen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jolanrensen Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

Add support for DataFrame `sum` operation with tests #1148

Add support for DataFrame `sum` operation with tests #1148

koperagen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

Jolanrensen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

koperagen Apr 24, 2025 •

edited

Loading

Jolanrensen Apr 24, 2025 •

edited

Loading