[Compiler plugin] Support more parameters in read* functions

Author: koperagen

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

@@ -8,12 +8,12 @@ public annotation class HasSchema(val schemaArg: Int)
  * 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)
+public annotation class Interpretable(val interpreter: String)
 
 /**
  * Compiler plugin will replace return type of calls to the annotated function
  */
-internal annotation class Refine
+public annotation class Refine
 
 internal annotation class OptInRefine
 

dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt

@@ -24,6 +24,8 @@ import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataColumn
 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.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.forEach
 import org.jetbrains.kotlinx.dataframe.api.select
@@ -143,6 +145,8 @@ public fun DataFrame.Companion.readExcel(
  * @param nameRepairStrategy handling of column names.
  * The default behavior is [NameRepairStrategy.CHECK_UNIQUE]
  */
+@Refine
+@Interpretable("ReadExcel")
 public fun DataFrame.Companion.readExcel(
     fileOrUrl: String,
     sheetName: String? = null,
@@ -209,7 +213,9 @@ public fun DataFrame.Companion.readExcel(
  * @param range comma separated list of Excel column letters and column ranges (e.g. “A:E” or “A,C,E:F”)
  */
 @JvmInline
-public value class StringColumns(public val range: String)
+public value class StringColumns
+    @Interpretable("StringColumns")
+    constructor(public val range: String)
 
 public fun StringColumns.toFormattingOptions(formatter: DataFormatter = DataFormatter()): FormattingOptions =
     FormattingOptions(range, formatter)

plugins/kotlin-dataframe/build.gradle.kts

@@ -37,6 +37,7 @@ dependencies {
     testRuntimeOnly("org.jetbrains.kotlin:kotlin-annotations-jvm:$kotlinVersion")
 
     implementation(project(":core"))
+    implementation(project(":dataframe-excel"))
     api(libs.kotlinLogging)
     api("org.jetbrains.kotlinx:kotlinx-serialization-json:1.4.0-RC")
 

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

@@ -7,11 +7,17 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
 import org.jetbrains.kotlinx.dataframe.api.schema
+import org.jetbrains.kotlinx.dataframe.io.JSON
+import org.jetbrains.kotlinx.dataframe.io.JSON.TypeClashTactic.ARRAY_AND_VALUE_COLUMNS
+import org.jetbrains.kotlinx.dataframe.io.NameRepairStrategy
+import org.jetbrains.kotlinx.dataframe.io.StringColumns
 import org.jetbrains.kotlinx.dataframe.io.read
 import org.jetbrains.kotlinx.dataframe.io.readCSV
 import org.jetbrains.kotlinx.dataframe.io.readDelimStr
+import org.jetbrains.kotlinx.dataframe.io.readExcel
 import org.jetbrains.kotlinx.dataframe.io.readJson
 import org.jetbrains.kotlinx.dataframe.io.readJsonStr
+import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.IoSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.deserializeToPluginDataFrameSchema
@@ -30,13 +36,17 @@ internal class Read0 : AbstractInterpreter<PluginDataFrameSchema>() {
 
 internal class ReadCSV0 : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.fileOrUrl: String by arg()
+    val Arguments.delimiter: Char by arg(defaultValue = Present(','))
+    val Arguments.skipLines: Int by arg(defaultValue = Present(0))
+    val Arguments.readLines: Int? by arg(defaultValue = Present(null))
+    val Arguments.duplicate: Boolean by arg(defaultValue = Present(true))
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
         val file = resolveFile(resolutionPath, fileOrUrl)
         val df = if (file != null && file.exists()) {
-            DataFrame.readCSV(file)
+            DataFrame.readCSV(file, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
         } else {
-            DataFrame.readCSV(fileOrUrl)
+            DataFrame.readCSV(fileOrUrl, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
         }
         return df.schema().toPluginDataFrameSchema()
     }
@@ -91,17 +101,42 @@ private fun resolveFile(resolutionPath: String?, path: String): File? {
 internal class ReadDelimStr : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.text: String by arg()
     val Arguments.delimiter: Char by arg(defaultValue = Present(','))
+    val Arguments.skipLines: Int by arg(defaultValue = Present(0))
+    val Arguments.readLines: Int? by arg(defaultValue = Present(null))
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
-        return DataFrame.readDelimStr(text, delimiter).schema().toPluginDataFrameSchema()
+        return DataFrame.readDelimStr(text, delimiter, skipLines = skipLines, readLines = readLines).schema().toPluginDataFrameSchema()
     }
 }
 
 internal class ReadJsonStr : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.text: String by arg()
+    val Arguments.typeClashTactic: JSON.TypeClashTactic by arg(defaultValue = Present(ARRAY_AND_VALUE_COLUMNS))
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
-        return DataFrame.readJsonStr(text).schema().toPluginDataFrameSchema()
+        return DataFrame.readJsonStr(text, typeClashTactic = typeClashTactic).schema().toPluginDataFrameSchema()
     }
 }
 
+internal class ReadExcel : AbstractSchemaModificationInterpreter() {
+    val Arguments.fileOrUrl: String by arg()
+    val Arguments.sheetName: String? by arg(defaultValue = Present(null))
+    val Arguments.skipRows: Int by arg(defaultValue = Present(0))
+    val Arguments.columns: String? by arg(defaultValue = Present(null))
+    val Arguments.stringColumns: StringColumns? by arg(defaultValue = Present(null))
+    val Arguments.rowsCount: Int? by arg(defaultValue = Present(null))
+    val Arguments.nameRepairStrategy: NameRepairStrategy by arg(defaultValue = Present(NameRepairStrategy.CHECK_UNIQUE))
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val df = DataFrame.readExcel(fileOrUrl, sheetName, skipRows, columns, stringColumns, rowsCount, nameRepairStrategy)
+        return df.schema().toPluginDataFrameSchema()
+    }
+}
+
+internal class StringColumnsConstructor : AbstractInterpreter<StringColumns>() {
+    val Arguments.range: String by arg()
+
+    override fun Arguments.interpret(): StringColumns {
+        return StringColumns(range)
+    }
+}

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

@@ -74,6 +74,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
@@ -183,6 +184,8 @@ internal inline fun <reified T> String.load(): T {
         "DataFrameOf0" -> DataFrameOf0()
         "DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
         "ToDataFrameColumn" -> ToDataFrameColumn()
+        "StringColumns" -> ToDataFrameColumn()
+        "ReadExcel" -> ReadExcel()
         else -> error("$this")
     } as T
 }