Allow reading ViewVarCharVector and ViewVarBinaryVector from Arrow files

Kopilov · Sep 3, 2024 · 1d35299 · 1d35299
1 parent ea7dbe8
commit 1d35299
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowReadingImpl.kt b/dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowReadingImpl.kt
@@ -37,6 +37,8 @@ import org.apache.arrow.vector.UInt8Vector
 import org.apache.arrow.vector.VarBinaryVector
 import org.apache.arrow.vector.VarCharVector
 import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.ViewVarBinaryVector
+import org.apache.arrow.vector.ViewVarCharVector
 import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.ipc.ArrowFileReader
 import org.apache.arrow.vector.ipc.ArrowReader
@@ -214,6 +216,24 @@ private fun VarCharVector.values(range: IntRange): List<String?> =
         }
     }
 
+private fun LargeVarCharVector.values(range: IntRange): List<String?> =
+    range.map {
+        if (isNull(it)) {
+            null
+        } else {
+            String(get(it))
+        }
+    }
+
+private fun ViewVarCharVector.values(range: IntRange): List<String?> =
+    range.map {
+        if (isNull(it)) {
+            null
+        } else {
+            String(get(it))
+        }
+    }
+
 private fun VarBinaryVector.values(range: IntRange): List<ByteArray?> =
     range.map {
         if (isNull(it)) {
@@ -232,12 +252,12 @@ private fun LargeVarBinaryVector.values(range: IntRange): List<ByteArray?> =
         }
     }
 
-private fun LargeVarCharVector.values(range: IntRange): List<String?> =
+private fun ViewVarBinaryVector.values(range: IntRange): List<ByteArray?> =
     range.map {
         if (isNull(it)) {
             null
         } else {
-            String(get(it))
+            get(it)
         }
     }
 
@@ -273,10 +293,14 @@ private fun readField(root: VectorSchemaRoot, field: Field, nullability: Nullabi
 
             is LargeVarCharVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
 
+            is ViewVarCharVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
+
             is VarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
 
             is LargeVarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
 
+            is ViewVarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
+
             is BitVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
 
             is SmallIntVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)