Skip to content

Commit

Permalink
Merge pull request #951 from Kotlin/valueCounts
Browse files Browse the repository at this point in the history
[Compiler plugin] Support valueCounts
  • Loading branch information
koperagen authored Nov 20, 2024
2 parents 741854c + 24ddca2 commit ec60fd4
Show file tree
Hide file tree
Showing 13 changed files with 93 additions and 93 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,14 +1,8 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl

import org.jetbrains.kotlinx.dataframe.plugin.impl.api.RenameClauseApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter.*
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnAccessorApproximation
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.data.DataFrameCallableId
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.InsertClauseApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.KPropertyApproximation
import kotlin.properties.PropertyDelegateProvider
import kotlin.properties.ReadOnlyProperty

Expand All @@ -18,63 +12,26 @@ fun <T> AbstractInterpreter<T>.dataFrame(
name: ArgumentName? = null
): ExpectedArgumentProvider<PluginDataFrameSchema> = arg(name, lens = Interpreter.Schema)

fun <T> AbstractInterpreter<T>.varargString(
name: ArgumentName? = null,
defaultValue: DefaultValue<List<String>> = Absent
): ExpectedArgumentProvider<List<String>> = arg(name, lens = Interpreter.Value, defaultValue = defaultValue)

fun <T> AbstractInterpreter<T>.renameClause(
name: ArgumentName? = null
): ExpectedArgumentProvider<RenameClauseApproximation> = arg(name, lens = Interpreter.Value)

fun <T> AbstractInterpreter<T>.columnsSelector(
name: ArgumentName? = null
): ExpectedArgumentProvider<List<String>> = arg(name, lens = Interpreter.Value)

fun <T> AbstractInterpreter<T>.type(
name: ArgumentName? = null
): ExpectedArgumentProvider<TypeApproximation> = arg(name, lens = Interpreter.ReturnType)

fun <T, E : Enum<E>> AbstractInterpreter<T>.enum(
name: ArgumentName? = null,
defaultValue: DefaultValue<E> = Absent
): ExpectedArgumentProvider<E> = argConvert(name = name, lens = Interpreter.Value, defaultValue = defaultValue) { it: DataFrameCallableId ->
): ExpectedArgumentProvider<E> = argConvert(name = name, defaultValue = defaultValue) { it: DataFrameCallableId ->
val forName: Class<*> = Class.forName("${it.packageName}.${it.className}")
@Suppress("UNCHECKED_CAST")
java.lang.Enum.valueOf(forName as Class<out Enum<*>>, it.callableName) as E
}

fun <T> AbstractInterpreter<T>.columnAccessor(
name: ArgumentName? = null
): ExpectedArgumentProvider<ColumnAccessorApproximation> = arg(name, lens = Interpreter.Value)

fun <T> AbstractInterpreter<T>.dataColumn(
name: ArgumentName? = null
): ExpectedArgumentProvider<SimpleCol> = arg(name, lens = Interpreter.Value)

fun <T> AbstractInterpreter<T>.insertClause(
name: ArgumentName? = null
): ExpectedArgumentProvider<InsertClauseApproximation> = arg(name, lens = Interpreter.Value)

internal fun <T> AbstractInterpreter<T>.columnPath(
name: ArgumentName? = null
): ExpectedArgumentProvider<ColumnPathApproximation> = arg(name, lens = Interpreter.Value)

internal fun <T> AbstractInterpreter<T>.columnWithPath(
name: ArgumentName? = null
): ExpectedArgumentProvider<ColumnWithPathApproximation> = arg(name, lens = Interpreter.Value)

fun <T> AbstractInterpreter<T>.kproperty(
name: ArgumentName? = null
): ExpectedArgumentProvider<KPropertyApproximation> = arg(name, lens = Interpreter.Value)

internal fun <T> AbstractInterpreter<T>.string(
name: ArgumentName? = null
): ExpectedArgumentProvider<String> =
arg(name, lens = Interpreter.Value)

internal fun <T> AbstractInterpreter<T>.dsl(
name: ArgumentName? = null
): 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))

Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,19 @@ import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
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.Interpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
import org.jetbrains.kotlinx.dataframe.plugin.impl.type

typealias TypeApproximation = Marker

class Add : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.name: String by string()
val Arguments.name: String by arg()
val Arguments.type: TypeApproximation by type(name("expression"))

override fun Arguments.interpret(): PluginDataFrameSchema {
Expand All @@ -27,8 +25,8 @@ class Add : AbstractSchemaModificationInterpreter() {
}

class From : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg(lens = Interpreter.Value)
val Arguments.receiver: String by string()
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: String by arg()
val Arguments.type: TypeApproximation by type(name("expression"))

override fun Arguments.interpret() {
Expand All @@ -37,9 +35,9 @@ class From : AbstractInterpreter<Unit>() {
}

class Into : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg(lens = Interpreter.Value)
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: TypeApproximation by type()
val Arguments.name: String by string()
val Arguments.name: String by arg()

override fun Arguments.interpret() {
dsl.columns += simpleColumnOf(name, receiver.type)
Expand All @@ -60,8 +58,8 @@ class AddWithDsl : AbstractSchemaModificationInterpreter() {
}

class AddDslStringInvoke : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg(lens = Interpreter.Value)
val Arguments.receiver: String by string()
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: String by arg()
val Arguments.body by dsl()

override fun Arguments.interpret() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api

import org.jetbrains.kotlinx.dataframe.api.Infer
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.impl.Absent
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
Expand All @@ -11,13 +12,10 @@ 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.SimpleFrameColumn
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.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString

internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {
val Arguments.columns: ColumnsResolver by arg()
Expand All @@ -31,7 +29,7 @@ internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {

class Convert2 : AbstractInterpreter<ConvertApproximation>() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: List<String> by varargString()
val Arguments.columns: List<String> by arg(defaultValue = Absent)

override fun Arguments.interpret(): ConvertApproximation {
return ConvertApproximation(receiver, columns.map { listOf(it) })
Expand All @@ -41,8 +39,8 @@ class Convert2 : AbstractInterpreter<ConvertApproximation>() {
class ConvertApproximation(val schema: PluginDataFrameSchema, val columns: List<List<String>>)

internal class Convert6 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.firstCol: String by string()
val Arguments.cols: List<String> by varargString(defaultValue = Present(emptyList()))
val Arguments.firstCol: String by arg()
val Arguments.cols: List<String> by arg(defaultValue = Present(emptyList()))
val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
val Arguments.expression: TypeApproximation by type()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,10 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
import org.jetbrains.kotlinx.dataframe.impl.api.withValuesImpl

class DataFrameOf0 : AbstractInterpreter<DataFrameBuilderApproximation>() {
val Arguments.header: List<String> by varargString()
val Arguments.header: List<String> by arg()

override fun Arguments.interpret(): DataFrameBuilderApproximation {
return DataFrameBuilderApproximation(header)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,27 +14,21 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.columnAccessor
import org.jetbrains.kotlinx.dataframe.plugin.impl.columnPath
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnAccessorApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.InsertClauseApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.KPropertyApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
import org.jetbrains.kotlinx.dataframe.plugin.impl.insertClause
import org.jetbrains.kotlinx.dataframe.plugin.impl.kproperty
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.type

/**
* @see DataFrame.insert
*/
internal class Insert0 : AbstractInterpreter<InsertClauseApproximation>() {
val Arguments.column: SimpleCol by dataColumn()
val Arguments.column: SimpleCol by arg()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()

override fun Arguments.interpret(): InsertClauseApproximation {
Expand All @@ -43,7 +37,7 @@ internal class Insert0 : AbstractInterpreter<InsertClauseApproximation>() {
}

internal class Insert1 : AbstractInterpreter<InsertClauseApproximation>() {
val Arguments.name: String by string()
val Arguments.name: String by arg()
val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
val Arguments.expression: TypeApproximation by type()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
Expand All @@ -54,7 +48,7 @@ internal class Insert1 : AbstractInterpreter<InsertClauseApproximation>() {
}

internal class Insert2 : AbstractInterpreter<InsertClauseApproximation>() {
val Arguments.column: ColumnAccessorApproximation by columnAccessor()
val Arguments.column: ColumnAccessorApproximation by arg()
val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
val Arguments.expression: TypeApproximation by type()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
Expand All @@ -65,7 +59,7 @@ internal class Insert2 : AbstractInterpreter<InsertClauseApproximation>() {
}

internal class Insert3 : AbstractInterpreter<InsertClauseApproximation>() {
val Arguments.column: KPropertyApproximation by kproperty()
val Arguments.column: KPropertyApproximation by arg()
val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
val Arguments.expression: TypeApproximation by type()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
Expand All @@ -77,7 +71,7 @@ internal class Insert3 : AbstractInterpreter<InsertClauseApproximation>() {

internal class Under0 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.column: SingleColumnApproximation by arg()
val Arguments.receiver: InsertClauseApproximation by insertClause()
val Arguments.receiver: InsertClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val path = column.resolve(receiver.df).single().path
Expand All @@ -88,8 +82,8 @@ internal class Under0 : AbstractInterpreter<PluginDataFrameSchema>() {
}

internal class Under1 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.columnPath: ColumnPathApproximation by columnPath()
val Arguments.receiver: InsertClauseApproximation by insertClause()
val Arguments.columnPath: ColumnPathApproximation by arg()
val Arguments.receiver: InsertClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame()
Expand All @@ -99,8 +93,8 @@ internal class Under1 : AbstractInterpreter<PluginDataFrameSchema>() {
}

internal class Under2 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.column: ColumnAccessorApproximation by columnAccessor()
val Arguments.receiver: InsertClauseApproximation by insertClause()
val Arguments.column: ColumnAccessorApproximation by arg()
val Arguments.receiver: InsertClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame()
Expand All @@ -110,8 +104,8 @@ internal class Under2 : AbstractInterpreter<PluginDataFrameSchema>() {
}

internal class Under3 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.column: KPropertyApproximation by kproperty()
val Arguments.receiver: InsertClauseApproximation by insertClause()
val Arguments.column: KPropertyApproximation by arg()
val Arguments.receiver: InsertClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame()
Expand All @@ -121,8 +115,8 @@ internal class Under3 : AbstractInterpreter<PluginDataFrameSchema>() {
}

internal class Under4 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.column: String by string()
val Arguments.receiver: InsertClauseApproximation by insertClause()
val Arguments.column: String by arg()
val Arguments.receiver: InsertClauseApproximation by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInt
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.IoSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.deserializeToPluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.toPluginDataFrameSchema
import java.io.File

internal class Read0 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.path by string()
val Arguments.path: String by arg()
val Arguments.header: List<String> by arg(defaultValue = Present(listOf()))

override fun Arguments.interpret(): PluginDataFrameSchema {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString

class Rename : AbstractInterpreter<RenameClauseApproximation>() {
private val Arguments.receiver by dataFrame()
Expand All @@ -24,7 +22,7 @@ class RenameClauseApproximation(val schema: PluginDataFrameSchema, val columns:

class RenameInto : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: RenameClauseApproximation by arg()
val Arguments.newNames: List<String> by varargString()
val Arguments.newNames: List<String> by arg()

override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = receiver.columns.resolve(receiver.schema)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,16 +22,16 @@ class PairConstructor : AbstractInterpreter<Pair<*, *>>() {
}

class TrimMargin : AbstractInterpreter<String>() {
val Arguments.receiver: String by arg(lens = Interpreter.Value)
val Arguments.marginPrefix: String by arg(lens = Interpreter.Value, defaultValue = Present("|"))
val Arguments.receiver: String by arg()
val Arguments.marginPrefix: String by arg(defaultValue = Present("|"))

override fun Arguments.interpret(): String {
return receiver.trimMargin(marginPrefix)
}
}

class TrimIndent : AbstractInterpreter<String>() {
val Arguments.receiver: String by arg(lens = Interpreter.Value)
val Arguments.receiver: String by arg()

override fun Arguments.interpret(): String {
return receiver.trimIndent()
Expand Down
Loading

0 comments on commit ec60fd4

Please sign in to comment.