Made sure aggregating an empty grouped dataframe still generates aggregation columns with test. Issue #1531

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt

@@ -8,7 +8,9 @@ import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedBody
 import org.jetbrains.kotlinx.dataframe.aggregation.NamedValue
 import org.jetbrains.kotlinx.dataframe.api.GroupBy
 import org.jetbrains.kotlinx.dataframe.api.GroupedRowFilter
+import org.jetbrains.kotlinx.dataframe.api.asFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.asGroupBy
+import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.convert
 import org.jetbrains.kotlinx.dataframe.api.getColumn
@@ -18,6 +20,7 @@ import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.api.rename
+import org.jetbrains.kotlinx.dataframe.api.take
 import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.AggregatableInternal
 import org.jetbrains.kotlinx.dataframe.impl.aggregation.GroupByReceiverImpl
@@ -27,8 +30,10 @@ import org.jetbrains.kotlinx.dataframe.impl.api.GroupedDataRowImpl
 import org.jetbrains.kotlinx.dataframe.impl.api.insertImpl
 import org.jetbrains.kotlinx.dataframe.impl.api.removeImpl
 import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.impl.schema.createEmptyDataFrame
 import org.jetbrains.kotlinx.dataframe.ncol
 import org.jetbrains.kotlinx.dataframe.nrow
+import org.jetbrains.kotlinx.dataframe.size
 import org.jetbrains.kotlinx.dataframe.values
 
 /**
@@ -74,14 +79,25 @@ internal fun <T, G, R> aggregateGroupBy(
     body: AggregateGroupedBody<G, R>,
 ): DataFrame<T> {
     val defaultAggregateName = "aggregated"
-
+    val groupedDfIsEmpty = df.size().nrow == 0
     val column = df.getColumn(selector)
-
     val removed = df.removeImpl(columns = selector)
-
     val hasKeyColumns = removed.df.ncol > 0
 
-    val groupedFrame = column.values.map {
+    val groups =
+        if (groupedDfIsEmpty) {
+            // if the grouped dataframe is empty, make sure the provided AggregateGroupedBody is called at least once
+            // to create aggregated columns. We empty them below.
+            listOf(
+                column.asFrameColumn().schema.value
+                    .createEmptyDataFrame()
+                    .cast(),
+            )
+        } else {
+            column.values
+        }
+
+    val groupedFrame = groups.map {
         if (it == null) {
             null
         } else {
@@ -101,6 +117,11 @@ internal fun <T, G, R> aggregateGroupBy(
             builder.compute()
         }
     }.concat()
+        .let {
+            // empty the aggregated columns that were created by calling the provided AggregateGroupedBody once
+            // if the grouped dataframe is empty
+            if (groupedDfIsEmpty) it.take(0) else it
+        }
 
     val removedNode = removed.removedColumns.single()
     val insertPath = removedNode.pathFromRoot().dropLast(1)

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt

@@ -62,6 +62,7 @@ import org.jetbrains.kotlinx.dataframe.api.fill
 import org.jetbrains.kotlinx.dataframe.api.fillNulls
 import org.jetbrains.kotlinx.dataframe.api.filter
 import org.jetbrains.kotlinx.dataframe.api.first
+import org.jetbrains.kotlinx.dataframe.api.firstOrNull
 import org.jetbrains.kotlinx.dataframe.api.forEach
 import org.jetbrains.kotlinx.dataframe.api.forEachIndexed
 import org.jetbrains.kotlinx.dataframe.api.frameColumn
@@ -94,14 +95,17 @@ import org.jetbrains.kotlinx.dataframe.api.match
 import org.jetbrains.kotlinx.dataframe.api.matches
 import org.jetbrains.kotlinx.dataframe.api.max
 import org.jetbrains.kotlinx.dataframe.api.maxBy
+import org.jetbrains.kotlinx.dataframe.api.maxByOrNull
 import org.jetbrains.kotlinx.dataframe.api.mean
 import org.jetbrains.kotlinx.dataframe.api.meanFor
 import org.jetbrains.kotlinx.dataframe.api.meanOf
 import org.jetbrains.kotlinx.dataframe.api.median
+import org.jetbrains.kotlinx.dataframe.api.medianOrNull
 import org.jetbrains.kotlinx.dataframe.api.merge
 import org.jetbrains.kotlinx.dataframe.api.min
 import org.jetbrains.kotlinx.dataframe.api.minBy
 import org.jetbrains.kotlinx.dataframe.api.minOf
+import org.jetbrains.kotlinx.dataframe.api.minOrNull
 import org.jetbrains.kotlinx.dataframe.api.minus
 import org.jetbrains.kotlinx.dataframe.api.move
 import org.jetbrains.kotlinx.dataframe.api.moveTo
@@ -710,6 +714,37 @@ class DataFrameTests : BaseTest() {
         res.size() shouldBe 2
     }
 
+    // Issue #1531
+    @Test
+    fun `groupBy empty df should generate empty aggregation cols`() {
+        val empty = typed.take(0)
+        val resDf = empty.groupBy { name }.aggregate {
+            count() into "n"
+            count { age > 25 } into "old count"
+            medianOrNull { age } into "median age"
+            minOrNull { age } into "min age"
+            all { weight != null } into "all with weights"
+            maxByOrNull { age }?.city into "oldest origin"
+            sortBy { age }.firstOrNull()?.city into "youngest origin"
+            pivot { city.map { "from $it" } }.count()
+            age.toList() into "ages"
+        }
+
+        resDf.columnNames() shouldBe listOf(
+            "name",
+            "n",
+            "old count",
+            "median age",
+            "min age",
+            "all with weights",
+            "oldest origin",
+            "youngest origin",
+            "ages",
+        )
+
+        resDf.alsoDebug()
+    }
+
     @Test
     fun `groupBy`() {
         fun AnyFrame.check() {