Compiler plugin

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
@@ -216,6 +217,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

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

@@ -8,6 +8,7 @@ import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.Check
 import org.jetbrains.kotlinx.dataframe.columns.ColumnAccessor
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -22,6 +23,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.TransformableColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.columns.TransformableSingleColumn
 import kotlin.reflect.typeOf
 
+@Check
 public fun <T> AnyFrame.cast(): DataFrame<T> = this as DataFrame<T>
 
 public inline fun <reified T> AnyFrame.cast(verify: Boolean = true): DataFrame<T> = if (verify) convertToImpl(

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

@@ -27,6 +27,8 @@ public fun <T> DataRow<T>.concat(vararg rows: DataRow<T>): DataFrame<T> = (listO
 
 public fun <T> DataFrame<T>.concat(vararg frames: DataFrame<T>): DataFrame<T> = concatImpl(listOf(this) + frames)
 
+public infix fun <T> DataFrame<T>.concat(frame: DataFrame<T>): DataFrame<T> = concatImpl(listOf(this) + frame)
+
 @JvmName("concatT")
 public fun <T> DataFrame<T>.concat(rows: Iterable<DataRow<T>>): DataFrame<T> = (rows() + rows).concat()
 

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.RowColumnExpression
 import org.jetbrains.kotlinx.dataframe.RowValueExpression
+import org.jetbrains.kotlinx.dataframe.annotations.*
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.dataTypes.IFRAME
@@ -40,12 +41,14 @@ import kotlin.reflect.full.isSubtypeOf
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+@Interpretable("Convert0")
 public fun <T, C> DataFrame<T>.convert(columns: ColumnsSelector<T, C>): Convert<T, C> =
     Convert(this, columns)
 
 public fun <T, C> DataFrame<T>.convert(vararg columns: KProperty<C>): Convert<T, C> =
     convert { columns.toColumnSet() }
 
+@Interpretable("Convert2")
 public fun <T> DataFrame<T>.convert(vararg columns: String): Convert<T, Any?> = convert { columns.toColumnSet() }
 
 public fun <T, C> DataFrame<T>.convert(vararg columns: ColumnReference<C>): Convert<T, C> =
@@ -67,6 +70,7 @@ public inline fun <T, C, reified R> DataFrame<T>.convert(
 ): DataFrame<T> =
     convert(*headPlusArray(firstCol, cols)).with(infer, expression)
 
+@Interpretable("Convert6")
 public inline fun <T, reified R> DataFrame<T>.convert(
     firstCol: String,
     vararg cols: String,
@@ -81,29 +85,38 @@ public inline fun <T, C, reified R> Convert<T, C?>.notNull(crossinline expressio
         else expression(this, it)
     }
 
+@HasSchema(schemaArg = 0)
 public data class Convert<T, out C>(val df: DataFrame<T>, val columns: ColumnsSelector<T, C>) {
     public fun <R> cast(): Convert<T, R> = Convert(df, columns as ColumnsSelector<T, R>)
 
+    @Interpretable("To0")
     public inline fun <reified D> to(): DataFrame<T> = to(typeOf<D>())
 }
 
 public fun <T> Convert<T, *>.to(type: KType): DataFrame<T> = to { it.convertTo(type) }
 
+public fun <T, C> Convert<T, C>.to(columnConverter: DataFrame<T>.(DataColumn<C>) -> AnyBaseCol): DataFrame<T> =
+    df.replace(columns).with { columnConverter(df, it) }
+
+@Interpretable("With0")
 public inline fun <T, C, reified R> Convert<T, C>.with(
     infer: Infer = Infer.Nulls,
     noinline rowConverter: RowValueExpression<T, C, R>,
 ): DataFrame<T> =
     withRowCellImpl(typeOf<R>(), infer, rowConverter)
 
+@Refine
+@Interpretable("With0")
+public inline fun <T, C, reified R> Convert<T, C>.with(
+    noinline rowConverter: RowValueExpression<T, C, R>
+): DataFrame<T> = with(Infer.Nulls, rowConverter)
+
 public inline fun <T, C, reified R> Convert<T, C>.perRowCol(
     infer: Infer = Infer.Nulls,
     noinline expression: RowColumnExpression<T, C, R>,
 ): DataFrame<T> =
     convertRowColumnImpl(typeOf<R>(), infer, expression)
 
-public fun <T, C> Convert<T, C>.to(columnConverter: DataFrame<T>.(DataColumn<C>) -> AnyBaseCol): DataFrame<T> =
-    df.replace(columns).with { columnConverter(df, it) }
-
 public inline fun <reified C> AnyCol.convertTo(): DataColumn<C> = convertTo(typeOf<C>()) as DataColumn<C>
 
 public fun AnyCol.convertTo(newType: KType): AnyCol {

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

@@ -4,6 +4,8 @@ import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 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.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
@@ -13,7 +15,8 @@ import kotlin.reflect.KProperty
 private val defaultExplodeColumns: ColumnsSelector<*, *> = { colsAtAnyDepth { it.isList() || it.isFrameColumn() } }
 
 // region explode DataFrame
-
+@Refine
+@Interpretable("Explode0")
 public fun <T> DataFrame<T>.explode(
     dropEmpty: Boolean = true,
     selector: ColumnsSelector<T, *> = defaultExplodeColumns

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

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.api
 import org.jetbrains.kotlinx.dataframe.ColumnsContainer
 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.ExprColumnsSelectionDsl.Grammar
 import org.jetbrains.kotlinx.dataframe.api.ExprColumnsSelectionDsl.Grammar.PlainDslName
 import org.jetbrains.kotlinx.dataframe.documentation.ColumnExpression
@@ -91,6 +92,7 @@ public interface ExprColumnsSelectionDsl {
  * @param [expression] An [AddExpression] to define what each new row of the temporary column should contain.
  * @see [ColumnsContainer.mapToColumn]
  */
+@Interpretable("Expr0")
 public inline fun <T, reified R> ColumnsSelectionDsl<T>.expr(
     name: String = "",
     infer: Infer = Infer.Nulls,

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

@@ -4,6 +4,8 @@ import org.jetbrains.kotlinx.dataframe.AnyColumnGroupAccessor
 import org.jetbrains.kotlinx.dataframe.AnyColumnReference
 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.ColumnWithPath
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.columnName
@@ -12,6 +14,7 @@ import kotlin.reflect.KProperty
 
 // region DataFrame
 
+@Interpretable("Group0")
 public fun <T, C> DataFrame<T>.group(columns: ColumnsSelector<T, C>): GroupClause<T, C> = GroupClause(this, columns)
 
 public fun <T> DataFrame<T>.group(vararg columns: String): GroupClause<T, Any?> = group { columns.toColumnSet() }
@@ -39,6 +42,8 @@ public fun <T, C> GroupClause<T, C>.into(column: ColumnsSelectionDsl<T>.(ColumnW
 public fun <T, C> GroupClause<T, C>.into(column: ColumnsSelectionDsl<T>.(ColumnWithPath<C>) -> AnyColumnReference): DataFrame<T> =
     df.move(columns).under(column)
 
+@Refine
+@Interpretable("Into0")
 public fun <T, C> GroupClause<T, C>.into(column: String): DataFrame<T> = into(columnGroup().named(column))
 
 public fun <T, C> GroupClause<T, C>.into(column: AnyColumnGroupAccessor): DataFrame<T> = df.move(columns).under(column)