Skip to content

Column Selection DSL improvements #318

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Apr 6, 2023
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public interface ColumnsContainer<out T> {
public fun getColumnOrNull(name: String): AnyCol?
public fun getColumnOrNull(index: Int): AnyCol?
public fun <R> getColumnOrNull(column: ColumnReference<R>): DataColumn<R>?
public fun <R> getColumnOrNull(column: KProperty<R>): DataColumn<R>?
public fun getColumnOrNull(path: ColumnPath): AnyCol?
public fun <R> getColumnOrNull(column: ColumnSelector<T, R>): DataColumn<R>?

Expand Down

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
import org.jetbrains.kotlinx.dataframe.impl.columns.addParentPath
import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
import org.jetbrains.kotlinx.dataframe.impl.columns.depth
import kotlin.reflect.KProperty

public interface ColumnWithPath<out T> : DataColumn<T> {

Expand All @@ -17,6 +18,8 @@ public interface ColumnWithPath<out T> : DataColumn<T> {
public fun <C> getChild(accessor: ColumnReference<C>): ColumnWithPath<C>? = asColumnGroup().getColumnOrNull(accessor)?.addPath(path + accessor.path())
public fun getChild(name: String): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(name)?.addParentPath(path)
public fun getChild(index: Int): ColumnWithPath<Any?>? = asColumnGroup().getColumnOrNull(index)?.addParentPath(path)
public fun <C> getChild(accessor: KProperty<C>): ColumnWithPath<C>? = asColumnGroup().getColumnOrNull(accessor)?.addParentPath(path)

public fun children(): List<ColumnWithPath<Any?>> = if (isColumnGroup()) data.asColumnGroup().columns().map { it.addParentPath(path) } else emptyList()

override fun path(): ColumnPath = path
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.getColumn
import org.jetbrains.kotlinx.dataframe.api.indices
import org.jetbrains.kotlinx.dataframe.api.name
import org.jetbrains.kotlinx.dataframe.api.toColumnAccessor
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
Expand All @@ -24,6 +25,7 @@ import org.jetbrains.kotlinx.dataframe.impl.aggregation.AggregatableInternal
import org.jetbrains.kotlinx.dataframe.impl.aggregation.GroupByReceiverImpl
import org.jetbrains.kotlinx.dataframe.impl.columns.resolveSingle
import org.jetbrains.kotlinx.dataframe.io.renderToString
import kotlin.reflect.KProperty

private const val unnamedColumnPrefix = "untitled"

Expand Down Expand Up @@ -116,6 +118,7 @@ internal open class DataFrameImpl<T>(cols: List<AnyCol>, val nrow: Int) : DataFr
).singleOrNull()

override fun <R> getColumnOrNull(column: ColumnReference<R>): DataColumn<R>? = column.resolveSingle(this, UnresolvedColumnsPolicy.Skip)?.data
override fun <R> getColumnOrNull(column: KProperty<R>): DataColumn<R>? = getColumnOrNull(column.toColumnAccessor())

override fun getColumnOrNull(path: ColumnPath): AnyCol? =
when (path.size) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
import org.jetbrains.kotlinx.dataframe.impl.columns.DataColumnGroup
import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
import kotlin.reflect.KProperty
import kotlin.reflect.KType

internal class MissingColumnGroup<T>(val path: ColumnPath, val host: ColumnsContainer<*>) : MissingDataColumn<DataRow<T>>(), DataColumnGroup<T> {
Expand Down Expand Up @@ -49,6 +50,8 @@ internal class MissingColumnGroup<T>(val path: ColumnPath, val host: ColumnsCont

override fun <R> getColumnOrNull(column: ColumnReference<R>) = MissingColumnGroup<Any>(path + column.name(), host).asDataColumn().cast<R>()

override fun <R> getColumnOrNull(column: KProperty<R>) = MissingColumnGroup<Any>(path + column.name, host).asDataColumn().cast<R>()

override fun getColumnOrNull(path: ColumnPath) = MissingColumnGroup<Any?>(this.path + path, host)

override fun <R> getColumnOrNull(column: ColumnSelector<T, R>) = MissingColumnGroup<Any>(path + "", host).asDataColumn().cast<R>()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -706,6 +706,41 @@ class Access : TestBase() {
// SampleEnd
}

@Test
fun columnSelectors_kproperties() {
// SampleStart
// by column name
df.select { it[Person::name] }
df.select { (Person::name)() }
df.select { col(Person::name) }

// by column path
df.select { it[Person::name][Name::firstName] }
df.select { Person::name[Name::firstName] }

// with a new name
df.select { Person::name named "Full Name" }

// converted
df.select { Person::name[Name::firstName].map { it.lowercase() } }

// column arithmetics
df.select { 2021 - (Person::age)() }

// two columns
df.select { Person::name and Person::age }

// range of columns
df.select { Person::name..Person::age }

// all children of ColumnGroup
df.select { Person::name.all() }

// dfs traversal of all children columns
df.select { Person::name.allDfs() }
// SampleEnd
}

@Test
fun columnSelectors_strings() {
// SampleStart
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1899,9 +1899,9 @@ class DataFrameTests : BaseTest() {
@Test
fun `select all after`() {
typed.select { allAfter(age) } shouldBe typed.select { city and weight }
typed.select { allSince(age) } shouldBe typed.select { age and city and weight }
typed.select { allFrom(age) } shouldBe typed.select { age and city and weight }
typed.select { allBefore(age) } shouldBe typed.select { name }
typed.select { allUntil(age) } shouldBe typed.select { name and age }
typed.select { allUpTo(age) } shouldBe typed.select { name and age }
}

@Test
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public interface ColumnsContainer<out T> {
public fun getColumnOrNull(name: String): AnyCol?
public fun getColumnOrNull(index: Int): AnyCol?
public fun <R> getColumnOrNull(column: ColumnReference<R>): DataColumn<R>?
public fun <R> getColumnOrNull(column: KProperty<R>): DataColumn<R>?
public fun getColumnOrNull(path: ColumnPath): AnyCol?
public fun <R> getColumnOrNull(column: ColumnSelector<T, R>): DataColumn<R>?

Expand Down
Loading