Skip to content

Commit

Permalink
Annotate some of the API overloads
Browse files Browse the repository at this point in the history
  • Loading branch information
koperagen committed Dec 16, 2024
1 parent 291371c commit d88b609
Show file tree
Hide file tree
Showing 45 changed files with 290 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package org.jetbrains.kotlinx.dataframe

import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.api.ColumnSelectionDsl
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
import org.jetbrains.kotlinx.dataframe.api.cast
Expand Down Expand Up @@ -42,8 +43,10 @@ public interface ColumnsContainer<out T> {

public fun getColumnOrNull(index: Int): AnyCol?

@AccessApiOverload
public fun <R> getColumnOrNull(column: ColumnReference<R>): DataColumn<R>?

@AccessApiOverload
public fun <R> getColumnOrNull(column: KProperty<R>): DataColumn<R>?

public fun getColumnOrNull(path: ColumnPath): AnyCol?
Expand All @@ -58,23 +61,32 @@ public interface ColumnsContainer<out T> {

public operator fun get(columnPath: ColumnPath): AnyCol = getColumn(columnPath)

@AccessApiOverload
public operator fun <R> get(column: DataColumn<R>): DataColumn<R> = getColumn(column.name()).cast()

@AccessApiOverload
public operator fun <R> get(column: DataColumn<DataRow<R>>): ColumnGroup<R> = getColumn(column)

@AccessApiOverload
public operator fun <R> get(column: DataColumn<DataFrame<R>>): FrameColumn<R> = getColumn(column)

@AccessApiOverload
public operator fun <R> get(column: ColumnReference<R>): DataColumn<R> = getColumn(column)

@AccessApiOverload
public operator fun <R> get(column: ColumnReference<DataRow<R>>): ColumnGroup<R> = getColumn(column)

@AccessApiOverload
public operator fun <R> get(column: ColumnReference<DataFrame<R>>): FrameColumn<R> = getColumn(column)

@AccessApiOverload
public operator fun <R> get(column: KProperty<R>): DataColumn<R> = get(column.columnName).cast()

@AccessApiOverload
public operator fun <R> get(column: KProperty<DataRow<R>>): ColumnGroup<R> =
get(column.columnName).asColumnGroup().cast()

@AccessApiOverload
public operator fun <R> get(column: KProperty<DataFrame<R>>): FrameColumn<R> =
get(column.columnName).asAnyFrameColumn().castFrameColumn()

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe

import org.jetbrains.kotlinx.dataframe.aggregation.Aggregatable
import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedBody
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
import org.jetbrains.kotlinx.dataframe.api.add
Expand Down Expand Up @@ -116,6 +117,7 @@ public interface DataFrame<out T> :
*/
public operator fun <T, C> DataFrame<T>.get(columns: ColumnsSelector<T, C>): List<DataColumn<C>> = this.get(columns)

@AccessApiOverload
public operator fun <T> DataFrame<T>.get(first: AnyColumnReference, vararg other: AnyColumnReference): DataFrame<T> =
select { (listOf(first) + other).toColumnSet() }

Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package org.jetbrains.kotlinx.dataframe

import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.api.next
import org.jetbrains.kotlinx.dataframe.api.prev
import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
Expand Down Expand Up @@ -31,6 +32,7 @@ public interface DataRow<out T> {

public operator fun <R> get(columns: List<ColumnReference<R>>): List<R> = columns.map { get(it) }

@AccessApiOverload
public operator fun <R> get(property: KProperty<R>): R = get(property.columnName) as R

public operator fun get(first: AnyColumnReference, vararg other: AnyColumnReference): DataRow<T> =
Expand Down Expand Up @@ -66,6 +68,7 @@ public interface DataRow<out T> {

public fun getOrNull(name: String): Any?

@AccessApiOverload
public fun <R> getValueOrNull(column: ColumnReference<R>): R?

// endregion
Expand All @@ -74,6 +77,7 @@ public interface DataRow<out T> {

public operator fun String.get(vararg path: String): ColumnPath = ColumnPath(listOf(this) + path)

@AccessApiOverload
public operator fun <R> ColumnReference<R>.invoke(): R = get(this)

public operator fun <R> String.invoke(): R = this@DataRow[this@invoke] as R
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnsSelector
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
Expand Down Expand Up @@ -63,12 +64,15 @@ public fun <T> ColumnsContainer<T>.getColumnGroup(columnPath: ColumnPath): Colum
public fun <T> ColumnsContainer<T>.getColumn(name: String): AnyCol =
getColumnOrNull(name) ?: throw IllegalArgumentException("Column not found: '$name'")

@AccessApiOverload
public fun <T, R> ColumnsContainer<T>.getColumn(column: ColumnReference<DataFrame<R>>): FrameColumn<R> =
getColumnOrNull(column)?.asFrameColumn() ?: throw IllegalArgumentException("FrameColumn not found: '$column'")

@AccessApiOverload
public fun <T, R> ColumnsContainer<T>.getColumn(column: ColumnReference<DataRow<R>>): ColumnGroup<R> =
getColumnOrNull(column)?.asColumnGroup() ?: throw IllegalArgumentException("ColumnGroup not found: '$column'")

@AccessApiOverload
public fun <T, R> ColumnsContainer<T>.getColumn(column: ColumnReference<R>): DataColumn<R> =
getColumnOrNull(column) ?: throw IllegalArgumentException("Column not found: '$column'")

Expand All @@ -89,9 +93,11 @@ public fun <T> ColumnsContainer<T>.getColumnGroup(index: Int): ColumnGroup<*> =

public fun <T> ColumnsContainer<T>.getColumnGroup(name: String): ColumnGroup<*> = getColumn(name).asColumnGroup()

@AccessApiOverload
public fun <T> ColumnsContainer<T>.getColumnGroup(column: KProperty<*>): ColumnGroup<*> =
getColumnGroup(column.columnName)

@AccessApiOverload
public fun <T, C> ColumnsContainer<T>.getColumnGroup(column: ColumnReference<DataRow<C>>): ColumnGroup<C> =
getColumn(column)

Expand All @@ -105,19 +111,24 @@ public fun <T, C> ColumnsContainer<T>.getColumnGroup(column: ColumnSelector<T, D
public fun <T> ColumnsContainer<T>.getColumnGroupOrNull(name: String): ColumnGroup<*>? =
getColumnOrNull(name)?.asColumnGroup()

@AccessApiOverload
public fun <T> ColumnsContainer<T>.getColumnGroupOrNull(column: KProperty<*>): ColumnGroup<*>? =
getColumnGroupOrNull(column.columnName)

// endregion

// region containsColumn

@AccessApiOverload
public fun <C> ColumnsContainer<*>.containsColumn(column: ColumnReference<C>): Boolean = getColumnOrNull(column) != null

@AccessApiOverload
public fun ColumnsContainer<*>.containsColumn(column: KProperty<*>): Boolean = containsColumn(column.columnName)

@AccessApiOverload
public operator fun ColumnsContainer<*>.contains(column: AnyColumnReference): Boolean = containsColumn(column)

@AccessApiOverload
public operator fun ColumnsContainer<*>.contains(column: KProperty<*>): Boolean = containsColumn(column)

// region rows
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.RowExpression
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
import org.jetbrains.kotlinx.dataframe.impl.columnName
Expand Down Expand Up @@ -60,10 +61,12 @@ public fun <T> AnyRow.getValue(columnName: String): T = get(columnName) as T

public fun <T> AnyRow.getValue(column: ColumnReference<T>): T = get(column)

@AccessApiOverload
public fun <T> AnyRow.getValue(column: KProperty<T>): T = get(column)

public fun <T> AnyRow.getValueOrNull(columnName: String): T? = getOrNull(columnName) as T?

@AccessApiOverload
public fun <T> AnyRow.getValueOrNull(column: KProperty<T>): T? = getValueOrNull<T>(column.columnName)

// endregion
Expand All @@ -74,10 +77,12 @@ public fun AnyRow.containsKey(columnName: String): Boolean = owner.containsColum

public fun AnyRow.containsKey(column: AnyColumnReference): Boolean = owner.containsColumn(column)

@AccessApiOverload
public fun AnyRow.containsKey(column: KProperty<*>): Boolean = owner.containsColumn(column)

public operator fun AnyRow.contains(column: AnyColumnReference): Boolean = containsKey(column)

@AccessApiOverload
public operator fun AnyRow.contains(column: KProperty<*>): Boolean = containsKey(column)

// endregion
Expand Down
13 changes: 13 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/Nulls.kt
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.AnyRow
import org.jetbrains.kotlinx.dataframe.ColumnsSelector
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
import org.jetbrains.kotlinx.dataframe.annotations.Refine
import org.jetbrains.kotlinx.dataframe.api.DropNA.DropNASelectingOptions
Expand Down Expand Up @@ -85,6 +86,7 @@ public fun <T> DataFrame<T>.fillNulls(vararg columns: String): Update<T, Any?> =
* @include [SelectingColumns.KProperties.WithExample] {@include [SetFillNullsOperationArg]}
* @include [Update.KPropertiesParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNulls(vararg columns: KProperty<C>): Update<T, C?> =
fillNulls { columns.toColumnSet() }

Expand All @@ -93,6 +95,7 @@ public fun <T, C> DataFrame<T>.fillNulls(vararg columns: KProperty<C>): Update<T
* @include [SelectingColumns.ColumnAccessors.WithExample] {@include [SetFillNullsOperationArg]}
* @include [Update.ColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNulls(vararg columns: ColumnReference<C>): Update<T, C?> =
fillNulls { columns.toColumnSet() }

Expand Down Expand Up @@ -176,13 +179,15 @@ public fun <T> DataFrame<T>.fillNaNs(vararg columns: String): Update<T, Any?> =
* @include [SelectingColumns.KProperties.WithExample] {@include [SetFillNaNsOperationArg]}
* @include [Update.KPropertiesParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNaNs(vararg columns: KProperty<C>): Update<T, C> = fillNaNs { columns.toColumnSet() }

/**
* @include [CommonFillNaNsFunctionDoc]
* @include [SelectingColumns.ColumnAccessors.WithExample] {@include [SetFillNaNsOperationArg]}
* @include [Update.ColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNaNs(vararg columns: ColumnReference<C>): Update<T, C> =
fillNaNs { columns.toColumnSet() }

Expand Down Expand Up @@ -246,13 +251,15 @@ public fun <T> DataFrame<T>.fillNA(vararg columns: String): Update<T, Any?> = fi
* @include [SelectingColumns.KProperties.WithExample] {@include [SetFillNAOperationArg]}
* @include [Update.KPropertiesParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNA(vararg columns: KProperty<C>): Update<T, C?> = fillNA { columns.toColumnSet() }

/**
* @include [CommonFillNAFunctionDoc]
* @include [SelectingColumns.ColumnAccessors.WithExample] {@include [SetFillNAOperationArg]}
* @include [Update.ColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T, C> DataFrame<T>.fillNA(vararg columns: ColumnReference<C>): Update<T, C?> =
fillNA { columns.toColumnSet() }

Expand Down Expand Up @@ -348,6 +355,7 @@ public fun <T> DataFrame<T>.dropNulls(whereAllNull: Boolean = false): DataFrame<
* @include [DropNulls.WhereAllNullParam]
* @include [DropKPropertiesParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNulls(vararg columns: KProperty<*>, whereAllNull: Boolean = false): DataFrame<T> =
dropNulls(whereAllNull) { columns.toColumnSet() }

Expand All @@ -368,6 +376,7 @@ public fun <T> DataFrame<T>.dropNulls(vararg columns: String, whereAllNull: Bool
* @include [DropNulls.WhereAllNullParam]
* @include [DropColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNulls(vararg columns: AnyColumnReference, whereAllNull: Boolean = false): DataFrame<T> =
dropNulls(whereAllNull) { columns.toColumnSet() }

Expand Down Expand Up @@ -445,6 +454,7 @@ public fun <T> DataFrame<T>.dropNA(whereAllNA: Boolean = false, columns: Columns
* @include [DropNA.WhereAllNAParam]
* @include [DropKPropertiesParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNA(vararg columns: KProperty<*>, whereAllNA: Boolean = false): DataFrame<T> =
dropNA(whereAllNA) { columns.toColumnSet() }

Expand All @@ -465,6 +475,7 @@ public fun <T> DataFrame<T>.dropNA(vararg columns: String, whereAllNA: Boolean =
* @include [DropNA.WhereAllNAParam]
* @include [DropColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNA(vararg columns: AnyColumnReference, whereAllNA: Boolean = false): DataFrame<T> =
dropNA(whereAllNA) { columns.toColumnSet() }

Expand Down Expand Up @@ -552,6 +563,7 @@ public fun <T> DataFrame<T>.dropNaNs(whereAllNaN: Boolean = false, columns: Colu
* @include [DropNaNs.WhereAllNaNParam]
* @include [DropKPropertiesParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNaNs(vararg columns: KProperty<*>, whereAllNaN: Boolean = false): DataFrame<T> =
dropNaNs(whereAllNaN) { columns.toColumnSet() }

Expand All @@ -572,6 +584,7 @@ public fun <T> DataFrame<T>.dropNaNs(vararg columns: String, whereAllNaN: Boolea
* @include [DropNaNs.WhereAllNaNParam]
* @include [DropColumnAccessorsParam]
*/
@AccessApiOverload
public fun <T> DataFrame<T>.dropNaNs(vararg columns: AnyColumnReference, whereAllNaN: Boolean = false): DataFrame<T> =
dropNaNs(whereAllNaN) { columns.toColumnSet() }

Expand Down
13 changes: 13 additions & 0 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.RowExpression
import org.jetbrains.kotlinx.dataframe.Selector
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
import org.jetbrains.kotlinx.dataframe.annotations.Refine
import org.jetbrains.kotlinx.dataframe.columns.BaseColumn
Expand Down Expand Up @@ -131,12 +132,14 @@ public inline fun <reified R, T> DataFrame<T>.add(
noinline expression: AddExpression<T, R>,
): DataFrame<T> = (this + mapToColumn(name, infer, expression))

@AccessApiOverload
public inline fun <reified R, T> DataFrame<T>.add(
property: KProperty<R>,
infer: Infer = Infer.Nulls,
noinline expression: AddExpression<T, R>,
): DataFrame<T> = (this + mapToColumn(property, infer, expression))

@AccessApiOverload
public inline fun <reified R, T> DataFrame<T>.add(
column: ColumnAccessor<R>,
infer: Infer = Infer.Nulls,
Expand Down Expand Up @@ -188,30 +191,38 @@ public class AddDsl<T>(
add(this, Infer.Nulls, expression)

// TODO: use path instead of name
@AccessApiOverload
public inline infix fun <reified R> ColumnAccessor<R>.from(noinline expression: RowExpression<T, R>): Boolean =
name().from(expression)

@AccessApiOverload
public inline infix fun <reified R> KProperty<R>.from(noinline expression: RowExpression<T, R>): Boolean =
add(name, Infer.Nulls, expression)

public infix fun String.from(column: AnyColumnReference): Boolean = add(column.rename(this))

@AccessApiOverload
public inline infix fun <reified R> ColumnAccessor<R>.from(column: ColumnReference<R>): Boolean = name() from column

@AccessApiOverload
public inline infix fun <reified R> KProperty<R>.from(column: ColumnReference<R>): Boolean = name from column

@Interpretable("Into")
public infix fun AnyColumnReference.into(name: String): Boolean = add(rename(name))

@AccessApiOverload
public infix fun <R> ColumnReference<R>.into(column: ColumnAccessor<R>): Boolean = into(column.name())

@AccessApiOverload
public infix fun <R> ColumnReference<R>.into(column: KProperty<R>): Boolean = into(column.name)

@Interpretable("AddDslStringInvoke")
public operator fun String.invoke(body: AddDsl<T>.() -> Unit): Unit = group(this, body)

@AccessApiOverload
public infix fun AnyColumnGroupAccessor.from(body: AddDsl<T>.() -> Unit): Unit = group(this, body)

@AccessApiOverload
public fun group(column: AnyColumnGroupAccessor, body: AddDsl<T>.() -> Unit): Unit = group(column.name(), body)

public fun group(name: String, body: AddDsl<T>.() -> Unit) {
Expand All @@ -224,6 +235,7 @@ public class AddDsl<T>(

public infix fun AddGroup<T>.into(groupName: String): Unit = group(groupName, body)

@AccessApiOverload
public infix fun AddGroup<T>.into(column: AnyColumnGroupAccessor): Unit = into(column.name())
}

Expand All @@ -241,6 +253,7 @@ public inline fun <reified R, T, G> GroupBy<T, G>.add(
noinline expression: RowExpression<G, R>,
): GroupBy<T, G> = updateGroups { add(name, infer, expression) }

@AccessApiOverload
public inline fun <reified R, T, G> GroupBy<T, G>.add(
column: ColumnAccessor<G>,
infer: Infer = Infer.Nulls,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.api
import org.jetbrains.kotlinx.dataframe.AnyCol
import org.jetbrains.kotlinx.dataframe.AnyFrame
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
import org.jetbrains.kotlinx.dataframe.annotations.Refine
import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
Expand All @@ -15,6 +16,7 @@ public fun AnyCol.addId(columnName: String = "id"): AnyFrame = toDataFrame().add

// region DataFrame

@AccessApiOverload
public fun <T> DataFrame<T>.addId(column: ColumnAccessor<Int>): DataFrame<T> = insert(column) { index() }.at(0)

@Refine
Expand Down
Loading

0 comments on commit d88b609

Please sign in to comment.