Support adding ColumnGroups in add DSL.

Author: nikitinas

src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyCol
+import org.jetbrains.kotlinx.dataframe.AnyColumnGroupAccessor
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.Column
@@ -88,6 +89,7 @@ public inline fun <reified R, T, G> GroupBy<T, G>.add(
 
 public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsContainer<T> by df, ColumnSelectionDsl<T> {
 
+    // TODO: support adding column into path
     internal val columns = mutableListOf<AnyCol>()
 
     public fun add(column: Column): Boolean = columns.add(column.resolveSingle(df)!!.data)
@@ -104,6 +106,8 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
     ): Boolean = add(df.map(name, infer, expression))
 
     public inline infix fun <reified R> String.from(noinline expression: RowExpression<T, R>): Boolean = add(this, Infer.Nulls, expression)
+
+    // TODO: use path instead of name
     public inline infix fun <reified R> ColumnAccessor<R>.from(noinline expression: RowExpression<T, R>): Boolean = name().from(expression)
     public inline infix fun <reified R> KProperty<R>.from(noinline expression: RowExpression<T, R>): Boolean = add(name, Infer.Nulls, expression)
 
@@ -114,4 +118,21 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
     public infix fun Column.into(name: String): Boolean = add(rename(name))
     public infix fun <R> ColumnReference<R>.into(column: ColumnAccessor<R>): Boolean = into(column.name())
     public infix fun <R> ColumnReference<R>.into(column: KProperty<R>): Boolean = into(column.name)
+
+    public operator fun String.invoke(body: AddDsl<T>.() -> Unit): Unit = group(this, body)
+    public infix fun AnyColumnGroupAccessor.from(body: AddDsl<T>.() -> Unit): Unit = group(this, body)
+
+    public fun group(column: AnyColumnGroupAccessor, body: AddDsl<T>.() -> Unit): Unit = group(column.name(), body)
+    public fun group(name: String, body: AddDsl<T>.() -> Unit) {
+        val dsl = AddDsl(df)
+        body(dsl)
+        add(dsl.columns.toColumnGroup(name))
+    }
+
+    public fun group(body: AddDsl<T>.() -> Unit): AddGroup<T> = AddGroup(body)
+
+    public infix fun AddGroup<T>.into(groupName: String): Unit = group(groupName, body)
+    public infix fun AddGroup<T>.into(column: AnyColumnGroupAccessor): Unit = into(column.name())
 }
+
+public data class AddGroup<T>(internal val body: AddDsl<T>.() -> Unit)

src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt

@@ -777,8 +777,10 @@ class Modify : TestBase() {
         df.add {
             "year of birth" from 2021 - age
             age gt 18 into "is adult"
-            name.lastName.length() into "last name length"
-            "full name" from { name.firstName + " " + name.lastName }
+            "details" {
+                name.lastName.length() into "last name length"
+                "full name" from { name.firstName + " " + name.lastName }
+            }
         }
         // SampleEnd
     }
@@ -792,14 +794,17 @@ class Modify : TestBase() {
         val isAdult = column<Boolean>("is adult")
         val fullName = column<String>("full name")
         val name by columnGroup()
+        val details by columnGroup()
         val firstName by name.column<String>()
         val lastName by name.column<String>()
 
         df.add {
             yob from 2021 - age
             age gt 18 into isAdult
-            lastName.length() into lastNameLength
-            fullName from { firstName() + " " + lastName() }
+            details from {
+                lastName.length() into lastNameLength
+                fullName from { firstName() + " " + lastName() }
+            }
         }
         // SampleEnd
     }
@@ -810,8 +815,10 @@ class Modify : TestBase() {
         df.add {
             "year of birth" from 2021 - "age"<Int>()
             "age"<Int>() gt 18 into "is adult"
-            "name"["lastName"]<String>().length() into "last name length"
-            "full name" from { "name"["firstName"]<String>() + " " + "name"["lastName"]<String>() }
+            "details" {
+                "name"["lastName"]<String>().length() into "last name length"
+                "full name" from { "name"["firstName"]<String>() + " " + "name"["lastName"]<String>() }
+            }
         }
         // SampleEnd
     }

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

@@ -63,6 +63,7 @@ import org.jetbrains.kotlinx.dataframe.api.frameColumn
 import org.jetbrains.kotlinx.dataframe.api.gather
 import org.jetbrains.kotlinx.dataframe.api.getColumn
 import org.jetbrains.kotlinx.dataframe.api.getColumnGroup
+import org.jetbrains.kotlinx.dataframe.api.getColumns
 import org.jetbrains.kotlinx.dataframe.api.getFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.getValue
 import org.jetbrains.kotlinx.dataframe.api.group
@@ -883,14 +884,32 @@ class DataFrameTests : BaseTest() {
     fun `add several columns`() {
         val now = 2020
         val expected = typed.rows().map { now - it.age }
-
-        fun AnyFrame.check() = (1..3).forEach { this["year$it"].toList() shouldBe expected }
-
-        typed.add {
-            "year1" from { now - age }
-            "year2" from now - age
-            now - age into "year3"
-        }.check()
+        val g by columnGroup()
+
+        val df = typed.add {
+            "a" from { now - age }
+            "b" from now - age
+            now - age into "c"
+            "d" {
+                "f" from { now - age }
+            }
+            group {
+                g from {
+                    add(age.map { now - it }.named("h"))
+                }
+            } into "e"
+        }.remove { allBefore("a") }
+
+        df.columnNames() shouldBe listOf("a", "b", "c", "d", "e")
+        df["d"].kind() shouldBe ColumnKind.Group
+        df["e"].kind() shouldBe ColumnKind.Group
+        df.getColumnGroup("d").columnNames() shouldBe listOf("f")
+        df.getColumnGroup("e").getColumnGroup("g").columnNames() shouldBe listOf("h")
+        val cols = df.getColumns { allDfs() }
+        cols.size shouldBe 5
+        cols.forEach {
+            it.toList() shouldBe expected
+        }
     }
 
     @Test