[Compiler plugin ] Handle "relative path" argument in read operations

Author: koperagen

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

@@ -1,6 +1,5 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
-import kotlinx.serialization.decodeFromString
 import kotlinx.serialization.json.Json
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
@@ -30,7 +29,11 @@ internal class Read0 : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.header: List<String> by arg(defaultValue = Present(listOf()))
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
-        return DataFrame.read(path).schema().toPluginDataFrameSchema()
+        val df = when (val source = tryResolveFile(resolutionPath, path)) {
+            is ResolutionDirFile -> DataFrame.read(source.file)
+            is UrlOrAbsolutePath -> DataFrame.read(source.path)
+        }
+        return df.schema().toPluginDataFrameSchema()
     }
 }
 
@@ -42,11 +45,13 @@ internal class ReadCSV0 : AbstractInterpreter<PluginDataFrameSchema>() {
     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, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
-        } else {
-            DataFrame.readCSV(fileOrUrl, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
+        val df = when (val source = tryResolveFile(resolutionPath, fileOrUrl)) {
+            is ResolutionDirFile -> {
+                DataFrame.readCSV(source.file, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
+            }
+            is UrlOrAbsolutePath -> {
+                DataFrame.readCSV(source.path, delimiter, skipLines = skipLines, readLines = readLines, duplicate = duplicate)
+            }
         }
         return df.schema().toPluginDataFrameSchema()
     }
@@ -74,17 +79,14 @@ internal class ReadJson0 : AbstractInterpreter<PluginDataFrameSchema>() {
 }
 
 fun readJson(resolutionPath: String?, path: String): DataFrame<Any?> {
-    val file = resolveFile(resolutionPath, path)
-    val df = if (file != null && file.exists()) {
-        DataFrame.readJson(file)
-    } else {
-        DataFrame.readJson(path)
+    return when (val source = tryResolveFile(resolutionPath, path)) {
+        is ResolutionDirFile -> DataFrame.readJson(source.file)
+        is UrlOrAbsolutePath -> DataFrame.readJson(source.path)
     }
-    return df
 }
 
-private fun resolveFile(resolutionPath: String?, path: String): File? {
-    return  resolutionPath?.let {
+private fun tryResolveFile(resolutionPath: String?, path: String): DataSource {
+    val file = resolutionPath?.let {
         try {
             val file = File(it)
             if (file.exists() && file.isDirectory) {
@@ -96,8 +98,17 @@ private fun resolveFile(resolutionPath: String?, path: String): File? {
             null
         }
     }
+    return if (file != null && file.exists()) {
+        ResolutionDirFile(file)
+    } else {
+        UrlOrAbsolutePath(path)
+    }
 }
 
+private sealed interface DataSource
+private class UrlOrAbsolutePath(val path: String) : DataSource
+private class ResolutionDirFile(val file: File) : DataSource
+
 internal class ReadDelimStr : AbstractInterpreter<PluginDataFrameSchema>() {
     val Arguments.text: String by arg()
     val Arguments.delimiter: Char by arg(defaultValue = Present(','))
@@ -128,7 +139,12 @@ internal class ReadExcel : AbstractSchemaModificationInterpreter() {
     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)
+        val df = when (val source = tryResolveFile(resolutionPath, fileOrUrl)) {
+            is ResolutionDirFile ->
+                DataFrame.readExcel(source.file, sheetName, skipRows, columns, stringColumns, rowsCount, nameRepairStrategy)
+            is UrlOrAbsolutePath ->
+                DataFrame.readExcel(source.path, sheetName, skipRows, columns, stringColumns, rowsCount, nameRepairStrategy)
+        }
         return df.schema().toPluginDataFrameSchema()
     }
 }

plugins/kotlin-dataframe/testData/box/readExcel.kt

@@ -0,0 +1,15 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = @Import DataFrame.read("testResources/sample.xls")
+    val d1: Double = df.col1[0]
+    val d2: Double = df.col2[0]
+
+    val df1 = @Import DataFrame.readExcel("testResources/sample.xls")
+    val d11: Double = df1.col1[0]
+    val d12: Double = df1.col2[0]
+    return "OK"
+}

plugins/kotlin-dataframe/testData/box/read_localFile.kt

@@ -0,0 +1,12 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = @Import DataFrame.read("testResources/achievements_all.json")
+
+    val df1 = df.explode { achievements }
+    df1.achievements.order
+    return "OK"
+}

plugins/kotlin-dataframe/testResources/sample.xls

@@ -340,6 +340,12 @@ public void testReadDelimStr_delimiter() {
     runTest("testData/box/readDelimStr_delimiter.kt");
   }
 
+  @Test
+  @TestMetadata("readExcel.kt")
+  public void testReadExcel() {
+    runTest("testData/box/readExcel.kt");
+  }
+
   @Test
   @TestMetadata("readJson.kt")
   public void testReadJson() {
@@ -364,6 +370,12 @@ public void testReadJsonStr_memberProperty() {
     runTest("testData/box/readJsonStr_memberProperty.kt");
   }
 
+  @Test
+  @TestMetadata("read_localFile.kt")
+  public void testRead_localFile() {
+    runTest("testData/box/read_localFile.kt");
+  }
+
   @Test
   @TestMetadata("remove.kt")
   public void testRemove() {