org.jetbrains.kotlinx.dataframe.api.update.kt Maven / Gradle / Ivy

Go to download
package org.jetbrains.kotlinx.dataframe.api

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.DataFrameExpression
import org.jetbrains.kotlinx.dataframe.DataRow
import org.jetbrains.kotlinx.dataframe.RowColumnExpression
import org.jetbrains.kotlinx.dataframe.RowValueExpression
import org.jetbrains.kotlinx.dataframe.RowValueFilter
import org.jetbrains.kotlinx.dataframe.Selector
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
import org.jetbrains.kotlinx.dataframe.impl.api.updateImpl
import org.jetbrains.kotlinx.dataframe.impl.api.updateWithValuePerColumnImpl
import org.jetbrains.kotlinx.dataframe.impl.columns.toColumnSet
import org.jetbrains.kotlinx.dataframe.impl.columns.toColumns
import org.jetbrains.kotlinx.dataframe.impl.headPlusArray
import org.jetbrains.kotlinx.dataframe.index
import kotlin.reflect.KProperty

public fun  DataFrame.update(columns: ColumnsSelector): Update =
    Update(this, null, columns)

public fun  DataFrame.update(columns: Iterable>): Update =
    update { columns.toColumnSet() }

public fun  DataFrame.update(vararg columns: String): Update = update { columns.toColumns() }
public fun  DataFrame.update(vararg columns: KProperty): Update = update { columns.toColumns() }
public fun  DataFrame.update(vararg columns: ColumnReference): Update =
    update { columns.toColumns() }

public data class Update(
    val df: DataFrame,
    val filter: RowValueFilter?,
    val columns: ColumnsSelector
) {
    public fun  cast(): Update =
        Update(df, filter as RowValueFilter?, columns as ColumnsSelector)
}

public fun  Update.where(predicate: RowValueFilter): Update =
    copy(filter = filter and predicate)

public fun  Update.at(rowIndices: Collection): Update = where { index in rowIndices }
public fun  Update.at(vararg rowIndices: Int): Update = at(rowIndices.toSet())
public fun  Update.at(rowRange: IntRange): Update = where { index in rowRange }

public infix fun  Update.perRowCol(expression: RowColumnExpression): DataFrame =
    updateImpl { row, column, _ -> expression(row, column) }

public typealias UpdateExpression = AddDataRow.(C) -> R

public infix fun  Update.with(expression: UpdateExpression): DataFrame =
    updateImpl { row, _, value ->
        expression(row, value)
    }

public infix fun  Update>.asFrame(expression: DataFrameExpression>): DataFrame =
    df.replace(columns).with { it.asColumnGroup().let { expression(it, it) }.asColumnGroup(it.name()) }

public fun  Update.asNullable(): Update = this as Update

public fun  Update.perCol(values: Map): DataFrame = updateWithValuePerColumnImpl {
    values[it.name()] ?: throw IllegalArgumentException("Update value for column ${it.name()} is not defined")
}

public fun  Update.perCol(values: AnyRow): DataFrame = perCol(values.toMap() as Map)

public fun  Update.perCol(valueSelector: Selector, C>): DataFrame =
    updateWithValuePerColumnImpl(valueSelector)

internal infix fun  RowValueFilter?.and(other: RowValueFilter): RowValueFilter {
    if (this == null) return other
    val thisExp = this
    return { thisExp(this, it) && other(this, it) }
}

public fun  Update.notNull(): Update =
    copy(filter = filter and { it != null }) as Update

public fun  Update.notNull(expression: RowValueExpression): DataFrame =
    notNull().updateImpl { row, column, value ->
        expression(row, value)
    }

public fun  DataFrame.update(
    firstCol: ColumnReference,
    vararg cols: ColumnReference,
    expression: RowValueExpression
): DataFrame =
    update(*headPlusArray(firstCol, cols)).with(expression)

public fun  DataFrame.update(
    firstCol: KProperty,
    vararg cols: KProperty,
    expression: RowValueExpression
): DataFrame =
    update(*headPlusArray(firstCol, cols)).with(expression)

public fun  DataFrame.update(
    firstCol: String,
    vararg cols: String,
    expression: RowValueExpression
): DataFrame =
    update(*headPlusArray(firstCol, cols)).with(expression)

public fun  Update.withNull(): DataFrame = asNullable().withValue(null)

public fun  Update.withZero(): DataFrame = updateWithValuePerColumnImpl { 0 as C }

public infix fun  Update.withValue(value: C): DataFrame = with { value }