added migrations for standard functions of groupBy with tests

Author: Jolanrensen

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

@@ -20,7 +20,10 @@ public inline fun <T> DataFrame<T>.forEach(action: RowExpression<T, Unit>): Unit
 
 // region GroupBy
 
-@Deprecated("Replaced with forEachEntry")
+@Deprecated(
+    "Replaced with forEachEntry",
+    ReplaceWith("forEachEntry { val key = it\nval group = it.group()\nbody(key, group) }"),
+)
 public inline fun <T, G> GroupBy<T, G>.forEach(body: (GroupBy.Entry<T, G>) -> Unit): Unit =
     keys.forEach { key ->
         val group = groups[key.index()]

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

@@ -122,7 +122,7 @@ public interface GroupBy<out T, out G> : Grouped<G> {
 
     public fun <R> updateGroups(transform: Selector<DataFrame<G>, DataFrame<R>>): GroupBy<T, R>
 
-    @Deprecated("Replaced by filterEntries")
+    @Deprecated("Replaced by filterEntries", ReplaceWith("filterEntries(predicate)"))
     public fun filter(predicate: GroupedRowFilter<T, G>): GroupBy<T, G>
 
     public fun filterEntries(predicate: GroupByEntryFilter<T, G>): GroupBy<T, G>

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

@@ -142,19 +142,28 @@ public inline fun <T> DataFrame<T>.mapToFrame(body: AddDsl<T>.() -> Unit): AnyFr
 
 // region GroupBy
 
-@Deprecated("Replaced by mapEntries")
+@Deprecated(
+    "Replaced by mapEntries",
+    ReplaceWith("mapEntries { val key = it\nval group = it.group()\nbody(key, group) }"),
+)
 public inline fun <T, G, R> GroupBy<T, G>.map(body: Selector<GroupWithKey<T, G>, R>): List<R> =
     keys.rows().mapIndexedNotNull { index, row ->
         val group = groups[index]
         val g = GroupWithKey(row, group)
         body(g, g)
     }
 
-@Deprecated("Replaced by mapEntriesToRows")
+@Deprecated(
+    "Replaced by mapEntriesToRows",
+    ReplaceWith("mapEntriesToRows { val key = it\nval group = it.group()\nbody(key, group) }"),
+)
 public fun <T, G> GroupBy<T, G>.mapToRows(body: Selector<GroupWithKey<T, G>, DataRow<G>?>): DataFrame<G> =
     map(body).concat()
 
-@Deprecated("Replaced by mapEntriesToFrames")
+@Deprecated(
+    "Replaced by mapEntriesToFrames",
+    ReplaceWith("mapEntriesToFrames { val key = it\nval group = it.group()\nbody(key, group) }"),
+)
 public fun <T, G> GroupBy<T, G>.mapToFrames(body: Selector<GroupWithKey<T, G>, DataFrame<G>>): FrameColumn<G> =
     DataColumn.createFrameColumn(groups.name, map(body))
 
@@ -164,10 +173,11 @@ public inline fun <T, G, R> GroupBy<T, G>.mapEntries(body: GroupByEntrySelector<
         body(entry, entry)
     }
 
-public fun <T, G> GroupBy<T, G>.mapEntriesToRows(body: GroupByEntrySelector<T, G, DataRow<G>?>): DataFrame<G> =
+public fun <T, G, R : Any> GroupBy<T, G>.mapEntriesToRows(body: GroupByEntrySelector<T, G, DataRow<R>?>): DataFrame<R> =
     mapEntries(body).concat()
 
-public fun <T, G> GroupBy<T, G>.mapEntriesToFrames(body: GroupByEntrySelector<T, G, DataFrame<G>>): FrameColumn<G> =
-    DataColumn.createFrameColumn(groups.name, mapEntries(body))
+public fun <T, G, R : Any> GroupBy<T, G>.mapEntriesToFrames(
+    body: GroupByEntrySelector<T, G, DataFrame<R>>,
+): FrameColumn<R> = DataColumn.createFrameColumn(groups.name, mapEntries(body))
 
 // endregion

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

@@ -47,9 +47,9 @@ internal class GroupByImpl<T, G>(
     override val keys by lazy { df.remove { groups } }
 
     @Suppress("UNCHECKED_CAST")
-    override fun <R> updateGroups(transform: Selector<DataFrame<G>, DataFrame<R>>) =
+    override fun <R> updateGroups(transform: Selector<DataFrame<G>, DataFrame<R>>): GroupBy<T, R> =
         df.convert { groups }.with { transform(it, it) }
-            .asGroupBy(groups.name()) as GroupBy<T, R>
+            .asGroupBy { frameCol<R>(groups.name()) }
 
     override fun toString() = df.toString()
 
@@ -62,15 +62,15 @@ internal class GroupByImpl<T, G>(
             val row = GroupedDataRowImpl(df[it], groups)
             predicate(row, row)
         }
-        return df[indices].asGroupBy { groups }
+        return df[indices].asGroupBy { frameCol<G>(groups.name()) }
     }
 
     override fun filterEntries(predicate: GroupByEntryFilter<T, G>): GroupBy<T, G> {
         val indices = (0 until df.nrow).filter {
             val row = GroupByEntryImpl(df[it], groups)
             predicate(row, row)
         }
-        return df[indices].asGroupBy { groups }
+        return df[indices].asGroupBy { frameCol<G>(groups.name()) }
     }
 
     override fun toDataFrame(groupedColumnName: String?): DataFrame<T> =

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/groupBy.kt

@@ -2,11 +2,13 @@ package org.jetbrains.kotlinx.dataframe.api
 
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.testSets.person.BaseTest
+import org.jetbrains.kotlinx.dataframe.testSets.person.age
 import org.junit.Test
 import kotlin.reflect.typeOf
 
 @Suppress("ktlint:standard:argument-list-wrapping")
-class GroupByTests {
+class GroupByTests : BaseTest() {
 
     @Test
     fun `groupBy values with nulls`() {
@@ -56,4 +58,80 @@ class GroupByTests {
             getFrameColumn("d") into "e"
         }["e"].type() shouldBe typeOf<List<AnyFrame>>()
     }
+
+    @Test
+    fun `groupBy forEachEntry`() {
+        val grouped = typed.groupBy { age }
+        val entries1 = buildList {
+            grouped.forEach { (key, group) ->
+                add(key.toMap() to group)
+            }
+        }
+        val entries2 = buildList {
+            grouped.forEachEntry {
+                add(it.toMap() to it.group())
+            }
+        }
+
+        entries1 shouldBe entries2
+    }
+
+    @Test
+    fun `groupBy mapEntries`() {
+        // old mapToRows, and mapToFrames stick to the same return type, so let's make the types Any? for the test
+        val grouped: GroupBy<Any?, Any?> = df.groupBy { age }
+        val entries1 = grouped.map { (key, group) ->
+            key.toMap() to group
+        }
+        val entries2 = grouped.mapEntries {
+            it.toMap() to it.group()
+        }
+        entries1 shouldBe entries2
+
+        val entries3 = grouped.mapToRows { (key, group) ->
+            listOf(key.toMap() to group).toDataFrame().single()
+        }
+        val entries4 = grouped.mapEntriesToRows {
+            listOf(it.toMap() to it.group()).toDataFrame().single()
+        }
+        entries3 shouldBe entries4
+
+        val entries5 = grouped.mapToFrames { (key, group) ->
+            listOf(key.toMap() to group).toDataFrame()
+        }
+        val entries6 = grouped.mapEntriesToFrames {
+            listOf(it.toMap() to it.group()).toDataFrame()
+        }
+        entries5 shouldBe entries6
+
+        // let's test the -Entries variants with typed versions
+        val grouped2 = typed.groupBy { age }
+
+        val entries7 = grouped2.mapEntries {
+            it.toMap() to it.group()
+        }
+        val entries8 = grouped2.mapEntriesToRows {
+            listOf(it.toMap() to it.group()).toDataFrame().single()
+        }.toList()
+        val entries9 = grouped2.mapEntriesToFrames {
+            listOf(it.toMap() to it.group()).toDataFrame()
+        }.map { it[0][0] to it[0][1] }.toList()
+        entries7 shouldBe entries8
+        entries8 shouldBe entries9
+    }
+
+    @Test
+    fun `groupBy filterEntries`() {
+        val grouped = typed.groupBy { age }
+
+        val entries1 = grouped.filter { age == 20 }
+            .mapEntries {
+                it.toMap() to it.group()
+            }
+        val entries2 = grouped.filterEntries { age == 20 }
+            .mapEntries {
+                it.toMap() to it.group()
+            }
+        entries1 shouldBe entries2
+    }
 }