Add tests for the count function

Author: Allex-Nik

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/count.kt

@@ -0,0 +1,175 @@
+package org.jetbrains.kotlinx.dataframe.api
+
+import io.kotest.matchers.shouldBe
+import org.jetbrains.kotlinx.dataframe.nrow
+import org.junit.Test
+
+class CountTests {
+
+    // Test data
+
+    val df = dataFrameOf(
+        "name" to columnOf("Alice", "Bob", "Charlie"),
+        "age" to columnOf(15, 20, 25),
+        "group" to columnOf(1, 1, 2)
+    )
+    val age = df["age"].cast<Int>()
+    val name = df["name"].cast<String>()
+    val grouped = df.groupBy("group")
+    val pivot = df.pivot("group")
+
+    val emptyDf = df.drop(df.nrow)
+
+    val dfWithNulls = df.append("Martin", null, null)
+    val ageWithNulls = dfWithNulls["age"].cast<Int?>()
+    val groupedWithNulls = dfWithNulls.groupBy("group")
+    val pivotWithNulls = dfWithNulls.pivot("group")
+
+    // DataColumn
+
+    @Test
+    fun `count on column`() {
+        age.count() shouldBe 3
+        age.count { it > 18 } shouldBe 2
+        name.count { it.startsWith("A") } shouldBe 1
+    }
+
+    @Test
+    fun `count on empty column`() {
+        emptyDf["name"].count() shouldBe 0
+        emptyDf["name"].count { it == "Alice" } shouldBe 0
+    }
+
+    @Test
+    fun `count on column with nulls`() {
+        ageWithNulls.count() shouldBe 4
+        ageWithNulls.count { it == null } shouldBe 1
+    }
+
+    // DataRow
+
+    @Test
+    fun `count on row`() {
+        val row = df[0]
+        row.count() shouldBe 3
+        (row.count { it is Number }) shouldBe 2
+    }
+
+    @Test
+    fun `count on row with nulls`() {
+        val row = dfWithNulls[3]
+        row.count() shouldBe 3
+        row.count { it == null } shouldBe 2
+    }
+
+    // DataFrame
+
+    @Test
+    fun `count on dataframe`() {
+        df.count() shouldBe 3
+        df.count { age > 18 } shouldBe 2
+        df.count { it["name"] == "Alice" } shouldBe 1
+    }
+
+    @Test
+    fun `count on empty dataframe`() {
+        emptyDf.count() shouldBe 0
+    }
+
+    @Test
+    fun `count on dataframe with nulls`() {
+        dfWithNulls.count() shouldBe 4
+        dfWithNulls.count { it["age"] != null } shouldBe 3
+    }
+
+    // GroupBy
+
+    @Test
+    fun `count on grouped dataframe`() {
+        val groupedCount = grouped.count()
+        val expected = dataFrameOf(
+            "group" to columnOf(1, 2),
+            "count" to columnOf(2, 1)
+        )
+        groupedCount shouldBe expected
+    }
+
+    @Test
+    fun `count on grouped dataframe with predicate`() {
+        val groupedCount = grouped.count { "age"<Int>() > 18 }
+        val expected = dataFrameOf(
+            "group" to columnOf(1, 2),
+            "count" to columnOf(1, 1)
+        )
+        groupedCount shouldBe expected
+    }
+
+    @Test
+    fun `count on empty grouped dataframe`() {
+        emptyDf.groupBy("group").count().count() shouldBe 0
+    }
+
+    @Test
+    fun `count on grouped dataframe with nulls`() {
+        val groupedWithNullsCount = groupedWithNulls.count()
+        val expected = dataFrameOf(
+            "group" to columnOf(1, 2, null),
+            "count" to columnOf(2, 1, 1)
+        )
+        groupedWithNullsCount shouldBe expected
+    }
+
+    @Test
+    fun `count on grouped dataframe with nulls and predicate`() {
+        val groupedWithNullsCount = groupedWithNulls.count { it["age"] != null }
+        val expected = dataFrameOf(
+            "group" to columnOf(1, 2, null),
+            "count" to columnOf(2, 1, 0)
+        )
+        groupedWithNullsCount shouldBe expected
+    }
+
+    // Pivot
+
+    @Test
+    fun `count on pivot`() {
+        val counted = pivot.count()
+        val expected = dataFrameOf(
+            "1" to columnOf(2),
+            "2" to columnOf(1)
+        )[0]
+        counted shouldBe expected
+    }
+
+    @Test
+    fun `count on pivot with predicate`() {
+        val counted = pivot.count { "group"<Int>() != 1 }
+        val expected = dataFrameOf(
+            "1" to columnOf(0),
+            "2" to columnOf(1)
+        )[0]
+        counted shouldBe expected
+    }
+
+    @Test
+    fun `count on pivot with nulls`() {
+        val counted = pivotWithNulls.count()
+        val expected = dataFrameOf(
+            "1" to columnOf(2),
+            "2" to columnOf(1),
+            "null" to columnOf(1)
+        )[0]
+        counted shouldBe expected
+    }
+
+    @Test
+    fun `count on pivot with nulls and predicate`() {
+        val counted = pivotWithNulls.count { it["age"] != null }
+        val expected = dataFrameOf(
+            "1" to columnOf(2),
+            "2" to columnOf(1),
+            "null" to columnOf(0)
+        )[0]
+        counted shouldBe expected
+    }
+}