Avoid 'expected: range(0, 32768)' exception when saving mixed data to…

… Arrow
Kopilov · Sep 3, 2024 · bc8da27 · bc8da27
1 parent ea7dbe8
commit bc8da27
Showing 1 changed file with 12 additions and 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
@@ -145,6 +145,16 @@ internal class ArrowWriterImpl(
         }
     }
 
+    private fun convertColumnToCompatible(column: AnyCol): Pair<AnyCol, Field> {
+        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!!)
             }
         }