[Compiler plugin] Support mapToFrame

Author: koperagen

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

@@ -7,6 +7,8 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.Selector
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.jetbrains.kotlinx.dataframe.impl.columnName
@@ -102,6 +104,8 @@ public fun <T, R> ColumnsContainer<T>.mapToColumn(
     body: AddExpression<T, R>,
 ): DataColumn<R> = mapToColumn(column.columnName, type, infer, body)
 
+@Refine
+@Interpretable("MapToFrame")
 public fun <T> DataFrame<T>.mapToFrame(body: AddDsl<T>.() -> Unit): AnyFrame {
     val dsl = AddDsl(this)
     body(dsl)

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

@@ -0,0 +1,18 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
+
+class MapToFrame : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        return PluginDataFrameSchema(addDsl.columns)
+    }
+}

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

@@ -79,6 +79,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
@@ -198,6 +199,7 @@ internal inline fun <reified T> String.load(): T {
         "FlattenDefault" -> FlattenDefault()
         "AddId" -> AddId()
         "AddDslStringInvoke" -> AddDslStringInvoke()
+        "MapToFrame" -> MapToFrame()
         else -> error("$this")
     } as T
 }

plugins/kotlin-dataframe/testData/box/mapToFrame.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 = dataFrameOf("a")(1).mapToFrame {
+        "id" from { it }
+        "group" {
+            "a" from { it }
+        }
+    }
+
+    df.id
+    df.group.a
+
+    df.compareSchemas(strict = true)
+    return "OK"
+}

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

@@ -244,6 +244,12 @@ public void testMain() {
     runTest("testData/box/main.kt");
   }
 
+  @Test
+  @TestMetadata("mapToFrame.kt")
+  public void testMapToFrame() {
+    runTest("testData/box/mapToFrame.kt");
+  }
+
   @Test
   @TestMetadata("nestedDataSchemaCodegen.kt")
   public void testNestedDataSchemaCodegen() {