Fix allocating VariableWidthVector in Arrow

Kotlin · Apr 20, 2023 · c3b9939 · c3b9939
1 parent 13f3c4b
commit c3b9939
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowWriterImpl.kt b/dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowWriterImpl.kt
@@ -53,6 +53,10 @@ import org.jetbrains.kotlinx.dataframe.api.map
 import org.jetbrains.kotlinx.dataframe.exceptions.CellConversionException
 import org.jetbrains.kotlinx.dataframe.exceptions.TypeConverterNotFoundException
 import org.jetbrains.kotlinx.dataframe.name
+import org.jetbrains.kotlinx.dataframe.values
+import java.nio.charset.Charset
+import kotlin.reflect.full.isSubtypeOf
+import kotlin.reflect.typeOf
 
 /**
  * Save [dataFrame] content in Apache Arrow format (can be written to File, ByteArray, OutputStream or raw Channel) with [targetSchema].
@@ -67,14 +71,25 @@ internal class ArrowWriterImpl(
 
     private val allocator = RootAllocator()
 
-    private fun allocateVector(vector: FieldVector, size: Int) {
+    private fun allocateVector(vector: FieldVector, size: Int, totalBytes: Long? = null) {
         when (vector) {
             is FixedWidthVector -> vector.allocateNew(size)
-            is VariableWidthVector -> vector.allocateNew(size)
+            is VariableWidthVector -> totalBytes?.let { vector.allocateNew(it, size) } ?: vector.allocateNew(size)
             else -> throw IllegalArgumentException("Can not allocate ${vector.javaClass.canonicalName}")
         }
     }
 
+    /**
+     * Calculate buffer size for VariableWidthVector (return null for FixedWidthVector)
+     */
+    private fun countTotalBytes(column: AnyCol): Long? {
+        val columnType = column.type()
+        return when {
+            columnType.isSubtypeOf(typeOf<String?>()) -> column.values.fold(0L) {totalBytes, value -> totalBytes + value.toString().length * 4}
+            else -> null
+        }
+    }
+
     private fun infillWithNulls(vector: FieldVector, size: Int) {
         when (vector) {
             is BaseFixedWidthVector -> for (i in 0 until size) { vector.setNull(i) }
@@ -189,11 +204,12 @@ internal class ArrowWriterImpl(
             actualField.createVector(allocator)!!
         }
 
-        allocateVector(vector, dataFrame.rowsCount())
         if (convertedColumn == null) {
             check(actualField.isNullable)
+            allocateVector(vector, dataFrame.rowsCount())
             infillWithNulls(vector, dataFrame.rowsCount())
         } else {
+            allocateVector(vector, dataFrame.rowsCount(), countTotalBytes(convertedColumn))
             infillVector(vector, convertedColumn)
         }
         return vector

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -23,7 +23,7 @@ openapi = "2.1.13"
 junit = "4.13.2"
 kotestAsserions = "4.6.3"
 jsoup = "1.14.3"
-arrow = "10.0.0"
+arrow = "11.0.0"
 docProcessor = "0.1.6"
 simpleGit = "2.0.1"