enabling column-parallel parsing with coroutines

Author: Jolanrensen

core/build.gradle.kts

@@ -73,6 +73,7 @@ dependencies {
     implementation(libs.fuel)
 
     api(libs.kotlin.datetimeJvm)
+    implementation(libs.kotlin.coroutinesCore)
     implementation(libs.kotlinpoet)
     implementation(libs.sl4j)
     implementation(libs.kotlinLogging)

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/parse.kt

@@ -58,7 +58,7 @@ public data class ParserOptions(
 public fun DataColumn<String?>.tryParse(options: ParserOptions? = null): DataColumn<*> = tryParseImpl(options)
 
 public fun <T> DataFrame<T>.parse(options: ParserOptions? = null): DataFrame<T> =
-    parse(options) {
+    parseImpl(options) {
         colsAtAnyDepth { !it.isColumnGroup() }
     }
 

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/parse.kt

@@ -1,5 +1,9 @@
 package org.jetbrains.kotlinx.dataframe.impl.api
 
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.runBlocking
 import kotlinx.datetime.Instant
 import kotlinx.datetime.LocalDate
 import kotlinx.datetime.LocalDateTime
@@ -18,13 +22,16 @@ import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.api.GlobalParserOptions
 import org.jetbrains.kotlinx.dataframe.api.ParserOptions
 import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
+import org.jetbrains.kotlinx.dataframe.api.asComparable
 import org.jetbrains.kotlinx.dataframe.api.asDataColumn
+import org.jetbrains.kotlinx.dataframe.api.asFrameColumn
 import org.jetbrains.kotlinx.dataframe.api.cast
-import org.jetbrains.kotlinx.dataframe.api.convert
+import org.jetbrains.kotlinx.dataframe.api.getColumns
 import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
 import org.jetbrains.kotlinx.dataframe.api.isFrameColumn
-import org.jetbrains.kotlinx.dataframe.api.parse
-import org.jetbrains.kotlinx.dataframe.api.to
+import org.jetbrains.kotlinx.dataframe.api.name
+import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.api.tryParse
 import org.jetbrains.kotlinx.dataframe.columns.size
 import org.jetbrains.kotlinx.dataframe.columns.values
@@ -37,6 +44,7 @@ import org.jetbrains.kotlinx.dataframe.impl.javaDurationCanParse
 import org.jetbrains.kotlinx.dataframe.io.isURL
 import org.jetbrains.kotlinx.dataframe.io.readJsonStr
 import org.jetbrains.kotlinx.dataframe.typeClass
+import org.jetbrains.kotlinx.dataframe.values
 import java.math.BigDecimal
 import java.net.URL
 import java.text.NumberFormat
@@ -541,19 +549,39 @@ internal fun <T> DataColumn<String?>.parse(parser: StringParser<T>, options: Par
     return DataColumn.createValueColumn(name(), parsedValues, parser.type.withNullability(hasNulls)) as DataColumn<T?>
 }
 
-internal fun <T> DataFrame<T>.parseImpl(options: ParserOptions?, columns: ColumnsSelector<T, Any?>) =
-    convert(columns).to {
-        when {
-            it.isFrameColumn() -> it.cast<AnyFrame?>().parse(options)
-
-            it.isColumnGroup() ->
-                it.asColumnGroup()
-                    .parse(options) { all() }
-                    .asColumnGroup(it.name())
-                    .asDataColumn()
-
-            it.typeClass == String::class -> it.cast<String?>().tryParse(options)
-
-            else -> it
-        }
+internal fun <T> DataFrame<T>.parseImpl(options: ParserOptions?, columns: ColumnsSelector<T, Any?>): DataFrame<T> =
+    runBlocking { parseParallel(options, columns) }
+
+private suspend fun <T> DataFrame<T>.parseParallel(
+    options: ParserOptions?,
+    columns: ColumnsSelector<T, Any?>,
+): DataFrame<T> =
+    coroutineScope {
+        getColumns(columns).map {
+            async {
+                when {
+                    it.isFrameColumn() ->
+                        it.asFrameColumn().values.map {
+                            async {
+                                it.parseParallel(options) {
+                                    colsAtAnyDepth { !it.isColumnGroup() }
+                                }
+                            }
+                        }.awaitAll()
+                            .toColumn(it.name)
+
+
+                    it.isColumnGroup() ->
+                        it.asColumnGroup().parseParallel(options) { all() }
+                            .asColumnGroup(it.name())
+                            .asDataColumn()
+
+                    it.typeClass == String::class -> it.cast<String?>().tryParse(options)
+
+                    else -> it
+                }
+            }
+        }.awaitAll()
+            .toDataFrame()
+            .cast()
     }

gradle/libs.versions.toml

@@ -31,6 +31,7 @@ postgresql = "42.7.4"
 sqlite = "3.46.1.0"
 jtsCore = "1.19.0"
 kotlinDatetime = "0.6.1"
+coroutines = "1.9.0"
 openapi = "2.1.22"
 kotlinLogging = "7.0.0"
 sl4j = "2.0.16"
@@ -84,6 +85,7 @@ jts = { group = "org.locationtech.jts", name = "jts-core", version.ref = "jtsCor
 
 poi-ooxml = { group = "org.apache.poi", name = "poi-ooxml", version.ref = "poi" }
 kotlin-datetimeJvm = { group = "org.jetbrains.kotlinx", name = "kotlinx-datetime-jvm", version.ref = "kotlinDatetime" }
+kotlin-coroutinesCore = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-core", version.ref = "coroutines" }
 
 junit = { group = "junit", name = "junit", version.ref = "junit" }