Merge pull request #1145 from Kotlin/cs-dsl-select

[Compiler plugin] Support SingleColumn.select

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/select.kt

@@ -178,6 +178,7 @@ public interface SelectColumnsSelectionDsl {
      *
      * `df.`[select][DataFrame.select]`  { myColGroup  `[`{`][SingleColumn.select]`  colA  `[and][ColumnsSelectionDsl.and]`  colB  `[`}`][SingleColumn.select]` }`
      */
+    @Interpretable("NestedSelect")
     public fun <C, R> SingleColumn<DataRow<C>>.select(selector: ColumnsSelector<C, R>): ColumnSet<R> =
         selectInternal(selector)
 

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt

@@ -693,7 +693,6 @@ internal class ValueCols2 : AbstractInterpreter<ColumnsResolver>() {
     }
 }
 
-
 internal class Named0 : AbstractInterpreter<ColumnsResolver>() {
     val Arguments.receiver: SingleColumnApproximation by arg()
     val Arguments.newName: String by arg()
@@ -702,3 +701,12 @@ internal class Named0 : AbstractInterpreter<ColumnsResolver>() {
         return columnsResolver { receiver named newName }
     }
 }
+
+internal class NestedSelect : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: SingleColumnApproximation by arg()
+    val Arguments.selector: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return columnsResolver { receiver.asColumnGroup().select { selector } }
+    }
+}

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt

@@ -174,6 +174,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NameStartsWith2
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Named0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.NestedSelect
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairToConstructor
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PerRowCol
@@ -316,6 +317,7 @@ internal inline fun <reified T> String.load(): T {
         "RenameMapping" -> RenameMapping()
         "Select0" -> Select0()
         "Distinct0" -> Select0()
+        "NestedSelect" -> NestedSelect()
         "Expr0" -> Expr0()
         "And0" -> And0()
         "Remove0" -> Remove0()

plugins/kotlin-dataframe/testData/box/csDsl/select.kt

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf(
+        "a" to listOf(1),
+        "b" to listOf(1),
+    ).group { a and b }.into("gr")
+    val res = df.select { gr.select { a and b } }
+    res.a
+    res.b
+    return "OK"
+}

plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

@@ -845,6 +845,12 @@ public void testRename() {
       runTest("testData/box/csDsl/rename.kt");
     }
 
+    @Test
+    @TestMetadata("select.kt")
+    public void testSelect() {
+      runTest("testData/box/csDsl/select.kt");
+    }
+
     @Test
     @TestMetadata("single.kt")
     public void testSingle() {