expanded cumSum with Short, Byte, and BigInteger

Author: Jolanrensen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/cumSum.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlinx.dataframe.math.cumSum
 import org.jetbrains.kotlinx.dataframe.math.defaultCumSumSkipNA
 import org.jetbrains.kotlinx.dataframe.typeClass
 import java.math.BigDecimal
+import java.math.BigInteger
 import kotlin.reflect.KProperty
 import kotlin.reflect.typeOf
 
@@ -22,15 +23,30 @@ public fun <T : Number?> DataColumn<T>.cumSum(skipNA: Boolean = defaultCumSumSki
         typeOf<Float?>() -> cast<Float?>().cumSum(skipNA).cast()
         typeOf<Int>() -> cast<Int>().cumSum().cast()
         typeOf<Int?>() -> cast<Int?>().cumSum(skipNA).cast()
+        typeOf<Byte>() -> cast<Byte>().cumSum().cast()
+        typeOf<Byte?>() -> cast<Byte?>().cumSum(skipNA).cast()
+        typeOf<Short>() -> cast<Short>().cumSum().cast()
+        typeOf<Short?>() -> cast<Short?>().cumSum(skipNA).cast()
         typeOf<Long>() -> cast<Long>().cumSum().cast()
         typeOf<Long?>() -> cast<Long?>().cumSum(skipNA).cast()
+        typeOf<BigInteger>() -> cast<BigInteger>().cumSum().cast()
+        typeOf<BigInteger?>() -> cast<BigInteger?>().cumSum(skipNA).cast()
         typeOf<BigDecimal>() -> cast<BigDecimal>().cumSum().cast()
         typeOf<BigDecimal?>() -> cast<BigDecimal?>().cumSum(skipNA).cast()
         typeOf<Number?>(), typeOf<Number>() -> convertToDouble().cumSum(skipNA).cast()
         else -> error("Cumsum for type ${type()} is not supported")
     }
 
-private val supportedClasses = setOf(Double::class, Float::class, Int::class, Long::class, BigDecimal::class)
+private val supportedClasses = setOf(
+    Double::class,
+    Float::class,
+    Int::class,
+    Byte::class,
+    Short::class,
+    Long::class,
+    BigInteger::class,
+    BigDecimal::class,
+)
 
 // endregion
 

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/math/cumsum.kt

@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.api.isNA
 import org.jetbrains.kotlinx.dataframe.api.map
 import java.math.BigDecimal
+import java.math.BigInteger
 
 internal val defaultCumSumSkipNA: Boolean = true
 
@@ -88,6 +89,66 @@ internal fun DataColumn<Int?>.cumSum(skipNA: Boolean = defaultCumSumSkipNA): Dat
     }
 }
 
+@JvmName("byteCumsum")
+internal fun DataColumn<Byte>.cumSum(): DataColumn<Byte> {
+    var sum = 0.toByte()
+    return map {
+        sum = (sum + it).toByte()
+        sum
+    }
+}
+
+@JvmName("cumsumByteNullable")
+internal fun DataColumn<Byte?>.cumSum(skipNA: Boolean = defaultCumSumSkipNA): DataColumn<Byte?> {
+    var sum = 0.toByte()
+    var fillNull = false
+    return map {
+        when {
+            it == null -> {
+                if (!skipNA) fillNull = true
+                null
+            }
+
+            fillNull -> null
+
+            else -> {
+                sum = (sum + it).toByte()
+                sum
+            }
+        }
+    }
+}
+
+@JvmName("shortCumsum")
+internal fun DataColumn<Short>.cumSum(): DataColumn<Short> {
+    var sum = 0.toShort()
+    return map {
+        sum = (sum + it).toShort()
+        sum
+    }
+}
+
+@JvmName("cumsumShortNullable")
+internal fun DataColumn<Short?>.cumSum(skipNA: Boolean = defaultCumSumSkipNA): DataColumn<Short?> {
+    var sum = 0.toShort()
+    var fillNull = false
+    return map {
+        when {
+            it == null -> {
+                if (!skipNA) fillNull = true
+                null
+            }
+
+            fillNull -> null
+
+            else -> {
+                sum = (sum + it).toShort()
+                sum
+            }
+        }
+    }
+}
+
 @JvmName("longCumsum")
 internal fun DataColumn<Long>.cumSum(): DataColumn<Long> {
     var sum = 0L
@@ -118,6 +179,36 @@ internal fun DataColumn<Long?>.cumSum(skipNA: Boolean = defaultCumSumSkipNA): Da
     }
 }
 
+@JvmName("bigIntegerCumsum")
+internal fun DataColumn<BigInteger>.cumSum(): DataColumn<BigInteger> {
+    var sum = BigInteger.ZERO
+    return map {
+        sum += it
+        sum
+    }
+}
+
+@JvmName("cumsumBigIntegerNullable")
+internal fun DataColumn<BigInteger?>.cumSum(skipNA: Boolean = defaultCumSumSkipNA): DataColumn<BigInteger?> {
+    var sum = BigInteger.ZERO
+    var fillNull = false
+    return map {
+        when {
+            it == null -> {
+                if (!skipNA) fillNull = true
+                null
+            }
+
+            fillNull -> null
+
+            else -> {
+                sum += it
+                sum
+            }
+        }
+    }
+}
+
 @JvmName("bigDecimalCumsum")
 internal fun DataColumn<BigDecimal>.cumSum(): DataColumn<BigDecimal> {
     var sum = BigDecimal.ZERO

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/cumsum.kt

@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.api.concat
 import org.jetbrains.kotlinx.dataframe.api.cumSum
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.groupBy
+import org.jetbrains.kotlinx.dataframe.api.map
 import org.junit.Test
 
 @Suppress("ktlint:standard:argument-list-wrapping")
@@ -22,6 +23,30 @@ class CumsumTests {
         col.cumSum(skipNA = false).toList() shouldBe expectedNoSkip
     }
 
+    @Test
+    fun `short column`() {
+        col.map { it?.toShort() }.cumSum().toList() shouldBe expected.map { it?.toShort() }
+        col.map { it?.toShort() }.cumSum(skipNA = false).toList() shouldBe expectedNoSkip.map { it?.toShort() }
+    }
+
+    @Test
+    fun `byte column`() {
+        col.map { it?.toByte() }.cumSum().toList() shouldBe expected.map { it?.toByte() }
+        col.map { it?.toByte() }.cumSum(skipNA = false).toList() shouldBe expectedNoSkip.map { it?.toByte() }
+    }
+
+    @Test
+    fun `big int column`() {
+        col.map { it?.toBigInteger() }.cumSum().toList() shouldBe expected.map { it?.toBigInteger() }
+        col.map { it?.toBigInteger() }.cumSum(skipNA = false).toList() shouldBe expectedNoSkip.map { it?.toBigInteger() }
+    }
+
+    @Test
+    fun `big decimal column`() {
+        col.map { it?.toBigDecimal() }.cumSum().toList() shouldBe expected.map { it?.toBigDecimal() }
+        col.map { it?.toBigDecimal() }.cumSum(skipNA = false).toList() shouldBe expectedNoSkip.map { it?.toBigDecimal() }
+    }
+
     @Test
     fun frame() {
         val str by columnOf("a", "b", "c", "d", "e")