[Compiler plugin] Partially support Update operation

Something like update { }.where { }.with { } still would return an empty schema, and supporting it requires a more thorough approach with fallback schemas

Author: koperagen

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

@@ -155,6 +155,7 @@ private interface UpdateWithNote
  * @include [SelectingColumns.Dsl.WithExample] {@include [SetSelectingColumnsOperationArg]}
  * @include [Update.DslParam]
  */
+@Interpretable("Update0")
 public fun <T, C> DataFrame<T>.update(columns: ColumnsSelector<T, C>): Update<T, C> = Update(this, null, columns)
 
 /**

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

@@ -16,13 +16,16 @@ class FillNulls0 : AbstractInterpreter<FillNullsApproximation>() {
     }
 }
 
-class FillNullsApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver)
+class FillNullsApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) : UpdateApproximation
 
 class UpdateWith0 : AbstractSchemaModificationInterpreter() {
-    val Arguments.receiver: FillNullsApproximation by arg()
+    val Arguments.receiver: UpdateApproximation by arg()
     val Arguments.expression: TypeApproximation by type()
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
-        return convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
+        return when (val receiver = receiver) {
+            is FillNullsApproximation -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
+            is UpdateApproximationImpl -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
+        }
     }
 }

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

@@ -0,0 +1,19 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+
+class Update0 : AbstractInterpreter<UpdateApproximationImpl>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): UpdateApproximationImpl {
+        return UpdateApproximationImpl(receiver, columns)
+    }
+}
+
+sealed interface UpdateApproximation
+
+class UpdateApproximationImpl(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) : UpdateApproximation

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

@@ -88,6 +88,7 @@ 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.ToDataFrameFrom
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
@@ -204,6 +205,7 @@ internal inline fun <reified T> String.load(): T {
         "MapToFrame" -> MapToFrame()
         "Move0" -> Move0()
         "ToTop" -> ToTop()
+        "Update0" -> Update0()
         else -> error("$this")
     } as T
 }

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

@@ -0,0 +1,11 @@
+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", "b")(1, null, null, "")
+    val df1 = df.update { b }.with { "empty" }
+    val b: DataColumn<String> = df1.b
+    return "OK"
+}

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

@@ -418,6 +418,12 @@ public void testUngroup() {
     runTest("testData/box/ungroup.kt");
   }
 
+  @Test
+  @TestMetadata("update.kt")
+  public void testUpdate() {
+    runTest("testData/box/update.kt");
+  }
+
   @Nested
   @TestMetadata("testData/box/colKinds")
   @TestDataPath("$PROJECT_ROOT")