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" }