Skip to content

Commit bf53ec1

Browse files
committed
Avoid 'expected: range(0, 32768)' exception when saving mixed data to Arrow
1 parent 8e1084d commit bf53ec1

File tree

1 file changed

+12
-2
lines changed

1 file changed

+12
-2
lines changed

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,16 @@ internal class ArrowWriterImpl(
145145
}
146146
}
147147

148+
private fun convertColumnToCompatible(column: AnyCol): Pair<AnyCol, Field> {
149+
val actualField = column.toArrowField(mismatchSubscriber)
150+
val result = try {
151+
convertColumnToTarget(column, actualField.type)!!
152+
} catch (e: Exception) {
153+
column
154+
}
155+
return result to actualField
156+
}
157+
148158
private fun infillVector(vector: FieldVector, column: AnyCol) {
149159
when (vector) {
150160
is VarCharVector ->
@@ -306,7 +316,7 @@ internal class ArrowWriterImpl(
306316
cause = e,
307317
),
308318
)
309-
column to column!!.toArrowField(mismatchSubscriber)
319+
convertColumnToCompatible(column!!)
310320
}
311321
} catch (e: TypeConverterNotFoundException) {
312322
if (strictType) {
@@ -317,7 +327,7 @@ internal class ArrowWriterImpl(
317327
} else {
318328
// If strictType is not enabled, use original data with its type. Target nullable is saved at this step.
319329
mismatchSubscriber(ConvertingMismatch.TypeConversionNotFound.ConversionNotFoundIgnored(field.name, e))
320-
column to column!!.toArrowField(mismatchSubscriber)
330+
convertColumnToCompatible(column!!)
321331
}
322332
}
323333

0 commit comments

Comments
 (0)