Restore invoke operator for ColumnReference in ColumnSelectionDsl. Add tests

Author: nikitinas

src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.ColumnSet
 import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
 import org.jetbrains.kotlinx.dataframe.columns.SingleColumn
 import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
 import org.jetbrains.kotlinx.dataframe.columns.renamedReference
@@ -35,6 +36,12 @@ import kotlin.reflect.typeOf
 
 public interface ColumnSelectionDsl<out T> : ColumnsContainer<T> {
 
+    public operator fun <C> ColumnReference<C>.invoke(): DataColumn<C> = get(this)
+
+    public operator fun <T> ColumnReference<DataRow<T>>.invoke(): ColumnGroup<T> = get(this)
+
+    public operator fun <T> ColumnReference<DataFrame<T>>.invoke(): FrameColumn<T> = get(this)
+
     public operator fun <C> ColumnPath.invoke(): DataColumn<C> = getColumn(this).cast()
 
     public operator fun <C> String.invoke(): DataColumn<C> = getColumn(this).cast()

tests/src/test/kotlin/org/jetbrains/kotlinx/dataframe/testSets/person/DataFrameTreeTests.kt

@@ -37,6 +37,7 @@ import org.jetbrains.kotlinx.dataframe.api.duplicate
 import org.jetbrains.kotlinx.dataframe.api.duplicateRows
 import org.jetbrains.kotlinx.dataframe.api.emptyDataFrame
 import org.jetbrains.kotlinx.dataframe.api.explode
+import org.jetbrains.kotlinx.dataframe.api.expr
 import org.jetbrains.kotlinx.dataframe.api.filter
 import org.jetbrains.kotlinx.dataframe.api.forEach
 import org.jetbrains.kotlinx.dataframe.api.forEachRow
@@ -61,6 +62,8 @@ import org.jetbrains.kotlinx.dataframe.api.join
 import org.jetbrains.kotlinx.dataframe.api.last
 import org.jetbrains.kotlinx.dataframe.api.map
 import org.jetbrains.kotlinx.dataframe.api.max
+import org.jetbrains.kotlinx.dataframe.api.maxBy
+import org.jetbrains.kotlinx.dataframe.api.median
 import org.jetbrains.kotlinx.dataframe.api.minus
 import org.jetbrains.kotlinx.dataframe.api.move
 import org.jetbrains.kotlinx.dataframe.api.moveTo
@@ -73,6 +76,7 @@ import org.jetbrains.kotlinx.dataframe.api.remove
 import org.jetbrains.kotlinx.dataframe.api.rename
 import org.jetbrains.kotlinx.dataframe.api.rows
 import org.jetbrains.kotlinx.dataframe.api.select
+import org.jetbrains.kotlinx.dataframe.api.single
 import org.jetbrains.kotlinx.dataframe.api.sortBy
 import org.jetbrains.kotlinx.dataframe.api.split
 import org.jetbrains.kotlinx.dataframe.api.sumOf
@@ -704,4 +708,22 @@ class DataFrameTreeTests : BaseTest() {
             it.concat(it).sortBy("id").remove("id")
         }
     }
+
+    @Test
+    fun `select column group`() {
+        typed2.aggregate {
+            nameAndCity()[2..3].name.distinct().single() into "name"
+        }["name"] shouldBe "Charlie"
+    }
+
+    @Test
+    fun `select frame column`() {
+        val group by frameColumn<GroupedPerson>()
+
+        typed2
+            .groupBy { expr { age > 30 } into "isOld" }.into(group)
+            .aggregate {
+                group().maxBy { rowsCount() }.weight.median() into "m"
+            }["m"] shouldBe 61
+    }
 }