Save to Arrow test

Author: Kopilov

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

@@ -245,6 +245,9 @@ public class ArrowWriter(
 
     }
 
+    /**
+     * Create Arrow FieldVector with [column] content cast to [field] type according to [strictType] and [strictNullable] settings.
+     */
     private fun allocateVectorAndInfill(field: Field, column: AnyCol?, strictType: Boolean, strictNullable: Boolean): FieldVector {
         val containNulls = (column == null || column.hasNulls())
         // Convert the column to type specified in field. (If we already have target type, convertTo will do nothing)
@@ -264,7 +267,7 @@ public class ArrowWriter(
         }
         val vector = if (!actualField.isNullable && containNulls) {
             if (strictNullable) {
-                throw Exception("${actualField.name} column contains nulls but should be not nullable")
+                throw IllegalArgumentException("${actualField.name} column contains nulls but should be not nullable")
             } else {
                 warningSubscriber("${actualField.name} column contains nulls but expected not nullable")
                 Field(actualField.name, FieldType(true, actualField.fieldType.type, actualField.fieldType.dictionary), actualField.children).createVector(allocator)!!
@@ -296,7 +299,7 @@ public class ArrowWriter(
             val column = dataFrame.getColumnOrNull(field.name)
             if (column == null && !field.isNullable) {
                 if (mode.restrictNarrowing) {
-                    throw Exception("${field.name} column is not presented")
+                    throw IllegalArgumentException("${field.name} column is not presented")
                 } else {
                     warningSubscriber("${field.name} column is not presented")
                     continue

dataframe-arrow/src/test/kotlin/ArrowKtTest.kt

@@ -1,15 +1,19 @@
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
 import org.apache.arrow.vector.util.Text
+import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.NullabilityOptions
 import org.jetbrains.kotlinx.dataframe.api.columnOf
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.io.arrowWriter
 import org.jetbrains.kotlinx.dataframe.io.readArrowFeather
 import org.jetbrains.kotlinx.dataframe.io.readArrowIPC
 import org.junit.Test
+import java.io.File
 import java.net.URL
+import java.time.LocalDate
 
 internal class ArrowKtTest {
 
@@ -87,4 +91,58 @@ internal class ArrowKtTest {
         assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test-illegal.arrow"), NullabilityOptions.Widening), true, true)
         assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test-illegal.arrow"), NullabilityOptions.Widening), true, true)
     }
+
+    val cities = dataFrameOf(
+        DataColumn.createValueColumn("name", listOf(
+            "Berlin",
+            "Hamburg",
+            "New York",
+            "Washington",
+            "Saint Petersburg"
+        )),
+        DataColumn.createValueColumn("affiliation", listOf(
+            "Germany",
+            "Germany",
+            "The USA",
+            "The USA",
+            "Russia"
+        )),
+        DataColumn.createValueColumn("is_capital", listOf(
+            true,
+            false,
+            false,
+            true,
+            false
+        )),
+        DataColumn.createValueColumn("population", listOf(
+            3_769_495,
+            1_845_229,
+            8_467_513,
+            689_545,
+            5_377_503
+        )),
+        DataColumn.createValueColumn("area", listOf(
+            891.7,
+            755.22,
+            1223.59,
+            177.0,
+            1439.0
+        )),
+//        DataColumn.createValueColumn("settled", listOf(
+//            LocalDate.of(1237, 1, 1),
+//            LocalDate.of(1189, 5, 7),
+//            LocalDate.of(1624, 1, 1),
+//            LocalDate.of(1790, 7, 16),
+//            LocalDate.of(1703, 5, 27)
+//        ))
+    )
+
+    @Test
+    fun testWritingGeneral() {
+        val testFile = File.createTempFile("cities", "arrow")
+        cities.arrowWriter().writeArrowFeather(testFile)
+        DataFrame.readArrowFeather(testFile).shouldBe(cities)
+    }
+
+
 }