Ensure a predictable order of columns after aggregation of GroupBy

Author: koperagen

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

@@ -356,12 +356,16 @@ public fun <T> DataFrame<T>.asColumnGroup(column: ColumnGroupAccessor<T>): Colum
 
 // region as GroupedDataFrame
 
-public fun <T> DataFrame<T>.asGroupBy(groupedColumnName: String): GroupBy<T, T> =
-    GroupByImpl(this, getFrameColumn(groupedColumnName).castFrameColumn()) { none() }
+public fun <T> DataFrame<T>.asGroupBy(groupedColumnName: String): GroupBy<T, T> {
+    val groups = getFrameColumn(groupedColumnName)
+    return asGroupBy { groups.cast() }
+}
 
 @AccessApiOverload
-public fun <T, G> DataFrame<T>.asGroupBy(groupedColumn: ColumnReference<DataFrame<G>>): GroupBy<T, G> =
-    GroupByImpl(this, getFrameColumn(groupedColumn.name()).castFrameColumn()) { none() }
+public fun <T, G> DataFrame<T>.asGroupBy(groupedColumn: ColumnReference<DataFrame<G>>): GroupBy<T, G> {
+    val groups = getFrameColumn(groupedColumn.name()).castFrameColumn<G>()
+    return asGroupBy { groups }
+}
 
 public fun <T> DataFrame<T>.asGroupBy(): GroupBy<T, T> {
     val groupCol = columns().single { it.isFrameColumn() }.asAnyFrameColumn().castFrameColumn<T>()
@@ -370,7 +374,7 @@ public fun <T> DataFrame<T>.asGroupBy(): GroupBy<T, T> {
 
 public fun <T, G> DataFrame<T>.asGroupBy(selector: ColumnSelector<T, DataFrame<G>>): GroupBy<T, G> {
     val column = getColumn(selector).asFrameColumn()
-    return GroupByImpl(this, column) { none() }
+    return GroupByImpl(this.move { column }.toEnd(), column) { none() }
 }
 
 // endregion

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/PlaylistJsonTest.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.io
 
+import io.kotest.assertions.withClue
 import io.kotest.matchers.shouldBe
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
@@ -220,8 +221,10 @@ class PlaylistJsonTest {
             minBy { snippet.publishedAt }.snippet into "earliest"
         }
 
-        res.columnsCount() shouldBe typed.columnsCount() + 1
-        res.getColumnIndex("earliest") shouldBe typed.getColumnIndex("items") + 1
+        withClue(res.columnNames() to typed.columnNames()) {
+            res.columnsCount() shouldBe typed.columnsCount() + 1
+            res.getColumnIndex("earliest") shouldBe res.columns().indices.last
+        }
 
         val expected = typed.items.map { it.snippet.minBy { publishedAt } }.toList()
         res["earliest"].toList() shouldBe expected

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

@@ -138,6 +138,28 @@ class DataFrameTreeTests : BaseTest() {
         res shouldBe typed
     }
 
+    @Test
+    fun asGroupByOverloads() {
+        val rowsColumn by columnOf(typed[0..3], typed[4..5], typed[6..6])
+        val df = dataFrameOf(rowsColumn)
+        val res = df.asGroupBy { rowsColumn }.max()
+        df.asGroupBy("rowsColumn").max() shouldBe res
+        df.asGroupBy(rowsColumn).max() shouldBe res
+    }
+
+    @Test
+    fun moveGroupedColumn() {
+        val df = dataFrameOf(
+            "group" to listOf(typed[0..3], typed[4..5], typed[6..6]),
+            "col" to listOf(1, 2, 3),
+        )
+
+        // We need to match the order of columns in the runtime and in compiler plugin
+        // GroupBy with the same schema should give the same result after aggregation, no matter how it was created
+        // We cannot track position of group in original df, so we align `asGroupBy` with `groupBy` and move `group` column to end
+        df.asGroupBy("group").toDataFrame().columnNames() shouldBe listOf("col", "group")
+    }
+
     @Test
     fun createFrameColumn2() {
         val id by column(typed.indices())