Implement format guessing with ServiceLoader

Previous approach with enum for all supported formats no longer works when arrow support is in the different module

Author: koperagen

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

@@ -49,6 +49,16 @@ import java.time.Duration
 import java.time.LocalDateTime
 import kotlin.reflect.typeOf
 
+public class ArrowFeather : SupportedFormat {
+    override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readArrowFeather(stream)
+
+    override fun readDataFrame(file: File, header: List<String>): AnyFrame = DataFrame.readArrowFeather(file)
+
+    override fun acceptsExtension(ext: String): Boolean = ext == "feather"
+
+    override val testOrder: Int = 50000
+}
+
 internal object Allocator {
     val ROOT by lazy {
         RootAllocator(Long.MAX_VALUE)

dataframe-arrow/src/main/resources/META-INF/services/org.jetbrains.kotlinx.dataframe.io.SupportedFormat

@@ -0,0 +1 @@
+org.jetbrains.kotlinx.dataframe.io.ArrowFeather

src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/SchemaReader.kt

@@ -10,7 +10,11 @@ import org.jetbrains.kotlinx.dataframe.codeGen.DefaultReadDfMethod
 import org.jetbrains.kotlinx.dataframe.codeGen.DefaultReadExcelMethod
 import org.jetbrains.kotlinx.dataframe.codeGen.DefaultReadJsonMethod
 import org.jetbrains.kotlinx.dataframe.codeGen.DefaultReadTsvMethod
+import org.jetbrains.kotlinx.dataframe.io.CSV
+import org.jetbrains.kotlinx.dataframe.io.Excel
+import org.jetbrains.kotlinx.dataframe.io.JSON
 import org.jetbrains.kotlinx.dataframe.io.SupportedFormats
+import org.jetbrains.kotlinx.dataframe.io.TSV
 import org.jetbrains.kotlinx.dataframe.io.guessFormat
 import org.jetbrains.kotlinx.dataframe.io.readCSV
 import org.jetbrains.kotlinx.dataframe.io.readExcel
@@ -38,11 +42,11 @@ public val CodeGenerator.Companion.urlReader: (url: URL, csvOptions: CsvOptions)
         fun readExcel(url: URL) = DfReadResult.Success(DataFrame.readExcel(url), SupportedFormats.EXCEL, csvOptions)
         try {
             val res = when (guessFormat(url.path)) {
-                SupportedFormats.CSV -> readCSV(url)
-                SupportedFormats.TSV -> readTSV(url)
-                SupportedFormats.JSON -> readJson(url)
+                is CSV -> readCSV(url)
+                is TSV -> readTSV(url)
+                is JSON -> readJson(url)
 //                SupportedFormats.ARROW -> readArrow(url)
-                SupportedFormats.EXCEL -> readExcel(url)
+                is Excel -> readExcel(url)
                 null -> try {
                     readExcel(url)
                 } catch (e: Exception) {
@@ -56,6 +60,7 @@ public val CodeGenerator.Companion.urlReader: (url: URL, csvOptions: CsvOptions)
                         }
                     }
                 }
+                else -> TODO()
             }
             res
         } catch (e: Throwable) {

src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/csv.kt

@@ -32,6 +32,16 @@ import kotlin.reflect.KClass
 import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
+public class CSV : SupportedFormat {
+    override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readCSV(stream, header = header)
+
+    override fun readDataFrame(file: File, header: List<String>): AnyFrame = DataFrame.readCSV(file, header = header)
+
+    override fun acceptsExtension(ext: String): Boolean = ext == "csv"
+
+    override val testOrder: Int = 20000
+}
+
 public enum class CSVType(public val format: CSVFormat) {
     DEFAULT(CSVFormat.DEFAULT.withAllowMissingColumnNames().withIgnoreSurroundingSpaces()),
     TDF(CSVFormat.TDF.withAllowMissingColumnNames())

src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/guess.kt

@@ -10,6 +10,19 @@ import java.io.File
 import java.io.FileNotFoundException
 import java.io.InputStream
 import java.net.URL
+import java.util.ServiceLoader
+
+public interface SupportedFormat {
+    public fun readDataFrame(stream: InputStream, header: List<String> = emptyList()): AnyFrame
+
+    public fun readDataFrame(file: File, header: List<String> = emptyList()): AnyFrame
+
+    public fun acceptsExtension(ext: String): Boolean
+
+    // `DataFrame.Companion.read` methods uses this to sort list of all supported formats in ascending order (-1, 2, 10)
+    // sorted list is used to test if any format can read given input
+    public val testOrder: Int
+}
 
 public enum class SupportedFormats {
     CSV {
@@ -56,21 +69,19 @@ public enum class SupportedFormats {
     internal abstract fun acceptsExtension(ext: String): Boolean
 }
 
-private val testOrder get() = listOf(
-    SupportedFormats.JSON,
-    SupportedFormats.CSV,
-    SupportedFormats.TSV,
-    SupportedFormats.EXCEL,
-//    SupportedFormats.ARROW,
-)
+internal val supportedFormats: List<SupportedFormat> by lazy {
+    ServiceLoader.load(SupportedFormat::class.java).toList()
+}
+
+internal val testOrder: List<SupportedFormat> by lazy { supportedFormats.sortedBy { it.testOrder } }
 
-internal fun guessFormatForExtension(ext: String) = SupportedFormats.values().firstOrNull { it.acceptsExtension(ext) }
+internal fun guessFormatForExtension(ext: String) = supportedFormats.firstOrNull { it.acceptsExtension(ext) }
 
-internal fun guessFormat(file: File): SupportedFormats? = file.extension.lowercase().let { guessFormatForExtension(it) }
+internal fun guessFormat(file: File): SupportedFormat? = file.extension.lowercase().let { guessFormatForExtension(it) }
 
-internal fun guessFormat(url: URL): SupportedFormats? = guessFormat(url.path)
+internal fun guessFormat(url: URL): SupportedFormat? = guessFormat(url.path)
 
-internal fun guessFormat(url: String): SupportedFormats? = guessFormatForExtension(url.substringAfterLast("."))
+internal fun guessFormat(url: String): SupportedFormat? = guessFormatForExtension(url.substringAfterLast("."))
 
 private class NotCloseableStream(val src: InputStream) : InputStream() {
     override fun read(): Int = src.read()
@@ -85,10 +96,10 @@ private class NotCloseableStream(val src: InputStream) : InputStream() {
 
 internal fun DataFrame.Companion.read(
     stream: InputStream,
-    format: SupportedFormats? = null,
+    format: SupportedFormat? = null,
     header: List<String> = emptyList()
-): AnyFrame {
-    if (format != null) return format.readDataFrame(stream, header = header)
+): ReadAnyFrame {
+    if (format != null) return format to format.readDataFrame(stream, header = header)
     val input = NotCloseableStream(if (stream.markSupported()) stream else BufferedInputStream(stream))
     try {
         val readLimit = 10000
@@ -97,7 +108,7 @@ internal fun DataFrame.Companion.read(
         testOrder.forEach {
             try {
                 input.reset()
-                return it.readDataFrame(input, header = header)
+                return it to it.readDataFrame(input, header = header)
             } catch (e: Exception) {
             }
         }
@@ -109,24 +120,28 @@ internal fun DataFrame.Companion.read(
 
 internal fun DataFrame.Companion.read(
     file: File,
-    format: SupportedFormats? = null,
+    format: SupportedFormat? = null,
     header: List<String> = emptyList()
-): AnyFrame {
-    if (format != null) return format.readDataFrame(file, header = header)
+): ReadAnyFrame {
+    if (format != null) return format to format.readDataFrame(file, header = header)
     testOrder.forEach {
         try {
-            return it.readDataFrame(file, header = header)
+            return it to it.readDataFrame(file, header = header)
         } catch (e: FileNotFoundException) { throw e } catch (e: Exception) { }
     }
     throw IllegalArgumentException("Unknown file format")
 }
 
-public fun DataFrame.Companion.read(file: File, header: List<String> = emptyList()): AnyFrame = read(file, guessFormat(file), header)
+internal data class ReadAnyFrame(val format: SupportedFormat, val df: AnyFrame)
+
+internal infix fun SupportedFormat.to(df: AnyFrame) = ReadAnyFrame(this, df)
+
+public fun DataFrame.Companion.read(file: File, header: List<String> = emptyList()): AnyFrame = read(file, guessFormat(file), header).df
 public fun DataRow.Companion.read(file: File, header: List<String> = emptyList()): AnyRow = DataFrame.read(file, header).single()
 
 public fun DataFrame.Companion.read(url: URL, header: List<String> = emptyList()): AnyFrame = when {
     isFile(url) -> read(urlAsFile(url), header)
-    isProtocolSupported(url) -> catchHttpResponse(url) { read(it, guessFormat(url), header) }
+    isProtocolSupported(url) -> catchHttpResponse(url) { read(it, guessFormat(url), header).df }
     else -> throw IllegalArgumentException("Invalid protocol for url $url")
 }
 

src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt

@@ -39,6 +39,16 @@ import kotlin.reflect.KTypeProjection
 import kotlin.reflect.full.createType
 import kotlin.reflect.typeOf
 
+public class JSON : SupportedFormat {
+    override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readJson(stream, header = header)
+
+    override fun readDataFrame(file: File, header: List<String>): AnyFrame = DataFrame.readJson(file, header = header)
+
+    override fun acceptsExtension(ext: String): Boolean = ext == "json"
+
+    override val testOrder: Int = 10000
+}
+
 public fun DataFrame.Companion.readJson(file: File, header: List<String> = emptyList()): AnyFrame = readJson(file.toURI().toURL(), header)
 public fun DataRow.Companion.readJson(file: File, header: List<String> = emptyList()): AnyRow = DataFrame.readJson(file, header).single()
 

src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/tsv.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.io
 
+import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.ParserOptions
 import java.io.File
@@ -8,6 +9,16 @@ import java.io.InputStream
 import java.net.URL
 import java.nio.charset.Charset
 
+public class TSV : SupportedFormat {
+    override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readTSV(stream, header = header)
+
+    override fun readDataFrame(file: File, header: List<String>): AnyFrame = DataFrame.readTSV(file, header = header)
+
+    override fun acceptsExtension(ext: String): Boolean = ext == "tsv"
+
+    override val testOrder: Int = 30000
+}
+
 private val tabChar = '\t'
 
 public fun DataFrame.Companion.readTSV(

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

@@ -29,6 +29,16 @@ import java.time.LocalDate
 import java.time.LocalDateTime
 import java.util.*
 
+public class Excel : SupportedFormat {
+    override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readExcel(stream)
+
+    override fun readDataFrame(file: File, header: List<String>): AnyFrame = DataFrame.readExcel(file)
+
+    override fun acceptsExtension(ext: String): Boolean = ext == "xls" || ext == "xlsx"
+
+    override val testOrder: Int = 40000
+}
+
 public fun DataFrame.Companion.readExcel(
     url: URL,
     sheetName: String? = null,

src/main/resources/META-INF/services/org.jetbrains.kotlinx.dataframe.io.SupportedFormat

@@ -0,0 +1,4 @@
+org.jetbrains.kotlinx.dataframe.io.CSV
+org.jetbrains.kotlinx.dataframe.io.Excel
+org.jetbrains.kotlinx.dataframe.io.JSON
+org.jetbrains.kotlinx.dataframe.io.TSV