Merge pull request #853 from Kopilov/convert-comparable-to-string

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

Author: koperagen

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!!)
             }
         }
 

dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowKtTest.kt

@@ -498,6 +498,30 @@ internal class ArrowKtTest {
         DataFrame.readArrowFeather(data) shouldBe dataFrame
     }
 
+    @Test
+    fun testBigMixedColumn() {
+        val dataFrame = dataFrameOf(bigMixedColumn)
+        val warnings = ArrayList<ConvertingMismatch>()
+        val writer = dataFrame.arrowWriter(
+            targetSchema = Schema(
+                listOf(
+                    Field("bigMixedColumn", FieldType.nullable(ArrowType.Int(64, true)), emptyList()),
+                ),
+            ),
+            mode = ArrowWriter.Mode.LOYAL,
+        ) {
+            warnings.add(it)
+        }
+        val stream = ByteArrayOutputStream()
+        writer.writeArrowFeather(stream)
+        val data = stream.toByteArray()
+
+        assert(warnings.filterIsInstance<ConvertingMismatch.TypeConversionFail.ConversionFailIgnored>().size == 1)
+        assert(warnings.filterIsInstance<ConvertingMismatch.SavedAsString>().size == 1)
+
+        DataFrame.readArrowFeather(data)["bigMixedColumn"] shouldBe dataFrame[bigMixedColumn].map { it.toString() }
+    }
+
     @Test
     fun testTimeStamp() {
         val dates = listOf(

dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/examplesToWrite.kt

@@ -192,3 +192,15 @@ val bigStringColumn = run {
     }
     DataColumn.createValueColumn("bigStringColumn", list)
 }
+
+val bigMixedColumn = run {
+    val list = ArrayList<Any>()
+    for (i in 0..32768) {
+        list.add(i * i)
+    }
+    list.add("Dirty data")
+    for (i in 32768 downTo 0) {
+        list.add(i * i)
+    }
+    DataColumn.createValueColumn("bigMixedColumn", list)
+}