Skip to content

Commit 0a5f526

Browse files
authored
Merge pull request #852 from Kopilov/ViewVarCharVector
Allow reading ViewVarCharVector and ViewVarBinaryVector from Arrow files
2 parents ecf051c + 1d35299 commit 0a5f526

File tree

1 file changed

+26
-2
lines changed

1 file changed

+26
-2
lines changed

dataframe-arrow/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/arrowReadingImpl.kt

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ import org.apache.arrow.vector.UInt8Vector
3737
import org.apache.arrow.vector.VarBinaryVector
3838
import org.apache.arrow.vector.VarCharVector
3939
import org.apache.arrow.vector.VectorSchemaRoot
40+
import org.apache.arrow.vector.ViewVarBinaryVector
41+
import org.apache.arrow.vector.ViewVarCharVector
4042
import org.apache.arrow.vector.complex.StructVector
4143
import org.apache.arrow.vector.ipc.ArrowFileReader
4244
import org.apache.arrow.vector.ipc.ArrowReader
@@ -214,6 +216,24 @@ private fun VarCharVector.values(range: IntRange): List<String?> =
214216
}
215217
}
216218

219+
private fun LargeVarCharVector.values(range: IntRange): List<String?> =
220+
range.map {
221+
if (isNull(it)) {
222+
null
223+
} else {
224+
String(get(it))
225+
}
226+
}
227+
228+
private fun ViewVarCharVector.values(range: IntRange): List<String?> =
229+
range.map {
230+
if (isNull(it)) {
231+
null
232+
} else {
233+
String(get(it))
234+
}
235+
}
236+
217237
private fun VarBinaryVector.values(range: IntRange): List<ByteArray?> =
218238
range.map {
219239
if (isNull(it)) {
@@ -232,12 +252,12 @@ private fun LargeVarBinaryVector.values(range: IntRange): List<ByteArray?> =
232252
}
233253
}
234254

235-
private fun LargeVarCharVector.values(range: IntRange): List<String?> =
255+
private fun ViewVarBinaryVector.values(range: IntRange): List<ByteArray?> =
236256
range.map {
237257
if (isNull(it)) {
238258
null
239259
} else {
240-
String(get(it))
260+
get(it)
241261
}
242262
}
243263

@@ -273,10 +293,14 @@ private fun readField(root: VectorSchemaRoot, field: Field, nullability: Nullabi
273293

274294
is LargeVarCharVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
275295

296+
is ViewVarCharVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
297+
276298
is VarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
277299

278300
is LargeVarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
279301

302+
is ViewVarBinaryVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
303+
280304
is BitVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)
281305

282306
is SmallIntVector -> vector.values(range).withTypeNullable(field.isNullable, nullability)

0 commit comments

Comments
 (0)