Merge pull request #350 from Kopilov/VariableWidthVector-fix

Variable width vector fix

Author: Jolanrensen

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

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

@@ -266,4 +266,11 @@ internal class ArrowKtTest {
             Locale.setDefault(currentLocale)
         }
     }
+
+    @Test
+    fun testBigStringColumn() {
+        val dataFrame = dataFrameOf(bigStringColumn)
+        val data = dataFrame.saveArrowFeatherToByteArray()
+        DataFrame.readArrowFeather(data) shouldBe dataFrame
+    }
 }

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

@@ -154,3 +154,39 @@ val citiesExampleSchema = """{
   } ]
 }
 """.trimIndent()
+
+/**
+ * String column (variable length vector) with size >1 MiB
+ */
+val bigStringColumn = run {
+    val list = ArrayList<String>()
+    for (i in 0 until 1024) {
+        val row = StringBuilder()
+        for (j in 0 until 64) {
+            row.append("abcd")
+        }
+        list.add(row.toString())
+    }
+    for (i in 0 until 1024) {
+        val row = StringBuilder()
+        for (j in 0 until 64) {
+            row.append("гдёж")
+        }
+        list.add(row.toString())
+    }
+    for (i in 0 until 1024) {
+        val row = StringBuilder()
+        for (j in 0 until 64) {
+            row.append("αβγδ")
+        }
+        list.add(row.toString())
+    }
+    for (i in 0 until 1024) {
+        val row = StringBuilder()
+        for (j in 0 until 64) {
+            row.append("正体字")
+        }
+        list.add(row.toString())
+    }
+    DataColumn.createValueColumn("bigStringColumn", list)
+}

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"