Merge pull request #318 from Kotlin/kproperties-columnsSelectionDsl

Column Selection DSL improvements

Author: Jolanrensen

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnsContainer.kt

@@ -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>?
 

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl.kt

@@ -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> {
 
@@ -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

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/columns/ColumnWithPath.kt

@@ -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
@@ -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"
 
@@ -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) {

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/DataFrameImpl.kt

@@ -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> {
@@ -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>()

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/missing/MissingColumnGroup.kt

@@ -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

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Access.kt

@@ -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

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTests.kt

@@ -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>?