From 06b9e20c1b363e22e6debbb57d1534514e16995d Mon Sep 17 00:00:00 2001 From: Aleksandr Nikolaev Date: Wed, 29 Oct 2025 15:22:26 +0100 Subject: [PATCH] Add tests for the `count` function --- .../jetbrains/kotlinx/dataframe/api/count.kt | 235 ++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/count.kt diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/count.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/count.kt new file mode 100644 index 0000000000..f11b5ac42b --- /dev/null +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/count.kt @@ -0,0 +1,235 @@ +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() + val name = df["name"].cast() + 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() + val groupedWithNulls = dfWithNulls.groupBy("group") + val pivotWithNulls = dfWithNulls.pivot("group") + + // DataColumn + + @Test + fun `count on DataColumn`() { + age.count() shouldBe 3 + age.count { it > 18 } shouldBe 2 + name.count { it.startsWith("A") } shouldBe 1 + } + + @Test + fun `count on empty DataColumn`() { + emptyDf["name"].count() shouldBe 0 + emptyDf["name"].count { it == "Alice" } shouldBe 0 + } + + @Test + fun `count on DataColumn with nulls`() { + ageWithNulls.count() shouldBe 4 + ageWithNulls.count { it == null } shouldBe 1 + } + + // DataRow + + @Test + fun `count on DataRow`() { + val row = df[0] + row.count() shouldBe 3 + (row.count { it is Number }) shouldBe 2 + } + + @Test + fun `count on DataRow 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"() > 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"() != 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 + } + + // PivotGroupBy + + @Test + fun `count on PivotGroupBy`() { + val pivotGrouped = pivot.groupBy("age") + val counted = pivotGrouped.count() + val expected = dataFrameOf( + "age" to columnOf(15, 20, 25), + "group" to columnOf( + "1" to columnOf(1, 1, 0), + "2" to columnOf(0, 0, 1), + ), + ) + counted shouldBe expected + } + + @Test + fun `count on PivotGroupBy with predicate`() { + val pivotGrouped = pivot.groupBy("age") + val counted = pivotGrouped.count { "name"() == "Alice" } + val expected = dataFrameOf( + "age" to columnOf(15, 20, 25), + "group" to columnOf( + "1" to columnOf(1, 0, 0), + "2" to columnOf(0, 0, 0), + ), + ) + counted shouldBe expected + } + + @Test + fun `count on PivotGroupBy with nulls`() { + val pivotGrouped = pivotWithNulls.groupBy("age") + val counted = pivotGrouped.count() + val expected = dataFrameOf( + "age" to columnOf(15, 20, 25, null), + "group" to columnOf( + "1" to columnOf(1, 1, 0, 0), + "2" to columnOf(0, 0, 1, 0), + "null" to columnOf(0, 0, 0, 1), + ), + ) + counted shouldBe expected + } + + @Test + fun `count PivotGroupBy with nulls and predicate`() { + val pivotGrouped = pivotWithNulls.groupBy("age") + val counted = pivotGrouped.count { it["age"] != null && "age"() > 15 } + val expected = dataFrameOf( + "age" to columnOf(15, 20, 25, null), + "group" to columnOf( + "1" to columnOf(0, 1, 0, 0), + "2" to columnOf(0, 0, 1, 0), + "null" to columnOf(0, 0, 0, 0), + ), + ) + counted shouldBe expected + } +}