New test draft, keep nullables in schema

Author: Kopilov

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

@@ -3,6 +3,8 @@ package org.jetbrains.kotlinx.dataframe.io
 import org.apache.arrow.memory.RootAllocator
 import org.apache.arrow.vector.BigIntVector
 import org.apache.arrow.vector.BitVector
+import org.apache.arrow.vector.DateDayVector
+import org.apache.arrow.vector.DateMilliVector
 import org.apache.arrow.vector.Decimal256Vector
 import org.apache.arrow.vector.DecimalVector
 import org.apache.arrow.vector.DurationVector
@@ -28,16 +30,21 @@ import org.apache.arrow.vector.complex.StructVector
 import org.apache.arrow.vector.ipc.ArrowFileReader
 import org.apache.arrow.vector.ipc.ArrowStreamReader
 import org.apache.arrow.vector.types.pojo.Field
+import org.apache.arrow.vector.util.DateUtility
 import org.apache.commons.compress.utils.SeekableInMemoryByteChannel
 import org.jetbrains.kotlinx.dataframe.AnyBaseCol
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.Infer
-import org.jetbrains.kotlinx.dataframe.api.concat
+import org.jetbrains.kotlinx.dataframe.api.cast
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
+import org.jetbrains.kotlinx.dataframe.api.emptyDataFrame
+import org.jetbrains.kotlinx.dataframe.api.getColumn
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.codeGen.AbstractDefaultReadMethod
 import org.jetbrains.kotlinx.dataframe.codeGen.DefaultReadDfMethod
+import org.jetbrains.kotlinx.dataframe.impl.asList
 import java.io.File
 import java.io.InputStream
 import java.math.BigDecimal
@@ -48,6 +55,7 @@ import java.nio.channels.ReadableByteChannel
 import java.nio.channels.SeekableByteChannel
 import java.nio.file.Files
 import java.time.Duration
+import java.time.LocalDate
 import java.time.LocalDateTime
 import kotlin.reflect.typeOf
 
@@ -75,6 +83,25 @@ internal object Allocator {
     }
 }
 
+/**
+ * same as [Iterable<DataFrame<T>>.concat()] without internal type guessing (all batches should have the same schema)
+ */
+internal fun <T> Iterable<DataFrame<T>>.concatKeepingSchema(): DataFrame<T> {
+    val dataFrames = asList()
+    when (dataFrames.size) {
+        0 -> return emptyDataFrame()
+        1 -> return dataFrames[0]
+    }
+
+    val columnNames = dataFrames.first().columnNames()
+
+    val columns = columnNames.map { name ->
+        val values = dataFrames.flatMap { it.getColumn(name).values() }
+        DataColumn.createValueColumn(name, values, dataFrames.first().getColumn(name).type())
+    }
+    return dataFrameOf(columns).cast()
+}
+
 /**
  * Read [Arrow interprocess streaming format](https://arrow.apache.org/docs/java/ipc.html#writing-and-reading-streaming-format) data from existing [channel]
  */
@@ -88,7 +115,7 @@ public fun DataFrame.Companion.readArrowIPC(channel: ReadableByteChannel, alloca
                 add(df)
             }
         }
-        return dfs.concat()
+        return dfs.concatKeepingSchema()
     }
 }
 
@@ -106,7 +133,7 @@ public fun DataFrame.Companion.readArrowFeather(channel: SeekableByteChannel, al
                 add(df)
             }
         }
-        return dfs.concat()
+        return dfs.concatKeepingSchema()
     }
 }
 
@@ -129,6 +156,11 @@ private fun Float4Vector.values(range: IntRange): List<Float?> = range.map { get
 private fun Float8Vector.values(range: IntRange): List<Double?> = range.map { getObject(it) }
 
 private fun DurationVector.values(range: IntRange): List<Duration?> = range.map { getObject(it) }
+private fun DateDayVector.values(range: IntRange): List<LocalDate?> = range.map {
+    DateUtility.getLocalDateTimeFromEpochMilli(getObject(it).toLong() * DateUtility.daysToStandardMillis).toLocalDate()
+}
+private fun DateMilliVector.values(range: IntRange): List<LocalDateTime?> = range.map { getObject(it) }
+
 private fun TimeNanoVector.values(range: IntRange): List<Long?> = range.map { getObject(it) }
 private fun TimeMicroVector.values(range: IntRange): List<Long?> = range.map { getObject(it) }
 private fun TimeMilliVector.values(range: IntRange): List<LocalDateTime?> = range.map { getObject(it) }
@@ -190,6 +222,8 @@ private fun readField(root: VectorSchemaRoot, field: Field): AnyBaseCol {
         is Float8Vector -> vector.values(range).withType()
         is Float4Vector -> vector.values(range).withType()
         is DurationVector -> vector.values(range).withType()
+        is DateDayVector -> vector.values(range).withType()
+        is DateMilliVector -> vector.values(range).withType()
         is TimeNanoVector -> vector.values(range).withType()
         is TimeMicroVector -> vector.values(range).withType()
         is TimeMilliVector -> vector.values(range).withType()
@@ -199,7 +233,7 @@ private fun readField(root: VectorSchemaRoot, field: Field): AnyBaseCol {
             TODO("not fully implemented")
         }
     }
-    return DataColumn.createValueColumn(field.name, list, type, Infer.Nulls)
+    return DataColumn.createValueColumn(field.name, list, type, Infer.None)
 }
 
 // IPC reading block

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

@@ -1,34 +1,44 @@
 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.columnOf
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.toColumn
 import org.jetbrains.kotlinx.dataframe.io.readArrowFeather
+import org.jetbrains.kotlinx.dataframe.io.readArrowIPC
 import org.junit.Test
 import java.net.URL
+import kotlin.reflect.typeOf
 
 internal class ArrowKtTest {
 
     fun testResource(resourcePath: String): URL = ArrowKtTest::class.java.classLoader.getResource(resourcePath)!!
 
     fun testArrowFeather(name: String) = testResource("$name.feather")
+    fun testArrowIPC(name: String) = testResource("$name.ipc")
 
     @Test
     fun testReadingFromFile() {
         val feather = testArrowFeather("data-arrow_2.0.0_uncompressed")
         val df = DataFrame.readArrowFeather(feather)
-        val a by columnOf("one")
-        val b by columnOf(2.0)
+        val a by listOf("one" as String?).toColumn()
+        val b by listOf(2.0 as Double?).toColumn()
         val c by listOf(
             mapOf(
                 "c1" to Text("inner"),
                 "c2" to 4.0,
                 "c3" to 50.0
-            ) as Map<String, Any?>
+            ) as Map<String, Any?>?
         ).toColumn()
-        val d by columnOf("four")
+        val d by listOf("four" as String?).toColumn()
         val expected = dataFrameOf(a, b, c, d)
         df shouldBe expected
     }
+
+    @Test
+    fun testReadingAllTypesAsEstimated() {
+        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test.arrow")))
+        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test.arrow")))
+    }
 }

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

@@ -0,0 +1,48 @@
+import io.kotest.matchers.shouldBe
+import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.DataColumn
+import org.jetbrains.kotlinx.dataframe.api.forEach
+import org.jetbrains.kotlinx.dataframe.api.forEachIndexed
+import kotlin.reflect.typeOf
+
+/**
+ * Assert that we have got the same data that was originally saved on example creation.
+ */
+internal fun assertEstimations(exampleFrame: AnyFrame) {
+    /**
+     * In [exampleFrame] we get two concatenated batches. To assert the estimations, we should transform frame row number to batch row number
+     */
+    fun iBatch(iFrame: Int): Int {
+        val firstBatchSize = 100;
+        return if (iFrame < firstBatchSize) iFrame else iFrame - firstBatchSize
+    }
+    val asciiStringCol = exampleFrame["asciiString"] as DataColumn<String?>
+    asciiStringCol.type() shouldBe typeOf<String?>()
+    asciiStringCol.forEachIndexed { i, element ->
+        element shouldBe "Test Example ${iBatch(i)}"
+    }
+
+    val utf8StringCol = exampleFrame["utf8String"]
+    val largeStringCol = exampleFrame["largeString"]
+
+    val booleanCol = exampleFrame["boolean"]
+
+    val byteCol = exampleFrame["byte"]
+    val shortCol = exampleFrame["short"]
+    val intCol = exampleFrame["int"]
+    val longIntCol = exampleFrame["longInt"]
+
+    val unsignedByteCol = exampleFrame["unsigned_byte"]
+    val unsignedShortCol = exampleFrame["unsigned_short"]
+    val unsignedIntCol = exampleFrame["unsigned_int"]
+    val unsignedLongIntCol = exampleFrame["unsigned_longInt"]
+
+    val dateCol = exampleFrame["date32"]
+    val datetimeCol = exampleFrame["date64"]
+
+    val timeSecCol = exampleFrame["time32_seconds"]
+    val timeMilliCol = exampleFrame["time32_milli"]
+
+    val timeMicroCol = exampleFrame["time64_micro"]
+    val timeNanoCol = exampleFrame["time64_nano"]
+}