Nullable checking beta

Author: Kopilov

dataframe-arrow/build.gradle.kts

@@ -12,6 +12,7 @@ dependencies {
     implementation(libs.arrow.format)
     implementation(libs.arrow.memory)
     implementation(libs.commonsCompress)
+    implementation(libs.kotlin.reflect)
 
     testApi(project(":core"))
     testImplementation(libs.junit)

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

@@ -58,6 +58,8 @@ import java.time.Duration
 import java.time.LocalDate
 import java.time.LocalDateTime
 import java.time.LocalTime
+import kotlin.reflect.KType
+import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
 public class ArrowFeather : SupportedFormat {
@@ -141,7 +143,7 @@ public fun DataFrame.Companion.readArrowFeather(channel: SeekableByteChannel, al
 private fun BitVector.values(range: IntRange): List<Boolean?> = range.map { getObject(it) }
 
 private fun UInt1Vector.values(range: IntRange): List<Short?> = range.map { getObjectNoOverflow(it) }
-private fun UInt2Vector.values(range: IntRange): List<Int?> = range.map { getObject(it).code }
+private fun UInt2Vector.values(range: IntRange): List<Int?> = range.map { getObject(it)?.code }
 private fun UInt4Vector.values(range: IntRange): List<Long?> = range.map { getObjectNoOverflow(it) }
 private fun UInt8Vector.values(range: IntRange): List<BigInteger?> = range.map { getObjectNoOverflow(it) }
 
@@ -158,14 +160,33 @@ private fun Float8Vector.values(range: IntRange): List<Double?> = range.map { ge
 
 private fun DurationVector.values(range: IntRange): List<Duration?> = range.map { getObject(it) }
 private fun DateDayVector.values(range: IntRange): List<LocalDate?> = range.map {
+    if (getObject(it) == null) null else
     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<LocalTime?> = range.map { LocalTime.ofNanoOfDay(get(it)) }
-private fun TimeMicroVector.values(range: IntRange): List<LocalTime?> = range.map { LocalTime.ofNanoOfDay(get(it) * 1000) }
-private fun TimeMilliVector.values(range: IntRange): List<LocalTime?> = range.map { LocalTime.ofNanoOfDay(get(it).toLong() * 1000_000) }
-private fun TimeSecVector.values(range: IntRange): List<LocalTime?> = range.map { LocalTime.ofSecondOfDay(get(it).toLong()) }
+private fun TimeNanoVector.values(range: IntRange): List<LocalTime?> = range.mapIndexed { i, it ->
+    if (isNull(i)) {
+        null
+    } else {
+        LocalTime.ofNanoOfDay(get(it))
+    }
+}
+private fun TimeMicroVector.values(range: IntRange): List<LocalTime?> = range.mapIndexed { i, it ->
+    if (isNull(i)) {
+        null
+    } else {
+        LocalTime.ofNanoOfDay(getObject(it) * 1000)
+    }
+}
+private fun TimeMilliVector.values(range: IntRange): List<LocalTime?> = range.mapIndexed { i, it ->
+    if (isNull(i)) {
+        null
+    } else {
+        LocalTime.ofNanoOfDay(get(it).toLong() * 1000_000)
+    }
+}
+private fun TimeSecVector.values(range: IntRange): List<LocalTime?> = range.map { getObject(it)?.let {LocalTime.ofSecondOfDay(it.toLong())}  }
 
 private fun StructVector.values(range: IntRange): List<Map<String, Any?>?> = range.map { getObject(it) }
 
@@ -201,36 +222,36 @@ private fun LargeVarCharVector.values(range: IntRange): List<String?> = range.ma
     }
 }
 
-private inline fun <reified T> List<T>.withType() = this to typeOf<T>()
+private inline fun <reified T> List<T>.withType(nullability: Boolean) = this to typeOf<T>().withNullability(nullability)
 
 private fun readField(root: VectorSchemaRoot, field: Field): AnyBaseCol {
     val range = 0 until root.rowCount
     val (list, type) = when (val vector = root.getVector(field)) {
-        is VarCharVector -> vector.values(range).withType()
-        is LargeVarCharVector -> vector.values(range).withType()
-        is VarBinaryVector -> vector.values(range).withType()
-        is LargeVarBinaryVector -> vector.values(range).withType()
-        is BitVector -> vector.values(range).withType()
-        is SmallIntVector -> vector.values(range).withType()
-        is TinyIntVector -> vector.values(range).withType()
-        is UInt1Vector -> vector.values(range).withType()
-        is UInt2Vector -> vector.values(range).withType()
-        is UInt4Vector -> vector.values(range).withType()
-        is UInt8Vector -> vector.values(range).withType()
-        is IntVector -> vector.values(range).withType()
-        is BigIntVector -> vector.values(range).withType()
-        is DecimalVector -> vector.values(range).withType()
-        is Decimal256Vector -> vector.values(range).withType()
-        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()
-        is TimeSecVector -> vector.values(range).withType()
-        is StructVector -> vector.values(range).withType()
+        is VarCharVector -> vector.values(range).withType(field.isNullable)
+        is LargeVarCharVector -> vector.values(range).withType(field.isNullable)
+        is VarBinaryVector -> vector.values(range).withType(field.isNullable)
+        is LargeVarBinaryVector -> vector.values(range).withType(field.isNullable)
+        is BitVector -> vector.values(range).withType(field.isNullable)
+        is SmallIntVector -> vector.values(range).withType(field.isNullable)
+        is TinyIntVector -> vector.values(range).withType(field.isNullable)
+        is UInt1Vector -> vector.values(range).withType(field.isNullable)
+        is UInt2Vector -> vector.values(range).withType(field.isNullable)
+        is UInt4Vector -> vector.values(range).withType(field.isNullable)
+        is UInt8Vector -> vector.values(range).withType(field.isNullable)
+        is IntVector -> vector.values(range).withType(field.isNullable)
+        is BigIntVector -> vector.values(range).withType(field.isNullable)
+        is DecimalVector -> vector.values(range).withType(field.isNullable)
+        is Decimal256Vector -> vector.values(range).withType(field.isNullable)
+        is Float8Vector -> vector.values(range).withType(field.isNullable)
+        is Float4Vector -> vector.values(range).withType(field.isNullable)
+        is DurationVector -> vector.values(range).withType(field.isNullable)
+        is DateDayVector -> vector.values(range).withType(field.isNullable)
+        is DateMilliVector -> vector.values(range).withType(field.isNullable)
+        is TimeNanoVector -> vector.values(range).withType(field.isNullable)
+        is TimeMicroVector -> vector.values(range).withType(field.isNullable)
+        is TimeMilliVector -> vector.values(range).withType(field.isNullable)
+        is TimeSecVector -> vector.values(range).withType(field.isNullable)
+        is StructVector -> vector.values(range).withType(field.isNullable)
         else -> {
             TODO("not fully implemented")
         }

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

@@ -38,7 +38,25 @@ internal class ArrowKtTest {
 
     @Test
     fun testReadingAllTypesAsEstimated() {
-        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test.arrow")))
-        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test.arrow")))
+        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test.arrow")), true, false)
+        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test.arrow")), true, false)
+    }
+
+    @Test
+    fun testReadingAllTypesAsEstimatedWithNulls() {
+        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test-with-nulls.arrow")), true, true)
+        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test-with-nulls.arrow")), true, true)
+    }
+
+    @Test
+    fun testReadingAllTypesAsEstimatedNotNullable() {
+        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test-not-nullable.arrow")), false, false)
+        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test-not-nullable.arrow")), false, false)
+    }
+
+    @Test
+    fun testReadingAllTypesAsEstimatedNotNullableWithNulls() {
+        assertEstimations(DataFrame.readArrowFeather(testArrowFeather("test-illegal.arrow")), false, true)
+        assertEstimations(DataFrame.readArrowIPC(testArrowIPC("test-illegal.arrow")), false, true)
     }
 }

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

@@ -9,137 +9,151 @@ import java.time.LocalTime
 import java.time.ZoneOffset
 import kotlin.math.absoluteValue
 import kotlin.math.pow
+import kotlin.reflect.full.withNullability
 import kotlin.reflect.typeOf
 
 /**
  * Assert that we have got the same data that was originally saved on example creation.
  */
-internal fun assertEstimations(exampleFrame: AnyFrame) {
+internal fun assertEstimations(exampleFrame: AnyFrame, nullable: Boolean, withNulls: Boolean) {
     /**
      * 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
     }
+
+    fun expectedNull(rowNumber: Int): Boolean {
+        return (rowNumber + 1) % 5 == 0;
+    }
+
+    fun assertValueOrNull(rowNumber: Int, actual: Any?, expected: Any) {
+        if (withNulls && expectedNull(rowNumber)) {
+            actual shouldBe null
+        } else {
+            actual shouldBe expected
+        }
+    }
+
     val asciiStringCol = exampleFrame["asciiString"] as DataColumn<String?>
-    asciiStringCol.type() shouldBe typeOf<String?>()
+    asciiStringCol.type() shouldBe typeOf<String>().withNullability(nullable)
     asciiStringCol.forEachIndexed { i, element ->
-        element shouldBe "Test Example ${iBatch(i)}"
+        assertValueOrNull(iBatch(i), element, "Test Example ${iBatch(i)}")
     }
 
     val utf8StringCol = exampleFrame["utf8String"] as DataColumn<String?>
-    utf8StringCol.type() shouldBe typeOf<String?>()
+    utf8StringCol.type() shouldBe typeOf<String>().withNullability(nullable)
     utf8StringCol.forEachIndexed { i, element ->
-        element shouldBe "Тестовый пример ${iBatch(i)}"
+        assertValueOrNull(iBatch(i), element, "Тестовый пример ${iBatch(i)}")
     }
 
     val largeStringCol = exampleFrame["largeString"] as DataColumn<String?>
-    largeStringCol.type() shouldBe typeOf<String?>()
+    largeStringCol.type() shouldBe typeOf<String>().withNullability(nullable)
     largeStringCol.forEachIndexed { i, element ->
-        element shouldBe "Test Example Should Be Large ${iBatch(i)}"
+        assertValueOrNull(iBatch(i), element, "Test Example Should Be Large ${iBatch(i)}")
     }
 
     val booleanCol = exampleFrame["boolean"] as DataColumn<Boolean?>
-    booleanCol.type() shouldBe typeOf<Boolean?>()
+    booleanCol.type() shouldBe typeOf<Boolean>().withNullability(nullable)
     booleanCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) % 2 == 0)
+        assertValueOrNull(iBatch(i), element, iBatch(i) % 2 == 0)
     }
 
     val byteCol = exampleFrame["byte"] as DataColumn<Byte?>
-    byteCol.type() shouldBe typeOf<Byte?>()
+    byteCol.type() shouldBe typeOf<Byte>().withNullability(nullable)
     byteCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) * 10).toByte()
+        assertValueOrNull(iBatch(i), element, (iBatch(i) * 10).toByte())
     }
 
     val shortCol = exampleFrame["short"] as DataColumn<Short?>
-    shortCol.type() shouldBe typeOf<Short?>()
+    shortCol.type() shouldBe typeOf<Short>().withNullability(nullable)
     shortCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) * 1000).toShort()
+        assertValueOrNull(iBatch(i), element, (iBatch(i) * 1000).toShort())
     }
 
     val intCol = exampleFrame["int"] as DataColumn<Int?>
-    intCol.type() shouldBe typeOf<Int?>()
+    intCol.type() shouldBe typeOf<Int>().withNullability(nullable)
     intCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) * 100000000).toInt()
+        assertValueOrNull(iBatch(i), element, iBatch(i) * 100000000)
     }
 
     val longCol = exampleFrame["longInt"] as DataColumn<Long?>
-    longCol.type() shouldBe typeOf<Long?>()
+    longCol.type() shouldBe typeOf<Long>().withNullability(nullable)
     longCol.forEachIndexed { i, element ->
-        element shouldBe iBatch(i) * 100000000000000000L
+        assertValueOrNull(iBatch(i), element, iBatch(i) * 100000000000000000L)
     }
 
     val unsignedByteCol = exampleFrame["unsigned_byte"] as DataColumn<Short?>
-    unsignedByteCol.type() shouldBe typeOf<Short?>()
+    unsignedByteCol.type() shouldBe typeOf<Short>().withNullability(nullable)
     unsignedByteCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) * 10 % (Byte.MIN_VALUE.toShort() * 2).absoluteValue).toShort()
+        assertValueOrNull(iBatch(i), element, (iBatch(i) * 10 % (Byte.MIN_VALUE.toShort() * 2).absoluteValue).toShort())
     }
 
     val unsignedShortCol = exampleFrame["unsigned_short"] as DataColumn<Int?>
-    unsignedShortCol.type() shouldBe typeOf<Int?>()
+    unsignedShortCol.type() shouldBe typeOf<Int>().withNullability(nullable)
     unsignedShortCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i) * 1000 % (Short.MIN_VALUE.toInt() * 2).absoluteValue)
+        assertValueOrNull(iBatch(i), element, iBatch(i) * 1000 % (Short.MIN_VALUE.toInt() * 2).absoluteValue)
     }
 
     val unsignedIntCol = exampleFrame["unsigned_int"] as DataColumn<Long?>
-    unsignedIntCol.type() shouldBe typeOf<Long?>()
+    unsignedIntCol.type() shouldBe typeOf<Long>().withNullability(nullable)
     unsignedIntCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i).toLong() * 100000000 % (Int.MIN_VALUE.toLong() * 2).absoluteValue)
+        assertValueOrNull(iBatch(i), element, iBatch(i).toLong() * 100000000 % (Int.MIN_VALUE.toLong() * 2).absoluteValue)
     }
 
     val unsignedLongIntCol = exampleFrame["unsigned_longInt"] as DataColumn<BigInteger?>
-    unsignedLongIntCol.type() shouldBe typeOf<BigInteger?>()
+    unsignedLongIntCol.type() shouldBe typeOf<BigInteger>().withNullability(nullable)
     unsignedLongIntCol.forEachIndexed { i, element ->
-        element shouldBe (iBatch(i).toBigInteger() * 100000000000000000L.toBigInteger() % (Long.MIN_VALUE.toBigInteger() * 2.toBigInteger()).abs())
+        assertValueOrNull(iBatch(i), element, iBatch(i).toBigInteger() * 100000000000000000L.toBigInteger() % (Long.MIN_VALUE.toBigInteger() * 2.toBigInteger()).abs())
     }
 
     val floatCol = exampleFrame["float"] as DataColumn<Float?>
-    floatCol.type() shouldBe typeOf<Float?>()
+    floatCol.type() shouldBe typeOf<Float>().withNullability(nullable)
     floatCol.forEachIndexed { i, element ->
-        element shouldBe (2.0f.pow(iBatch(i).toFloat()))
+        assertValueOrNull(iBatch(i), element, 2.0f.pow(iBatch(i).toFloat()))
     }
 
     val doubleCol = exampleFrame["double"] as DataColumn<Double?>
-    doubleCol.type() shouldBe typeOf<Double?>()
+    doubleCol.type() shouldBe typeOf<Double>().withNullability(nullable)
     doubleCol.forEachIndexed { i, element ->
-        element shouldBe (2.0.pow(iBatch(i).toDouble()))
+        assertValueOrNull(iBatch(i), element, 2.0.pow(iBatch(i)))
     }
 
     val dateCol = exampleFrame["date32"] as DataColumn<LocalDate?>
-    dateCol.type() shouldBe typeOf<LocalDate?>()
+    dateCol.type() shouldBe typeOf<LocalDate>().withNullability(nullable)
     dateCol.forEachIndexed { i, element ->
-        element shouldBe LocalDate.ofEpochDay(iBatch(i).toLong() * 30)
+        assertValueOrNull(iBatch(i), element, LocalDate.ofEpochDay(iBatch(i).toLong() * 30))
     }
 
     val datetimeCol = exampleFrame["date64"] as DataColumn<LocalDateTime?>
-    datetimeCol.type() shouldBe typeOf<LocalDateTime?>()
+    datetimeCol.type() shouldBe typeOf<LocalDateTime>().withNullability(nullable)
     datetimeCol.forEachIndexed { i, element ->
-        element shouldBe LocalDateTime.ofEpochSecond(iBatch(i).toLong() * 60 * 60 * 24 * 30, 0, ZoneOffset.UTC)
+        assertValueOrNull(iBatch(i), element, LocalDateTime.ofEpochSecond(iBatch(i).toLong() * 60 * 60 * 24 * 30, 0, ZoneOffset.UTC))
     }
 
     val timeSecCol = exampleFrame["time32_seconds"] as DataColumn<LocalTime?>
-    timeSecCol.type() shouldBe typeOf<LocalTime?>()
+    timeSecCol.type() shouldBe typeOf<LocalTime>().withNullability(nullable)
     timeSecCol.forEachIndexed { i, element ->
-        element shouldBe LocalTime.ofSecondOfDay(iBatch(i).toLong())
+        assertValueOrNull(iBatch(i), element, LocalTime.ofSecondOfDay(iBatch(i).toLong()))
     }
 
     val timeMilliCol = exampleFrame["time32_milli"] as DataColumn<LocalTime?>
-    timeMilliCol.type() shouldBe typeOf<LocalTime?>()
+    timeMilliCol.type() shouldBe typeOf<LocalTime>().withNullability(nullable)
     timeMilliCol.forEachIndexed { i, element ->
-        element shouldBe LocalTime.ofNanoOfDay(iBatch(i).toLong() * 1000_000)
+        assertValueOrNull(iBatch(i), element, LocalTime.ofNanoOfDay(iBatch(i).toLong() * 1000_000))
     }
 
     val timeMicroCol = exampleFrame["time64_micro"] as DataColumn<LocalTime?>
-    timeMicroCol.type() shouldBe typeOf<LocalTime?>()
+    timeMicroCol.type() shouldBe typeOf<LocalTime>().withNullability(nullable)
     timeMicroCol.forEachIndexed { i, element ->
-        element shouldBe LocalTime.ofNanoOfDay(iBatch(i).toLong() * 1000)
+        assertValueOrNull(iBatch(i), element, LocalTime.ofNanoOfDay(iBatch(i).toLong() * 1000))
     }
 
     val timeNanoCol = exampleFrame["time64_nano"] as DataColumn<LocalTime?>
-    timeNanoCol.type() shouldBe typeOf<LocalTime?>()
+    timeNanoCol.type() shouldBe typeOf<LocalTime>().withNullability(nullable)
     timeNanoCol.forEachIndexed { i, element ->
-        element shouldBe LocalTime.ofNanoOfDay(iBatch(i).toLong())
+        assertValueOrNull(iBatch(i), element, LocalTime.ofNanoOfDay(iBatch(i).toLong()))
     }
 
 }