Automated commit of generated code

Author: github-actions[bot]

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

@@ -0,0 +1,63 @@
+package org.jetbrains.kotlinx.dataframe.api
+
+import org.jetbrains.kotlinx.dataframe.ColumnSelector
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
+import org.jetbrains.kotlinx.dataframe.impl.api.requireImpl
+import kotlin.reflect.typeOf
+
+/**
+ * Resolves [column] in this [DataFrame] and checks that its runtime type is a subtype of [C].
+ * Throws if the column can't be resolved or if its type doesn't match.
+ *
+ * From the compiler plugin perspective, a new column will appear in the compile-time schema as a result of this operation.
+ *
+ * The aim here is to help incrementally migrate workflows to extension properties API.
+ *
+ * We recommend considering declaring a [DataSchema] and use [cast] or [convertTo] if you end up with more than a few `requireColumn` calls.
+ *
+ * Example:
+ *
+ * ```kotlin
+ * val repos = DataFrame
+ *     .readCsv("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
+ *
+ * repos
+ *     .filter { "stargazers_count"<Int>() > 100 }
+ *     .sortByDesc("stargazers_count")
+ *     .select("full_name", "stargazers_count")
+ * ```
+ *
+ * Notice how `stargazers_count` String is repeated three times. We can refactor this code using `requireColumn`:
+ *
+ * ```
+ * val repos = DataFrame
+ *     .readCsv("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
+ *     .requireColumn { "stargazers_count"<Int>() }
+ *
+ * repos
+ *     .filter { stargazers_count > 100 }
+ *     .sortByDesc { stargazers_count }
+ *     .select { "full_name" and stargazers_count }
+ * ```
+ *
+ * This way code becomes a bit more robust. For example, usages of a renamed column will become compile time errors that are easy to spot and update:
+ * ```kotlin
+ * val repos = DataFrame
+ *     .readCsv("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv")
+ *     .requireColumn { "stargazers_count"<Int>() }
+ *     .rename { stargazers_count }.into("stars")
+ *
+ * repos
+ *     .filter { stars > 100 }
+ *     .sortByDesc { stars }
+ *     .select { "full_name" and stars }
+ * ```
+ *
+ */
+@Refine
+@Interpretable("Require0")
+public inline fun <T, reified C> DataFrame<T>.requireColumn(noinline column: ColumnSelector<T, C>): DataFrame<T> =
+    requireImpl(column, typeOf<C>())

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

@@ -8,6 +8,7 @@ import org.jetbrains.kotlinx.dataframe.api.asDataColumn
 import org.jetbrains.kotlinx.dataframe.api.cast
 import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.pathOf
+import org.jetbrains.kotlinx.dataframe.api.toPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -21,6 +22,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columns.addPath
 import org.jetbrains.kotlinx.dataframe.impl.columns.missing.MissingColumnGroup
 import org.jetbrains.kotlinx.dataframe.impl.columns.missing.MissingDataColumn
 import org.jetbrains.kotlinx.dataframe.nrow
+import kotlin.collections.map
 
 private fun <T> DataFrame<T>.unbox(): DataFrame<T> =
     when (this) {
@@ -47,9 +49,7 @@ internal open class DataFrameReceiver<T>(
                         host = this@DataFrameReceiver,
                     ).asDataColumn().cast()
 
-                UnresolvedColumnsPolicy.Fail -> error(
-                    "Column '${path.joinToString()}' not found among ${df.columnNames()}.",
-                )
+                UnresolvedColumnsPolicy.Fail -> error(formatMissingColumnMessage(path))
             }
 
             is MissingDataColumn -> this
@@ -59,6 +59,43 @@ internal open class DataFrameReceiver<T>(
             else -> this
         }
 
+    // Context:
+    // it's strange that we have to reverse-search why the column is missing
+    // would be nice to "fail fast" exactly where resolve failed, knowing the current path and parent.
+    // but it's very unclear what to do with resolveSingle.
+    // at first glance: a lot of changes.
+    @Suppress("FoldInitializerAndIfToElvis")
+    private fun formatMissingColumnMessage(path: ColumnPath): String {
+        val fullPath = path.joinToString()
+
+        for (depth in path.indices) {
+            val currentPath = path.slice(0..depth).toPath()
+            val currentPathString = currentPath.joinToString()
+            val column = df.getColumnOrNull(currentPath)
+            if (column == null) {
+                return if (depth == 0) {
+                    "Column '$currentPathString' not found among ${df.columnNames()}."
+                } else {
+                    val parentPath = currentPath.dropLast()
+                    val parentPathString = parentPath.joinToString()
+                    val parentColumn = df.getColumnOrNull(parentPath)
+                    if (parentColumn != null && parentColumn.isColumnGroup()) {
+                        "Column '$currentPathString' not found among columns of '$parentPathString': ${parentColumn.columnNames()}."
+                    } else {
+                        "Column '$currentPathString' not found among ${df.columnNames()}."
+                    }
+                }
+            }
+
+            if (depth != path.lastIndex) {
+                if (!column.isColumnGroup()) {
+                    return "Column '$fullPath' cannot be resolved: '$currentPathString' is not a column group."
+                }
+            }
+        }
+        return "Column '$fullPath' not found among ${df.columnNames()}."
+    }
+
     override fun getColumnOrNull(name: String) = df.getColumnOrNull(name).check(pathOf(name))
 
     override fun getColumnOrNull(index: Int) = df.getColumnOrNull(index).check(pathOf(""))

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

@@ -0,0 +1,18 @@
+package org.jetbrains.kotlinx.dataframe.impl.api
+
+import org.jetbrains.kotlinx.dataframe.ColumnSelector
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.api.getColumnWithPath
+import org.jetbrains.kotlinx.dataframe.api.isSubtypeOf
+import org.jetbrains.kotlinx.dataframe.type
+import kotlin.reflect.KType
+
+@PublishedApi
+internal fun <T, C> DataFrame<T>.requireImpl(column: ColumnSelector<T, C>, type: KType): DataFrame<T> {
+    val resolvedColumn = getColumnWithPath(column)
+    val actualType = resolvedColumn.data.type
+    require(resolvedColumn.data.isSubtypeOf(type)) {
+        "Column '${resolvedColumn.path.joinToString()}' has type '$actualType', which is not subtype of required '$type' type."
+    }
+    return this
+}