From b0300e69ed95cbc9d9a8388e1457e788561fd363 Mon Sep 17 00:00:00 2001 From: Kopilov Aleksandr Date: Mon, 2 Sep 2024 17:04:11 +0300 Subject: [PATCH] Avoid 'expected: range(0, 32768)' exception when saving mixed data to Arrow --- .../kotlinx/dataframe/io/ArrowWriterImpl.kt | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 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 index e1e56061cc..b91835c435 100644 --- 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 @@ -146,6 +146,16 @@ internal class ArrowWriterImpl( } } + private fun convertColumnToCompatible(column: AnyCol): Pair { + val actualField = column.toArrowField(mismatchSubscriber) + val result = try { + convertColumnToTarget(column, actualField.type)!! + } catch (e: Exception) { + column + } + return result to actualField + } + private fun infillVector(vector: FieldVector, column: AnyCol) { when (vector) { is VarCharVector -> @@ -307,7 +317,7 @@ internal class ArrowWriterImpl( cause = e, ), ) - column to column!!.toArrowField(mismatchSubscriber) + convertColumnToCompatible(column!!) } } catch (e: TypeConverterNotFoundException) { if (strictType) { @@ -318,7 +328,7 @@ internal class ArrowWriterImpl( } else { // If strictType is not enabled, use original data with its type. Target nullable is saved at this step. mismatchSubscriber(ConvertingMismatch.TypeConversionNotFound.ConversionNotFoundIgnored(field.name, e)) - column to column!!.toArrowField(mismatchSubscriber) + convertColumnToCompatible(column!!) } }