ColumnDataCollector now uses ColumnDataHolder for collecting data dir…

…ectly into primitive arrays
Kotlin · Aug 20, 2024 · 63d5e0a · 63d5e0a
1 parent 1ceb69f
commit 63d5e0a
Show file tree

Hide file tree

Showing 10 changed files with 1,157 additions and 545 deletions.
diff --git a/core/build.gradle.kts b/core/build.gradle.kts
@@ -83,7 +83,8 @@ dependencies {
     testImplementation(libs.kotlin.scriptingJvm)
     testImplementation(libs.jsoup)
 
-    testImplementation("org.openjdk.jol:jol-core:0.10")
+//    testImplementation("org.openjdk.jol:jol-core:0.10")
+    implementation("org.openjdk.jol:jol-core:0.10")
     implementation("it.unimi.dsi:fastutil:8.5.14")
 }
 

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnDataHolder.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/ColumnDataHolder.kt
@@ -31,7 +31,9 @@ public interface ColumnDataHolder<T> : List<T> {
 
     public operator fun get(range: IntRange): List<T>
 
-    public fun add(element: T): Boolean
+    public fun add(element: T)
+
+    public fun canAdd(element: T): Boolean
 
     public val distinct: Lazy<Set<T>>
 

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ColumnDataCollector.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/ColumnDataCollector.kt
@@ -2,11 +2,14 @@ package org.jetbrains.kotlinx.dataframe.impl
 
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
+import org.jetbrains.kotlinx.dataframe.ColumnDataHolder
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.impl.columns.empty
+import org.jetbrains.kotlinx.dataframe.impl.columns.emptyForType
 import org.jetbrains.kotlinx.dataframe.impl.columns.guessColumnType
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
@@ -28,7 +31,7 @@ internal abstract class DataCollectorBase<T>(initCapacity: Int) : DataCollector<
 
     override var hasNulls = false
 
-    override val data = ArrayList<T?>(initCapacity)
+    override val data = ColumnDataHolder.empty<T>(initCapacity)
 
     val values: List<T?>
         get() = data
@@ -62,8 +65,15 @@ internal class TypedColumnDataCollector<T>(initCapacity: Int = 0, val type: KTyp
 
     internal val kclass = type.jvmErasure
 
+    override val data: ColumnDataHolder<T?> =
+        ColumnDataHolder.emptyForType(
+            type = type,
+            initCapacity = initCapacity,
+            strictTypes = checkTypes,
+        )
+
     override fun add(value: T?) {
-        if (checkTypes && value != null && !value.javaClass.kotlin.isSubclassOf(kclass)) {
+        if (checkTypes && data.canAdd(value) && value != null && !value.javaClass.kotlin.isSubclassOf(kclass)) {
             throw IllegalArgumentException(
                 "Can not add value of class ${value.javaClass.kotlin.qualifiedName} to column of type $type. Value = $value",
             )