[Compiler plugin] fillNulls { }.with { }

`with` used to have C? in UpdateExpression return type position, and so it was always inferred as nullable. Even for fillNulls { }.with { 123 }

Author: koperagen

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

@@ -69,6 +69,7 @@ private interface CommonFillNullsFunctionDoc
  * @include [SelectingColumns.Dsl.WithExample] {@include [SetFillNullsOperationArg]}
  * @include [Update.DslParam]
  */
+@Interpretable("FillNulls0")
 public fun <T, C> DataFrame<T>.fillNulls(columns: ColumnsSelector<T, C?>): Update<T, C?> =
     update(columns).where { it == null }
 

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

@@ -8,6 +8,8 @@ import org.jetbrains.kotlinx.dataframe.DataFrameExpression
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowColumnExpression
 import org.jetbrains.kotlinx.dataframe.RowValueFilter
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.api.Update.Grammar
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -273,7 +275,9 @@ public typealias UpdateExpression<T, C, R> = AddDataRow<T>.(C) -> R
  * - {@include [SeeAlsoPerRowCol]}
  * @param [expression] The {@include [ExpressionsGivenRow.RowValueExpressionLink]} to update the rows with.
  */
-public fun <T, C> Update<T, C>.with(expression: UpdateExpression<T, C, C?>): DataFrame<T> =
+@Refine
+@Interpretable("UpdateWith0")
+public fun <T, C, R : C?> Update<T, C>.with(expression: UpdateExpression<T, C, R>): DataFrame<T> =
     updateImpl { row, _, value ->
         expression(row, value)
     }

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

@@ -0,0 +1,28 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
+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.type
+
+class FillNulls0 : AbstractInterpreter<FillNullsApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): FillNullsApproximation {
+        return FillNullsApproximation(receiver, columns)
+    }
+}
+
+class FillNullsApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver)
+
+class UpdateWith0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: FillNullsApproximation 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)
+    }
+}

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

@@ -73,13 +73,15 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
 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
 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.UpdateWith0
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
     val symbol =
@@ -186,6 +188,8 @@ internal inline fun <reified T> String.load(): T {
         "ToDataFrameColumn" -> ToDataFrameColumn()
         "StringColumns" -> ToDataFrameColumn()
         "ReadExcel" -> ReadExcel()
+        "FillNulls0" -> FillNulls0()
+        "UpdateWith0" -> UpdateWith0()
         else -> error("$this")
     } as T
 }

plugins/kotlin-dataframe/testData/box/fillNulls.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.fillNulls { b }.with { "empty" }
+    val b: DataColumn<String> = df1.b
+    return "OK"
+}

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

@@ -136,6 +136,12 @@ public void testExtractPluginSchemaWithUnfold() {
     runTest("testData/box/extractPluginSchemaWithUnfold.kt");
   }
 
+  @Test
+  @TestMetadata("fillNulls.kt")
+  public void testFillNulls() {
+    runTest("testData/box/fillNulls.kt");
+  }
+
   @Test
   @TestMetadata("flexibleReturnType.kt")
   public void testFlexibleReturnType() {