[Compiler plugin] DSL for creating nested column groups in toDataFrame

Author: koperagen

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

@@ -83,6 +83,7 @@ public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T>
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)
 
+@AccessApiOverload
 public fun <T> InsertClause<T>.after(column: ColumnAccessor<*>): DataFrame<T> = after(column.path())
 
 @AccessApiOverload

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

@@ -202,6 +202,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {
 
     public inline fun <reified R> inferType(noinline expression: (T) -> R): InferType<T, R> = InferType(expression)
 
+    @Interpretable("ToDataFrameDslStringInvoke")
     public abstract operator fun String.invoke(builder: CreateDataFrameDsl<T>.() -> Unit)
 }
 

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

@@ -131,6 +131,18 @@ class Properties0 : AbstractInterpreter<Unit>() {
     }
 }
 
+class ToDataFrameDslStringInvoke : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.builder by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = CreateDataFrameDslImplApproximation()
+        builder(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
+    }
+}
+
 class CreateDataFrameConfiguration {
     var maxDepth = DEFAULT_MAX_DEPTH
     var traverseConfiguration: TraverseConfiguration = TraverseConfiguration()

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

@@ -103,6 +103,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
@@ -242,6 +243,7 @@ internal inline fun <reified T> String.load(): T {
         "toDataFrameDsl" -> ToDataFrameDsl()
         "toDataFrame" -> ToDataFrame()
         "toDataFrameDefault" -> ToDataFrameDefault()
+        "ToDataFrameDslStringInvoke" -> ToDataFrameDslStringInvoke()
         "DataFrameOf0" -> DataFrameOf0()
         "DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
         "ToDataFrameColumn" -> ToDataFrameColumn()

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

@@ -0,0 +1,19 @@
+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 = listOf("a").toDataFrame {
+        "group" {
+            "anotherGroup" {
+                "a" from { 1 }
+            }
+            "b" from { "c" }
+        }
+    }
+
+    df.group.b
+    df.group.anotherGroup.a
+    return "OK"
+}

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

@@ -508,6 +508,12 @@ public void testToDataFrame_from() {
     runTest("testData/box/toDataFrame_from.kt");
   }
 
+  @Test
+  @TestMetadata("toDataFrame_nested.kt")
+  public void testToDataFrame_nested() {
+    runTest("testData/box/toDataFrame_nested.kt");
+  }
+
   @Test
   @TestMetadata("toDataFrame_nullableList.kt")
   public void testToDataFrame_nullableList() {