From 2a36ca4de2139e71bde1934e0b678fd9a884e38e Mon Sep 17 00:00:00 2001
From: Nikita Klimenko <nikita.klimenko@jetbrains.com>
Date: Mon, 25 Aug 2025 13:24:03 +0300
Subject: [PATCH] Add annotations to support types extending DataRow in the
 compiler plugin

---
 .../jetbrains/kotlinx/dataframe/DataRow.kt    |  2 ++
 .../kotlinx/dataframe/annotations/Plugin.kt   | 20 +++++++++++++++++++
 .../jetbrains/kotlinx/dataframe/api/add.kt    |  2 ++
 3 files changed, 24 insertions(+)

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataRow.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataRow.kt
index b0e1c4415..b89437c97 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataRow.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/DataRow.kt
@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.dataframe
 
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
+import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
 import org.jetbrains.kotlinx.dataframe.api.next
 import org.jetbrains.kotlinx.dataframe.api.prev
 import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
@@ -18,6 +19,7 @@ import kotlin.reflect.KProperty
  *
  * @param T Schema marker. See [DataFrame] for details
  */
+@HasSchema(schemaArg = 0)
 public interface DataRow<out T> {
 
     public fun index(): Int
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt
index 42f4c8eec..da499b1fc 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt
@@ -2,6 +2,26 @@ package org.jetbrains.kotlinx.dataframe.annotations
 
 import kotlin.reflect.KClass
 
+/**
+ * Matches the type parameter of the annotated class to DataRow/DataFrame type parameter T.
+ *
+ * Annotate public API classes that implement DataRow/DataFrame interface to enable "extract schema/create column from values" operation:
+ * ```kotlin
+ * df.add {
+ *   "col" from { it }
+ * }
+ * ```
+ * Result before:
+ * `col: DataColumn<AddDataRow<MySchema>>`
+ *
+ * Result after:
+ *
+ * ```
+ * col:
+ *   col1: Int
+ *   col2: String
+ * ```
+ */
 @Target(AnnotationTarget.CLASS)
 public annotation class HasSchema(val schemaArg: Int)
 
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
index f9efa1e18..ae5754dce 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
@@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.Selector
 import org.jetbrains.kotlinx.dataframe.annotations.AccessApiOverload
+import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
 import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.api.add
@@ -99,6 +100,7 @@ public fun <T> DataFrame<T>.addAll(dataFrames: Iterable<AnyFrame>): DataFrame<T>
  * Receiver that is used by the [AddExpression] (for instance in the [add] and [update] operations)
  * to access new (added or updated) column value in preceding row.
  */
+@HasSchema(schemaArg = 0)
 public interface AddDataRow<out T> : DataRow<T> {
 
     /**