[Compiler plugin] Support creating groups in AddDsl

Author: koperagen

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

@@ -207,6 +207,7 @@ public class AddDsl<T>(
 
     public infix fun <R> ColumnReference<R>.into(column: KProperty<R>): Boolean = into(column.name)
 
+    @Interpretable("AddDslStringInvoke")
     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)

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt

@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
@@ -57,3 +58,15 @@ class AddWithDsl : AbstractSchemaModificationInterpreter() {
         return PluginDataFrameSchema(addDsl.columns)
     }
 }
+
+class AddDslStringInvoke : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg(lens = Interpreter.Value)
+    val Arguments.receiver: String by string()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
+    }
+}

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt

@@ -67,6 +67,7 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
@@ -196,6 +197,7 @@ internal inline fun <reified T> String.load(): T {
         "Flatten0" -> Flatten0()
         "FlattenDefault" -> FlattenDefault()
         "AddId" -> AddId()
+        "AddDslStringInvoke" -> AddDslStringInvoke()
         else -> error("$this")
     } as T
 }

plugins/kotlin-dataframe/testData/box/addDsl.kt

@@ -0,0 +1,16 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("a")(1).add {
+        "id" from { it }
+        "group" {
+            "a" from { it }
+        }
+    }
+
+    df.group.a
+    return "OK"
+}

plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

@@ -17,6 +17,12 @@
 @TestMetadata("testData/box")
 @TestDataPath("$PROJECT_ROOT")
 public class DataFrameBlackBoxCodegenTestGenerated extends AbstractDataFrameBlackBoxCodegenTest {
+  @Test
+  @TestMetadata("addDsl.kt")
+  public void testAddDsl() {
+    runTest("testData/box/addDsl.kt");
+  }
+
   @Test
   @TestMetadata("addId.kt")
   public void testAddId() {