[Compiler plugin] toDataFrame from support

Author: koperagen

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

@@ -175,6 +175,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {
     public inline fun <reified R> add(name: String, noinline expression: (T) -> R): Unit =
         add(source.map { expression(it) }.toColumn(name, Infer.Nulls))
 
+    @Interpretable("ToDataFrameFrom0")
     public inline infix fun <reified R> String.from(noinline expression: (T) -> R): Unit =
         add(this, expression)
 

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

@@ -175,6 +175,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {
     public inline fun <reified R> add(name: String, noinline expression: (T) -> R): Unit =
         add(source.map { expression(it) }.toColumn(name, Infer.Nulls))
 
+    @Interpretable("ToDataFrameFrom0")
     public inline infix fun <reified R> String.from(noinline expression: (T) -> R): Unit =
         add(this, expression)
 

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.types.ConeTypeProjection
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.resolvedType
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.CreateDataFrameDslImplApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PluginDataFrameSchema
 
 fun KotlinTypeFacade.analyzeRefinedCallShape(call: FirFunctionCall, reporter: InterpretationErrorReporter): CallResult? {
@@ -43,17 +44,17 @@ fun KotlinTypeFacade.analyzeRefinedCallShape(call: FirFunctionCall, reporter: In
                 val lambda = (list.arguments.singleOrNull() as? FirAnonymousFunctionExpression)?.anonymousFunction
                 val statements = lambda?.body?.statements
                 if (statements != null) {
-                    val receiver = CreateDataFrameConfiguration()
-                    statements.filterIsInstance<FirFunctionCall>().forEach { call ->
-                        val schemaProcessor = call.loadInterpreter() ?: return@forEach
+                    val receiver = CreateDataFrameDslImplApproximation()
+                    statements.filterIsInstance<FirFunctionCall>().forEach {
+                        val schemaProcessor = it.loadInterpreter() ?: return@forEach
                         interpret(
-                            call,
+                            it,
                             schemaProcessor,
-                            mapOf("dsl" to Interpreter.Success(receiver)),
+                            mapOf("dsl" to Interpreter.Success(receiver), "call" to Interpreter.Success(call)),
                             reporter
                         )
                     }
-                    toDataFrame(receiver.maxDepth, call, receiver.traverseConfiguration)
+                    PluginDataFrameSchema(receiver.columns)
                 } else {
                     PluginDataFrameSchema(emptyList())
                 }

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

@@ -1,5 +1,9 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.isNullable
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
@@ -13,6 +17,8 @@ import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumnsToInsert
 import org.jetbrains.kotlinx.dataframe.impl.api.DataFrameLikeContainer
 import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumnGroup
 import org.jetbrains.kotlinx.dataframe.impl.api.insertImplGenericContainer
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
+import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnAccessorApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
@@ -28,6 +34,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.insertClause
 import org.jetbrains.kotlinx.dataframe.plugin.impl.kproperty
 import org.jetbrains.kotlinx.dataframe.plugin.impl.string
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
+import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 
 /**
  * @see DataFrame.insert
@@ -265,6 +273,34 @@ class SimpleColumnGroup(
     }
 }
 
+fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
+    return if (type.classId == Names.DATA_ROW_CLASS_ID) {
+        val schema = pluginDataFrameSchema(type)
+        val group = SimpleColumnGroup(name, schema.columns(), anyRow)
+        val column = if (type.isNullable) {
+            makeNullable(group)
+        } else {
+            group
+        }
+        column
+    } else if (type.classId == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
+        val schema = pluginDataFrameSchema(type)
+        SimpleFrameColumn(name, schema.columns(), anyDataFrame)
+    } else {
+        SimpleCol(name, type.wrap())
+    }
+}
+
+private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
+    return when (column) {
+        is SimpleColumnGroup -> {
+            SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) }, anyRow)
+        }
+        is SimpleFrameColumn -> column
+        else -> SimpleCol(column.name, column.type.changeNullability { true })
+    }
+}
+
 @PublishedApi
 internal fun PluginDataFrameSchema.insertImpl(
     columns: List<GenericColumnsToInsert<SimpleCol>>,

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

@@ -48,16 +48,21 @@ import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
+import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 import java.util.*
 
 class Properties0 : AbstractInterpreter<Unit>() {
-    val Arguments.dsl: CreateDataFrameConfiguration by arg()
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.call: FirFunctionCall by arg()
     val Arguments.maxDepth: Int by arg()
-    val Arguments.body: (Any) -> Unit by arg(lens = Interpreter.Dsl)
+    val Arguments.body: (Any) -> Unit by arg(lens = Interpreter.Dsl, defaultValue = Present(value = {}))
 
     override fun Arguments.interpret() {
-        dsl.maxDepth = maxDepth
-        body(dsl.traverseConfiguration)
+        dsl.configuration.maxDepth = maxDepth
+        body(dsl.configuration.traverseConfiguration)
+        val schema = toDataFrame(dsl.configuration.maxDepth, call, dsl.configuration.traverseConfiguration)
+        dsl.columns.addAll(schema.columns())
     }
 }
 
@@ -238,3 +243,17 @@ internal fun KotlinTypeFacade.toDataFrame(
         }
     }
 }
+
+class CreateDataFrameDslImplApproximation {
+    val configuration: CreateDataFrameConfiguration = CreateDataFrameConfiguration()
+    val columns: MutableList<SimpleCol> = mutableListOf()
+}
+
+class ToDataFrameFrom : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.expression: TypeApproximation by type()
+    override fun Arguments.interpret() {
+        dsl.columns += simpleColumnOf(receiver, expression.type)
+    }
+}

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

@@ -86,6 +86,10 @@ fun <T> KotlinTypeFacade.interpret(
 
     val defaultArguments = processor.expectedArguments.filter { it.defaultValue is Present }.map { it.name }.toSet()
     val actualArgsMap = refinedArguments.associateBy { it.name.identifier }.toSortedMap()
+    val conflictingKeys = additionalArguments.keys intersect actualArgsMap.keys
+    if (conflictingKeys.isNotEmpty()) {
+        error("Conflicting keys: $conflictingKeys")
+    }
     val expectedArgsMap = processor.expectedArguments
         .filterNot { it.name.startsWith("typeArg") }
         .associateBy { it.name }.toSortedMap().minus(additionalArguments.keys)

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

@@ -67,6 +67,7 @@ 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.ToDataFrameFrom
 
 internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
     val symbol =
@@ -159,6 +160,7 @@ internal inline fun <reified T> String.load(): T {
         "ReadJsonStr" -> ReadJsonStr()
         "ReadDelimStr" -> ReadDelimStr()
         "GroupByToDataFrame" -> GroupByToDataFrame()
+        "ToDataFrameFrom0" -> ToDataFrameFrom()
         else -> error("$this")
     } as T
 }