Automated commit of generated code

Author: github-actions[bot]

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/concat.kt

@@ -3,9 +3,12 @@ package org.jetbrains.kotlinx.dataframe.api
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.values
 import org.jetbrains.kotlinx.dataframe.impl.api.concatImpl
 import org.jetbrains.kotlinx.dataframe.impl.asList
+import org.jetbrains.kotlinx.dataframe.type
 
 // region DataColumn
 
@@ -40,6 +43,64 @@ public fun <T> DataFrame<T>.concat(frames: Iterable<DataFrame<T>>): DataFrame<T>
 
 public fun <T, G> GroupBy<T, G>.concat(): DataFrame<G> = groups.concat()
 
+/**
+ * Concatenates all groups in this [GroupBy] into a single [DataFrame],
+ * preserving and including all grouping key columns that are not present in the group's columns.
+ *
+ * Doesn't affect key columns that have the same name as columns inside the groups (even if their content differs).
+ *
+ * This function is especially useful when grouping by expressions or renamed columns,
+ * and you want the resulting [DataFrame] to include those keys as part of the output.
+ *
+ * ### Example
+ *
+ * ```kotlin
+ * val df = dataFrameOf(
+ *     "value" to listOf(1, 2, 3, 3),
+ *     "type" to listOf("a", "b", "a", "b")
+ * )
+ *
+ * val gb = df.groupBy { expr { "Category: ${type.uppercase()}" } named "category" }
+ * ```
+ *
+ * A regular `concat()` will return a [DataFrame] similar to the original `df`
+ * (with the same columns and rows but in the different orders):
+ *
+ * ```
+ * gb.concat()
+ * ```
+ * | value | type |
+ * | :---- | :--- |
+ * | 1     | a    |
+ * | 3     | a    |
+ * | 2     | b    |
+ * | 3     | b    |
+ *
+ * But `concatWithKeys()` will include the new "category" key column:
+ *
+ * ```
+ * gb.concatWithKeys()
+ * ```
+ * | value | type | category      |
+ * | :---- | :--- | :------------ |
+ * | 1     | a    | Category: A   |
+ * | 3     | a    | Category: A   |
+ * | 2     | b    | Category: B   |
+ * | 3     | b    | Category: B   |
+ *
+ * @return A new [DataFrame] where all groups are combined and additional key columns are included in each row.
+ */
+@Refine
+@Interpretable("ConcatWithKeys")
+public fun <T, G> GroupBy<T, G>.concatWithKeys(): DataFrame<G> =
+    mapToFrames {
+        val rowsCount = group.rowsCount()
+        val keyColumns = keys.columns().filter { it.name !in group.columnNames() }.map { keyColumn ->
+            DataColumn.createByType(keyColumn.name, List(rowsCount) { key[keyColumn] }, keyColumn.type)
+        }
+        group.addAll(keyColumns)
+    }.concat()
+
 // endregion
 
 // region ReducedGroupBy

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/concat.kt

@@ -11,4 +11,16 @@ class ConcatTests {
         val b by columnOf(3.0, null)
         a.concat(b) shouldBe columnOf(1, 2, 3.0, null).named("a")
     }
+
+    @Test
+    fun `concat with keys`() {
+        val df = dataFrameOf(
+            "value" to listOf(1, 2, 3, 3),
+            "type" to listOf("a", "b", "a", "b"),
+        )
+        val gb = df.groupBy { expr { "Category: ${(this["type"] as String).uppercase()}" } named "category" }
+        val dfWithCategory = gb.concatWithKeys()
+
+        dfWithCategory.columnNames() shouldBe listOf("value", "type", "category")
+    }
 }