[Compiler plugin] Support valueCounts operation

Author: koperagen

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

@@ -5,6 +5,8 @@ import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.nameGenerator
 import kotlin.reflect.KProperty
@@ -50,6 +52,8 @@ public fun <T> DataColumn<T>.valueCounts(
 
 // region DataFrame
 
+@Refine
+@Interpretable("ValueCounts")
 public fun <T> DataFrame<T>.valueCounts(
     sort: Boolean = true,
     ascending: Boolean = false,

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

@@ -78,3 +78,8 @@ internal fun <T> AbstractInterpreter<T>.dsl(
 ): ExpectedArgumentProvider<(Any, Map<String, Interpreter.Success<Any?>>) -> Unit> =
     arg(name, lens = Interpreter.Dsl, defaultValue = Present(value = {_, _ -> }))
 
+internal fun <T> AbstractInterpreter<T>.ignore(
+    name: ArgumentName? = null
+): ExpectedArgumentProvider<Nothing?> =
+    arg(name, lens = Interpreter.Id, defaultValue = Present(null))
+

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

@@ -0,0 +1,29 @@
+@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.ValueCount
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
+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.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
+
+class ValueCounts : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.dropNA by ignore()
+    val Arguments.ascending by ignore()
+    val Arguments.sort by ignore()
+    val Arguments.resultColumn: String by arg(defaultValue = Present(ValueCount::count.name))
+    val Arguments.columns: ColumnsResolver? by arg(defaultValue = Present(null))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val res = columns?.resolve(receiver)?.map { it.column } ?: receiver.columns()
+        val generator = org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator(res.map { it.name })
+        val count = SimpleDataColumn(generator.addUnique(resultColumn), session.builtinTypes.intType.type.wrap())
+        return PluginDataFrameSchema(res + count)
+    }
+}

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

@@ -100,6 +100,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ValueCounts
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
@@ -248,6 +249,7 @@ internal inline fun <reified T> String.load(): T {
         "Update0" -> Update0()
         "Aggregate" -> Aggregate()
         "DataFrameOf3" -> DataFrameOf3()
+        "ValueCounts" -> ValueCounts()
         else -> error("$this")
     } as T
 }

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

@@ -0,0 +1,16 @@
+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", "1", "2", "3", "1", "4", "1"))
+    val counts = df.valueCounts(ascending = true, sort = false, dropNA = false) { a }
+    counts.a
+    counts.count
+
+    val counts1 = df.valueCounts(resultColumn = "a") { a }
+    val col: DataColumn<String> = counts1.a
+    val col1: DataColumn<Int> counts1.a1
+    return "OK"
+}

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

@@ -514,6 +514,12 @@ public void testUpdate() {
     runTest("testData/box/update.kt");
   }
 
+  @Test
+  @TestMetadata("valueCounts.kt")
+  public void testValueCounts() {
+    runTest("testData/box/valueCounts.kt");
+  }
+
   @Test
   @TestMetadata("wrongReceiver.kt")
   public void testWrongReceiver() {