Merge pull request #729 from Kotlin/compiler-plugin

Compiler plugin

Author: koperagen

build.gradle.kts

@@ -52,11 +52,11 @@ dependencies {
     api(project(":dataframe-jdbc"))
 }
 
-private enum class Version : Comparable<Version> {
+enum class Version : Comparable<Version> {
     SNAPSHOT, DEV, ALPHA, BETA, RC, STABLE;
 }
 
-private fun String.findVersion(): Version {
+fun String.findVersion(): Version {
     val version = this.lowercase()
     return when {
         "snapshot" in version -> Version.SNAPSHOT

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe
 
 import org.jetbrains.kotlinx.dataframe.aggregation.Aggregatable
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedBody
+import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
 import org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl
 import org.jetbrains.kotlinx.dataframe.api.add
 import org.jetbrains.kotlinx.dataframe.api.cast
@@ -29,6 +30,7 @@ import kotlin.reflect.KType
  *
  * @param T Schema marker. It identifies column schema and is used to generate schema-specific extension properties for typed data access. It is covariant, so `DataFrame<A>` is assignable to variable of type `DataFrame<B>` if `A` is a subtype of `B`.
  */
+@HasSchema(schemaArg = 0)
 public interface DataFrame<out T> : Aggregatable<T>, ColumnsContainer<T> {
 
     public companion object {

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

@@ -43,6 +43,17 @@ public interface DataRow<out T> {
         return value as AnyRow
     }
 
+    public fun getFrameColumn(columnName: String): AnyFrame {
+        val value = get(columnName)
+        if (value == null) {
+            val kind = df()[columnName].kind()
+            if (kind != ColumnKind.Frame) {
+                error("Cannot cast null value of a $kind to a ${DataFrame::class}")
+            }
+        }
+        return value as AnyFrame
+    }
+
     public fun getOrNull(name: String): Any?
     public fun <R> getValueOrNull(column: ColumnReference<R>): R?
 

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/aggregation/AggregateDsl.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe.aggregation
 
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.ColumnSelectionDsl
 import org.jetbrains.kotlinx.dataframe.api.pathOf
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
@@ -12,6 +13,7 @@ import kotlin.reflect.typeOf
 
 public abstract class AggregateDsl<out T> : DataFrame<T>, ColumnSelectionDsl<T> {
 
+    @Interpretable("GroupByInto")
     public inline infix fun <reified R> R.into(name: String): NamedValue = internal().yield(pathOf(name), this, typeOf<R>())
 
     public inline infix fun <reified R> R.into(column: ColumnAccessor<R>): NamedValue = internal().yield(pathOf(column.name()), this, typeOf<R>())

core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt

@@ -0,0 +1,32 @@
+package org.jetbrains.kotlinx.dataframe.annotations
+
+@Target(AnnotationTarget.CLASS)
+public annotation class HasSchema(val schemaArg: Int)
+
+/**
+ * Compiler plugin will evaluate compile time value of the annotated function.
+ * Needed because some function calls only serve as a part of overall compile time DataSchema evaluation
+ * There's no need to update return type of such calls
+ */
+internal annotation class Interpretable(val interpreter: String)
+
+/**
+ * Compiler plugin will replace return type of calls to the annotated function
+ */
+internal annotation class Refine
+
+internal annotation class OptInRefine
+
+@Retention(AnnotationRetention.SOURCE)
+@Target(AnnotationTarget.FILE, AnnotationTarget.EXPRESSION)
+public annotation class DisableInterpretation
+
+@Retention(AnnotationRetention.SOURCE)
+@Target(AnnotationTarget.EXPRESSION)
+public annotation class Import
+
+@Target(AnnotationTarget.PROPERTY)
+public annotation class Order(val order: Int)
+
+@Target(AnnotationTarget.FUNCTION)
+internal annotation class Check

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

@@ -0,0 +1,11 @@
+package org.jetbrains.kotlinx.dataframe.api
+
+import org.jetbrains.kotlinx.dataframe.DataFrame
+
+public interface DataRowSchema
+
+public inline fun <reified T : DataRowSchema> dataFrameOf(vararg rows: T): DataFrame<T> =
+    rows.asIterable().toDataFrame()
+
+public inline fun <reified T : DataRowSchema> DataFrame<T>.append(vararg rows: T): DataFrame<T> =
+    concat(dataFrameOf(*rows))

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

@@ -1,6 +1,8 @@
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.api.Update.UpdateOperationArg
 import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -918,6 +920,8 @@ public fun <T, C> DataFrame<T>.fillNA(vararg columns: ColumnReference<C>): Updat
  *   If `false`, rows are dropped if any of the selected cells is `null`.
  * @param columns The [Columns Selector][org.jetbrains.kotlinx.dataframe.ColumnsSelector] used to select the columns of this [DataFrame][org.jetbrains.kotlinx.dataframe.DataFrame] to drop rows in.
  */
+@Refine
+@Interpretable("DropNulls0")
 public fun <T> DataFrame<T>.dropNulls(whereAllNull: Boolean = false, columns: ColumnsSelector<T, *>): DataFrame<T> {
     val cols = this[columns]
     return if (whereAllNull) drop { row -> cols.all { col -> col[row] == null } }

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

@@ -12,6 +12,8 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.RowExpression
 import org.jetbrains.kotlinx.dataframe.Selector
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.BaseColumn
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
@@ -121,6 +123,8 @@ public typealias AddExpression<T, R> = Selector<AddDataRow<T>, R>
  * @return new [DataFrame] with added column
  * @throws DuplicateColumnNamesException if [DataFrame] already contains a column with given [name]
  */
+@Refine
+@Interpretable("Add")
 public inline fun <reified R, T> DataFrame<T>.add(
     name: String,
     infer: Infer = Infer.Nulls,
@@ -178,7 +182,10 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
         return df.mapToColumn("", infer, expression)
     }
 
-    public inline infix fun <reified R> String.from(noinline expression: RowExpression<T, R>): Boolean =
+    @Interpretable("From")
+    public inline infix fun <reified R> String.from(
+        noinline expression: RowExpression<T, R>
+    ): Boolean =
         add(this, Infer.Nulls, expression)
 
     // TODO: use path instead of name
@@ -192,6 +199,7 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
     public inline infix fun <reified R> ColumnAccessor<R>.from(column: ColumnReference<R>): Boolean = name() from column
     public inline infix fun <reified R> KProperty<R>.from(column: ColumnReference<R>): Boolean = name from column
 
+    @Interpretable("Into")
     public infix fun AnyColumnReference.into(name: String): Boolean = add(rename(name))
     public infix fun <R> ColumnReference<R>.into(column: ColumnAccessor<R>): Boolean = into(column.name())
     public infix fun <R> ColumnReference<R>.into(column: KProperty<R>): Boolean = into(column.name)
@@ -212,6 +220,8 @@ public class AddDsl<T>(@PublishedApi internal val df: DataFrame<T>) : ColumnsCon
     public infix fun AddGroup<T>.into(column: AnyColumnGroupAccessor): Unit = into(column.name())
 }
 
+@Refine
+@Interpretable("AddWithDsl")
 public fun <T> DataFrame<T>.add(body: AddDsl<T>.() -> Unit): DataFrame<T> {
     val dsl = AddDsl(this)
     body(dsl)

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

@@ -1,11 +1,22 @@
 package org.jetbrains.kotlinx.dataframe.api
 
+import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.aggregation.AggregateDsl
+import org.jetbrains.kotlinx.dataframe.aggregation.AggregateGroupedBody
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
+import org.jetbrains.kotlinx.dataframe.impl.aggregateGroupBy
 
 // region Pivot
 
 public fun <T, R> Pivot<T>.aggregate(separate: Boolean = false, body: Selector<AggregateDsl<T>, R>): DataRow<T> = delegate { aggregate(separate, body) }
 
 // endregion
+
+@Refine
+@Interpretable("Aggregate")
+public fun <T, R> Grouped<T>.aggregate(body: AggregateGroupedBody<T, R>): DataFrame<T> {
+    return aggregateGroupBy((this as GroupBy<*, *>).toDataFrame(), { groups.cast() }, removeColumns = true, body).cast<T>()
+}

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.api.AndColumnsSelectionDsl.Grammar
 import org.jetbrains.kotlinx.dataframe.api.AndColumnsSelectionDsl.Grammar.InfixName
 import org.jetbrains.kotlinx.dataframe.api.AndColumnsSelectionDsl.Grammar.Name
@@ -210,6 +211,7 @@ public interface AndColumnsSelectionDsl {
      * @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] that contains all the columns from the [ColumnsResolvers][org.jetbrains.kotlinx.dataframe.columns.ColumnsResolver] on the left
      *   and right side of the [and][org.jetbrains.kotlinx.dataframe.api.AndColumnsSelectionDsl.and] operator.
      */
+    @Interpretable("And0")
     public infix fun <C> ColumnsResolver<C>.and(other: ColumnsResolver<C>): ColumnSet<C> = ColumnsList(this, other)
 
     /** ## And Operator