testBigMixedColumn

Kopilov · Sep 4, 2024 · 62d0459 · 62d0459
1 parent bc8da27
commit 62d0459
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowKtTest.kt b/dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ArrowKtTest.kt
@@ -498,6 +498,28 @@ 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(

diff --git a/dataframe-arrow/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/examplesToWrite.kt b/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)
+}