From bc8da27a5598ebd9ac82acc24f085a60cb3845ab 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 9c658ffabf..ff3e7ea064 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 @@ -145,6 +145,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 -> @@ -306,7 +316,7 @@ internal class ArrowWriterImpl( cause = e, ), ) - column to column!!.toArrowField(mismatchSubscriber) + convertColumnToCompatible(column!!) } } catch (e: TypeConverterNotFoundException) { if (strictType) { @@ -317,7 +327,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!!) } }