[Compiler plugin] Support GroupBy.add

Kotlin · Feb 3, 2025 · 16f5d51 · 16f5d51
1 parent 2143dfc
commit 16f5d51
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 0 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
@@ -248,6 +248,8 @@ public fun <T> DataFrame<T>.add(body: AddDsl<T>.() -> Unit): DataFrame<T> {
     return dataFrameOf(this@add.columns() + dsl.columns).cast()
 }
 
+@Refine
+@Interpretable("GroupByAdd")
 public inline fun <reified R, T, G> GroupBy<T, G>.add(
     name: String,
     infer: Infer = Infer.Nulls,

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt
@@ -27,6 +27,10 @@ data class PluginDataFrameSchema(
     }
 }
 
+fun PluginDataFrameSchema.add(name: String, type: ConeKotlinType, context: KotlinTypeFacade): PluginDataFrameSchema {
+    return PluginDataFrameSchema(columns() + context.simpleColumnOf(name, type))
+}
+
 private fun List<SimpleCol>.asString(indent: String = ""): String {
     return joinToString("\n") {
         val col = when (it) {

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt
@@ -19,6 +19,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.add
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
@@ -156,3 +157,13 @@ class GroupByToDataFrame : AbstractSchemaModificationInterpreter() {
         )
     }
 }
+
+class GroupByAdd : AbstractInterpreter<GroupBy>() {
+    val Arguments.receiver: GroupBy by groupBy()
+    val Arguments.name: String by arg()
+    val Arguments.type: TypeApproximation by type(name("expression"))
+
+    override fun Arguments.interpret(): GroupBy {
+        return GroupBy(receiver.keys, receiver.groups.add(name, type.type, context = this))
+    }
+}
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -88,6 +88,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.GroupByAdd
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MoveAfter0
@@ -275,6 +276,7 @@ internal inline fun <reified T> String.load(): T {
         "MoveToLeft1" -> MoveToLeft1()
         "MoveToRight0" -> MoveToRight0()
         "MoveAfter0" -> MoveAfter0()
+        "GroupByAdd" -> GroupByAdd()
         else -> error("$this")
     } as T
 }
diff --git a/plugins/kotlin-dataframe/testData/box/groupByAdd.kt b/plugins/kotlin-dataframe/testData/box/groupByAdd.kt
@@ -0,0 +1,42 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.api.groupBy
+import org.jetbrains.kotlinx.dataframe.io.*
+
+enum class State {
+    Idle,
+    Productive,
+    Maintenance,
+}
+
+class Event(val toolId: String, val state: State, val timestamp: Long)
+
+fun box(): String {
+    val tool1 = "tool_1"
+    val tool2 = "tool_2"
+    val tool3 = "tool_3"
+
+    val events = listOf(
+        Event(tool1, State.Idle, 0),
+        Event(tool1, State.Productive, 5),
+        Event(tool2, State.Idle, 0),
+        Event(tool2, State.Maintenance, 10),
+        Event(tool2, State.Idle, 20),
+        Event(tool3, State.Idle, 0),
+        Event(tool3, State.Productive, 25),
+    ).toDataFrame()
+
+    val lastTimestamp = events.maxOf { timestamp }
+    val groupBy = events
+        .groupBy { toolId }
+        .sortBy { timestamp }
+        .add("stateDuration") {
+            (next()?.timestamp ?: lastTimestamp) - timestamp
+        }.toDataFrame()
+
+    groupBy.group[0].stateDuration
+
+    groupBy.compareSchemas(strict = true)
+    return "OK"
+}
diff --git a/...e/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/...e/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
@@ -214,6 +214,12 @@ public void testGroupBy() {
     runTest("testData/box/groupBy.kt");
   }
 
+  @Test
+  @TestMetadata("groupByAdd.kt")
+  public void testGroupByAdd() {
+    runTest("testData/box/groupByAdd.kt");
+  }
+
   @Test
   @TestMetadata("groupBy_DataRow.kt")
   public void testGroupBy_DataRow() {