[Compiler plugin] Support colsOf, colsAtAnyDepth, frameCols

koperagen · koperagen · commit 48247e11c654 · 2024-07-09T20:27:02.000+03:00
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/all.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.Predicate
 import org.jetbrains.kotlinx.dataframe.RowFilter
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.BehaviorArg
 import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ColumnDoesNotExistArg
 import org.jetbrains.kotlinx.dataframe.api.AllColumnsSelectionDsl.CommonAllSubsetDocs.ExampleArg
@@ -300,6 +301,7 @@ public interface AllColumnsSelectionDsl<out _UNUSED> {
      *
      * `df.`[select][DataFrame.select]`  {  `[all][ColumnsSelectionDsl.all]`() }`
      */
+    @Interpretable("All0")
     public fun ColumnsSelectionDsl<*>.all(): TransformableColumnSet<*> =
         asSingleColumn().allColumnsInternal()
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsAtAnyDepth.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar
 import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnGroupName
 import org.jetbrains.kotlinx.dataframe.api.ColsAtAnyDepthColumnsSelectionDsl.Grammar.ColumnSetName
@@ -138,6 +139,7 @@ public interface ColsAtAnyDepthColumnsSelectionDsl {
      *
      * `df.`[select][DataFrame.select]`  {  `[colsAtAnyDepth][ColumnsSelectionDsl.colsAtAnyDepth]` { !it.`[isColumnGroup][DataColumn.isColumnGroup]` } }`
      */
+    @Interpretable("ColsAtAnyDepth0")
     public fun ColumnsSelectionDsl<*>.colsAtAnyDepth(predicate: ColumnFilter<*> = { true }): ColumnSet<*> =
         asSingleColumn().colsAtAnyDepthInternal(predicate)
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/colsOf.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnFilter
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar
 import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnGroupName
 import org.jetbrains.kotlinx.dataframe.api.ColsOfColumnsSelectionDsl.Grammar.ColumnSetName
@@ -203,6 +204,7 @@ public fun <C> ColumnSet<*>.colsOf(
  * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam]
  * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return]
  */
+@Interpretable("ColsOf1")
 public inline fun <reified C> ColumnSet<*>.colsOf(
     noinline filter: ColumnFilter<C> = { true },
 ): TransformableColumnSet<C> = colsOf(typeOf<C>(), filter)
@@ -228,6 +230,7 @@ public fun <C> ColumnsSelectionDsl<*>.colsOf(
  * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.FilterParam]
  * @include [ColsOfColumnsSelectionDsl.CommonColsOfDocs.Return]
  */
+@Interpretable("ColsOf0")
 public inline fun <reified C> ColumnsSelectionDsl<*>.colsOf(
     noinline filter: ColumnFilter<C> = { true },
 ): TransformableColumnSet<C> = asSingleColumn().colsOf(typeOf<C>(), filter)
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/frameCols.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.Predicate
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnGroupName
 import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.ColumnSetName
 import org.jetbrains.kotlinx.dataframe.api.FrameColsColumnsSelectionDsl.Grammar.PlainDslName
@@ -111,6 +112,7 @@ public interface FrameColsColumnsSelectionDsl {
      *
      * `df.`[select][DataFrame.select]`  {  `[frameCols][ColumnsSelectionDsl.frameCols]` { it.`[name][ColumnReference.name]`.`[startsWith][String.startsWith]`("my") } }`
      */
+    @Interpretable("FrameCols0")
     public fun ColumnSet<*>.frameCols(filter: Predicate<FrameColumn<*>> = { true }): TransformableColumnSet<DataFrame<*>> =
         frameColumnsInternal(filter)
 
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExpressionAnalysisAdditionalChecker.kt
@@ -90,8 +90,8 @@ private class Checker(val cache: FirCache<String, PluginDataFrameSchema, KotlinT
             val targetProjection = expression.typeArguments.getOrNull(0) as? FirTypeProjectionWithVariance ?: return
             val targetType = targetProjection.typeRef.coneType as? ConeClassLikeType ?: return
             val target = pluginDataFrameSchema(targetType)
-            val sourceColumns = source.flatten()
-            val targetColumns = target.flatten()
+            val sourceColumns = source.flatten(includeFrames = true)
+            val targetColumns = target.flatten(includeFrames = true)
             val sourceMap = sourceColumns.associate { it.path.path to it.column }
             val missingColumns = mutableListOf<String>()
             var valid = true
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/flattenPluginSchema.kt
@@ -8,24 +8,24 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 
-fun PluginDataFrameSchema.flatten(): List<ColumnWithPathApproximation> {
+fun PluginDataFrameSchema.flatten(includeFrames: Boolean): List<ColumnWithPathApproximation> {
     if (columns().isEmpty()) return emptyList()
     val columns = mutableListOf<ColumnWithPathApproximation>()
-    flattenImpl(columns(), emptyList(), columns)
+    flattenImpl(columns(), emptyList(), columns, includeFrames)
     return columns
 }
 
-fun flattenImpl(columns: List<SimpleCol>, path: List<String>, flatList: MutableList<ColumnWithPathApproximation>) {
+fun flattenImpl(columns: List<SimpleCol>, path: List<String>, flatList: MutableList<ColumnWithPathApproximation>, includeFrames: Boolean) {
     columns.forEach { column ->
         val fullPath = path + listOf(column.name)
         when (column) {
             is SimpleColumnGroup -> {
                 flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
-                flattenImpl(column.columns(), fullPath, flatList)
+                flattenImpl(column.columns(), fullPath, flatList, includeFrames)
             }
             is SimpleFrameColumn -> {
                 flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
-                flattenImpl(column.columns(), fullPath, flatList)
+                flattenImpl(column.columns(), fullPath, flatList, includeFrames)
             }
             is SimpleDataColumn -> {
                 flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/select.kt
@@ -1,11 +1,14 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.isSubtypeOf
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
 import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
@@ -27,7 +30,12 @@ internal class Expr0 : AbstractInterpreter<ColumnsResolver>() {
     val Arguments.expression: TypeApproximation by type()
 
     override fun Arguments.interpret(): ColumnsResolver {
-        return SingleColumnApproximation(ColumnWithPathApproximation(ColumnPathApproximation(listOf(name)), SimpleDataColumn(name, expression)))
+        return SingleColumnApproximation(
+            ColumnWithPathApproximation(
+                ColumnPathApproximation(listOf(name)),
+                SimpleDataColumn(name, expression)
+            )
+        )
     }
 }
 
@@ -43,3 +51,75 @@ internal class And0 : AbstractInterpreter<ColumnsResolver>() {
         }
     }
 }
+
+internal class All0 : AbstractInterpreter<ColumnsResolver>() {
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return df.columns().map {
+                    val path = ColumnPathApproximation(listOf(it.name))
+                    ColumnWithPathApproximation(path, it)
+                }
+            }
+        }
+    }
+}
+
+internal class ColsOf0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.typeArg0: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                val cols = df.columns().map {
+                    val path = ColumnPathApproximation(listOf(it.name))
+                    ColumnWithPathApproximation(path, it)
+                }
+                return colsOf(cols, typeArg0.type)
+            }
+        }
+    }
+
+}
+
+private fun Arguments.colsOf(cols: List<ColumnWithPathApproximation>, type: ConeKotlinType) =
+    cols
+        .filter {
+            val column = it.column
+            column is SimpleDataColumn && column.type.type.isSubtypeOf(type, session)
+        }
+
+internal class ColsAtAnyDepth0 : AbstractInterpreter<ColumnsResolver>() {
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return df.flatten(includeFrames = false)
+            }
+        }
+    }
+}
+
+internal class ColsOf1 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+    val Arguments.typeArg0: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return colsOf(receiver.resolve(df), typeArg0.type)
+            }
+        }
+    }
+}
+
+internal class FrameCols0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): ColumnsResolver {
+        return object : ColumnsResolver {
+            override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
+                return receiver.resolve(df).filter { it.column is SimpleFrameColumn }
+            }
+        }
+    }
+}
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -67,6 +67,11 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
@@ -161,6 +166,11 @@ internal inline fun <reified T> String.load(): T {
         "ReadDelimStr" -> ReadDelimStr()
         "GroupByToDataFrame" -> GroupByToDataFrame()
         "ToDataFrameFrom0" -> ToDataFrameFrom()
+        "All0" -> All0()
+        "ColsOf0" -> ColsOf0()
+        "ColsOf1" -> ColsOf1()
+        "ColsAtAnyDepth0" -> ColsAtAnyDepth0()
+        "FrameCols0" -> FrameCols0()
         else -> error("$this")
     } as T
 }
diff --git a/plugins/kotlin-dataframe/testData/box/selectionDsl.kt b/plugins/kotlin-dataframe/testData/box/selectionDsl.kt
@@ -0,0 +1,17 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+data class Nested(val d: List<Double>)
+
+data class Record(val a: String, val b: Int, val nested: List<Nested>)
+
+fun box(): String {
+    val df = listOf(Record("112", 42, listOf(Nested(listOf(3.0))))).toDataFrame(maxDepth = 1)
+
+    df.group { nested }.into("group").convert { colsAtAnyDepth().frameCols() }.with { 1 }.compareSchemas()
+    df.group { b }.into("group").convert { colsAtAnyDepth().colsOf<Int>() }.with { "" }.compareSchemas()
+    df.convert { all().frameCols() }.with { 1 }.compareSchemas()
+    return "OK"
+}
diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
@@ -310,6 +310,12 @@ public void testSelectThis() {
     runTest("testData/box/selectThis.kt");
   }
 
+  @Test
+  @TestMetadata("selectionDsl.kt")
+  public void testSelectionDsl() {
+    runTest("testData/box/selectionDsl.kt");
+  }
+
   @Test
   @TestMetadata("toDataFrame.kt")
   public void testToDataFrame() {
