Implicit receiver fix

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/aggregation/AggregateGroupedDsl.kt

@@ -1,3 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.aggregation
 
+import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
+
+@HasSchema(schemaArg = 0)
 public abstract class AggregateGroupedDsl<out T> : AggregateDsl<T>()

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

@@ -435,6 +435,8 @@ private interface CommonDropNAFunctionDoc
  * @include [DropNA.WhereAllNAParam]
  * @include [DropDslParam]
  */
+@Refine
+@Interpretable("DropNa0")
 public fun <T> DataFrame<T>.dropNA(whereAllNA: Boolean = false, columns: ColumnsSelector<T, *>): DataFrame<T> {
     val cols = this[columns]
     return if (whereAllNA) {

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dropNulls.kt → plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/Nulls.kt

@@ -9,10 +9,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
 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.Present
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
-import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 
 class DropNulls0 : AbstractSchemaModificationInterpreter() {
@@ -24,6 +24,17 @@ class DropNulls0 : AbstractSchemaModificationInterpreter() {
     }
 }
 
+class DropNa0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.whereAllNA: Boolean by arg(defaultValue = Present(false))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        if (whereAllNA) return receiver
+        return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
+    }
+}
+
 fun KotlinTypeFacade.fillNullsImpl(
     columns: List<SimpleCol>,
     paths: Set<List<String>>,

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 
 internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
@@ -25,6 +26,7 @@ internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
 }
 
 internal class Expr0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     val Arguments.name: String by arg(defaultValue = Present("untitled"))
     val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
     val Arguments.expression: TypeApproximation by type()
@@ -53,6 +55,7 @@ internal class And0 : AbstractInterpreter<ColumnsResolver>() {
 }
 
 internal class All0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     override fun Arguments.interpret(): ColumnsResolver {
         return object : ColumnsResolver {
             override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
@@ -66,6 +69,7 @@ internal class All0 : AbstractInterpreter<ColumnsResolver>() {
 }
 
 internal class ColsOf0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     val Arguments.typeArg0: TypeApproximation by arg()
 
     override fun Arguments.interpret(): ColumnsResolver {
@@ -90,6 +94,7 @@ private fun Arguments.colsOf(cols: List<ColumnWithPathApproximation>, type: Cone
         }
 
 internal class ColsAtAnyDepth0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     override fun Arguments.interpret(): ColumnsResolver {
         return object : ColumnsResolver {
             override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {

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

@@ -505,7 +505,7 @@ internal fun FirFunctionCall.collectArgumentExpressions(): RefinedArguments {
     val refinedArgument = mutableListOf<RefinedArgument>()
 
     val parameterName = Name.identifier("receiver")
-    explicitReceiver?.let {
+    (explicitReceiver ?: extensionReceiver)?.let {
         if (it is FirResolvedQualifier && it.resolvedToCompanionObject) {
             return@let
         }

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

@@ -83,6 +83,7 @@ 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.DataFrameOf3
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataRowReadJsonStr
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DropNa0
 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
@@ -223,6 +224,7 @@ internal inline fun <reified T> String.load(): T {
         "Into0" -> Into0()
         "Ungroup0" -> Ungroup0()
         "DropNulls0" -> DropNulls0()
+        "DropNa0" -> DropNa0()
         "Properties0" -> Properties0()
         "Preserve0" -> Preserve0()
         "Preserve1" -> Preserve1()

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

@@ -0,0 +1,17 @@
+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" to listOf(1, null, 3),
+        "b" to listOf(null, 5, 6)
+    )
+    val df1 = df.dropNA { a and b }
+    df1.compareSchemas(strict = true)
+
+    val df2 = df.dropNA(whereAllNA = true) { a and b }
+    df2.compareSchemas(strict = true)
+    return "OK"
+}

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

@@ -0,0 +1,26 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.io.*
+
+data class Name(val firstName: String, val lastName: String)
+
+data class Score(val subject: String, val value: Int)
+
+data class Student(val name: Name, val age: Int, val scores: List<Score>)
+
+fun box(): String {
+    val students = listOf(
+        Student(Name("Alice", "Cooper"), 15, listOf(Score("math", 4), Score("biology", 3))),
+        Student(Name("Bob", "Marley"), 20, listOf(Score("music", 5))),
+    )
+
+    val df = students.toDataFrame().groupBy { expr { name.firstName} }
+        .aggregate {
+            remove { age } into "a"
+        }
+
+    df.compareSchemas(strict = true)
+    return "OK"
+}

plugins/kotlin-dataframe/testData/box/schemaFromImplicitReceiver.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.*
+
+data class Nested(val d: Double)
+
+data class Record(val a: String, val b: Int, val nested: Nested)
+
+fun box(): String {
+    val df = dataFrameOf("a", "b", "c")(1, 2, 3)
+
+    df.groupBy { a }
+        .updateGroups { remove { a } }
+        .aggregate { c into "c" }
+    return "OK"
+}
+
+

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

@@ -0,0 +1,23 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+interface JoinLeaf {
+    val something: Int
+    val somethingElse: String
+}
+
+@DataSchema
+interface Join2 {
+    val c: DataRow<JoinLeaf>
+}
+
+fun selectionDsl(df: DataFrame<Join2>) {
+    df.ungroup { c }.select { colsOf<String>() }.somethingElse
+}
+
+fun box(): String {
+    return "OK"
+}

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

@@ -130,6 +130,12 @@ public void testDiff() {
     runTest("testData/box/diff.kt");
   }
 
+  @Test
+  @TestMetadata("dropNA.kt")
+  public void testDropNA() {
+    runTest("testData/box/dropNA.kt");
+  }
+
   @Test
   @TestMetadata("dropNulls.kt")
   public void testDropNulls() {
@@ -286,6 +292,12 @@ public void testMapToFrame() {
     runTest("testData/box/mapToFrame.kt");
   }
 
+  @Test
+  @TestMetadata("modifySchemaInAggregate.kt")
+  public void testModifySchemaInAggregate() {
+    runTest("testData/box/modifySchemaInAggregate.kt");
+  }
+
   @Test
   @TestMetadata("moveAfter.kt")
   public void testMoveAfter() {
@@ -448,12 +460,24 @@ public void testSchema() {
     runTest("testData/box/Schema.kt");
   }
 
+  @Test
+  @TestMetadata("schemaFromImplicitReceiver.kt")
+  public void testSchemaFromImplicitReceiver() {
+    runTest("testData/box/schemaFromImplicitReceiver.kt");
+  }
+
   @Test
   @TestMetadata("select.kt")
   public void testSelect() {
     runTest("testData/box/select.kt");
   }
 
+  @Test
+  @TestMetadata("selectColsOf.kt")
+  public void testSelectColsOf() {
+    runTest("testData/box/selectColsOf.kt");
+  }
+
   @Test
   @TestMetadata("selectIt.kt")
   public void testSelectIt() {