[Compiler plugin] Support flatten via adapter for structure-related, data-agnostic operations

Author: koperagen

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

@@ -2,16 +2,22 @@ package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.api.flattenImpl
 import kotlin.reflect.KProperty
 
 // region DataFrame
 
+@Refine
+@Interpretable("FlattenDefault")
 public fun <T> DataFrame<T>.flatten(keepParentNameForColumns: Boolean = false, separator: String = "."): DataFrame<T> =
     flatten(keepParentNameForColumns, separator) { all() }
 
+@Refine
+@Interpretable("Flatten0")
 public fun <T, C> DataFrame<T>.flatten(
     keepParentNameForColumns: Boolean = false,
     separator: String = ".",

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

@@ -0,0 +1,38 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl
+
+import org.jetbrains.kotlinx.dataframe.AnyCol
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.api.cast
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
+import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
+import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
+
+fun PluginDataFrameSchema.asDataFrame(): DataFrame<ConeTypesAdapter> = columns().map()
+
+fun DataFrame<ConeTypesAdapter>.toPluginDataFrameSchema() = PluginDataFrameSchema(columns().mapBack())
+
+interface ConeTypesAdapter
+
+private fun List<SimpleCol>.map(): DataFrame<ConeTypesAdapter> = map {
+    when (it) {
+        is SimpleDataColumn -> DataColumn.create(it.name, listOf(it.type))
+        is SimpleColumnGroup -> DataColumn.createColumnGroup(it.name, it.columns().map())
+        is SimpleFrameColumn -> DataColumn.createFrameColumn(it.name, listOf(it.columns().map()))
+    }
+}.let { dataFrameOf(it).cast() }
+
+private fun List<AnyCol>.mapBack(): List<SimpleCol> = map {
+    when (it) {
+        is ColumnGroup<*> -> {
+            SimpleColumnGroup(it.name(), it.columns().mapBack())
+        }
+        is FrameColumn<*> -> {
+            SimpleFrameColumn(it.name(), it[0].columns().mapBack())
+        }
+        else -> {
+            SimpleDataColumn(it.name(), it[0] as TypeApproximation)
+        }
+    }
+}

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

@@ -0,0 +1,38 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.flatten
+import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
+
+class FlattenDefault : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
+    val Arguments.separator: String by arg(defaultValue = Present("."))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return receiver.asDataFrame().flatten(keepParentNameForColumns, separator).toPluginDataFrameSchema()
+    }
+}
+
+class Flatten0 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver by dataFrame()
+    val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
+    val Arguments.separator: String by arg(defaultValue = Present("."))
+    val Arguments.columns: ColumnsResolver by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val columns = columns.resolve(receiver).map { pathOf(*it.path.path.toTypedArray()) }
+        return receiver
+            .asDataFrame()
+            .flatten(keepParentNameForColumns, separator) { columns.toColumnSet() }
+            .toPluginDataFrameSchema()
+    }
+}
+

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

@@ -74,6 +74,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
@@ -190,6 +192,8 @@ internal inline fun <reified T> String.load(): T {
         "ReadExcel" -> ReadExcel()
         "FillNulls0" -> FillNulls0()
         "UpdateWith0" -> UpdateWith0()
+        "Flatten0" -> Flatten0()
+        "FlattenDefault" -> FlattenDefault()
         else -> error("$this")
     } as T
 }

plugins/kotlin-dataframe/testData/box/flatten.kt

@@ -0,0 +1,23 @@
+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", "b", "c", "d")(1, 2, 3, 4)
+    val grouped = df
+        .group { a and b }.into("e")
+        .group { e and c }.into("f")
+
+    grouped.flatten().compareSchemas(strict = true)
+    val flattened = grouped.flatten { f.e }
+    flattened.compareSchemas(strict = true)
+    flattened.ungroup { f }.compareSchemas(strict = true)
+
+    grouped.flatten { f.e and f }.compareSchemas(strict = true)
+    return "OK"
+}
+
+
+
+

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

@@ -142,6 +142,12 @@ public void testFillNulls() {
     runTest("testData/box/fillNulls.kt");
   }
 
+  @Test
+  @TestMetadata("flatten.kt")
+  public void testFlatten() {
+    runTest("testData/box/flatten.kt");
+  }
+
   @Test
   @TestMetadata("flexibleReturnType.kt")
   public void testFlexibleReturnType() {