Skip to content

Commit 06b9e20

Browse files
Allex-NikAllex-Nik
committed
Add tests for the count function
1 parent a2c7f5d commit 06b9e20

File tree

1 file changed

+235
-0
lines changed
  • core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api

1 file changed

+235
-0
lines changed

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

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
package org.jetbrains.kotlinx.dataframe.api
2+
3+
import io.kotest.matchers.shouldBe
4+
import org.jetbrains.kotlinx.dataframe.nrow
5+
import org.junit.Test
6+
7+
class CountTests {
8+
9+
// Test data
10+
11+
val df = dataFrameOf(
12+
"name" to columnOf("Alice", "Bob", "Charlie"),
13+
"age" to columnOf(15, 20, 25),
14+
"group" to columnOf(1, 1, 2),
15+
)
16+
val age = df["age"].cast<Int>()
17+
val name = df["name"].cast<String>()
18+
val grouped = df.groupBy("group")
19+
val pivot = df.pivot("group")
20+
21+
val emptyDf = df.drop(df.nrow)
22+
23+
val dfWithNulls = df.append("Martin", null, null)
24+
val ageWithNulls = dfWithNulls["age"].cast<Int?>()
25+
val groupedWithNulls = dfWithNulls.groupBy("group")
26+
val pivotWithNulls = dfWithNulls.pivot("group")
27+
28+
// DataColumn
29+
30+
@Test
31+
fun `count on DataColumn`() {
32+
age.count() shouldBe 3
33+
age.count { it > 18 } shouldBe 2
34+
name.count { it.startsWith("A") } shouldBe 1
35+
}
36+
37+
@Test
38+
fun `count on empty DataColumn`() {
39+
emptyDf["name"].count() shouldBe 0
40+
emptyDf["name"].count { it == "Alice" } shouldBe 0
41+
}
42+
43+
@Test
44+
fun `count on DataColumn with nulls`() {
45+
ageWithNulls.count() shouldBe 4
46+
ageWithNulls.count { it == null } shouldBe 1
47+
}
48+
49+
// DataRow
50+
51+
@Test
52+
fun `count on DataRow`() {
53+
val row = df[0]
54+
row.count() shouldBe 3
55+
(row.count { it is Number }) shouldBe 2
56+
}
57+
58+
@Test
59+
fun `count on DataRow with nulls`() {
60+
val row = dfWithNulls[3]
61+
row.count() shouldBe 3
62+
row.count { it == null } shouldBe 2
63+
}
64+
65+
// DataFrame
66+
67+
@Test
68+
fun `count on DataFrame`() {
69+
df.count() shouldBe 3
70+
df.count { age > 18 } shouldBe 2
71+
df.count { it["name"] == "Alice" } shouldBe 1
72+
}
73+
74+
@Test
75+
fun `count on empty DataFrame`() {
76+
emptyDf.count() shouldBe 0
77+
}
78+
79+
@Test
80+
fun `count on DataFrame with nulls`() {
81+
dfWithNulls.count() shouldBe 4
82+
dfWithNulls.count { it["age"] != null } shouldBe 3
83+
}
84+
85+
// GroupBy
86+
87+
@Test
88+
fun `count on grouped DataFrame`() {
89+
val groupedCount = grouped.count()
90+
val expected = dataFrameOf(
91+
"group" to columnOf(1, 2),
92+
"count" to columnOf(2, 1),
93+
)
94+
groupedCount shouldBe expected
95+
}
96+
97+
@Test
98+
fun `count on grouped DataFrame with predicate`() {
99+
val groupedCount = grouped.count { "age"<Int>() > 18 }
100+
val expected = dataFrameOf(
101+
"group" to columnOf(1, 2),
102+
"count" to columnOf(1, 1),
103+
)
104+
groupedCount shouldBe expected
105+
}
106+
107+
@Test
108+
fun `count on empty grouped DataFrame`() {
109+
emptyDf.groupBy("group").count().count() shouldBe 0
110+
}
111+
112+
@Test
113+
fun `count on grouped DataFrame with nulls`() {
114+
val groupedWithNullsCount = groupedWithNulls.count()
115+
val expected = dataFrameOf(
116+
"group" to columnOf(1, 2, null),
117+
"count" to columnOf(2, 1, 1),
118+
)
119+
groupedWithNullsCount shouldBe expected
120+
}
121+
122+
@Test
123+
fun `count on grouped DataFrame with nulls and predicate`() {
124+
val groupedWithNullsCount = groupedWithNulls.count { it["age"] != null }
125+
val expected = dataFrameOf(
126+
"group" to columnOf(1, 2, null),
127+
"count" to columnOf(2, 1, 0),
128+
)
129+
groupedWithNullsCount shouldBe expected
130+
}
131+
132+
// Pivot
133+
134+
@Test
135+
fun `count on Pivot`() {
136+
val counted = pivot.count()
137+
val expected = dataFrameOf(
138+
"1" to columnOf(2),
139+
"2" to columnOf(1),
140+
)[0]
141+
counted shouldBe expected
142+
}
143+
144+
@Test
145+
fun `count on Pivot with predicate`() {
146+
val counted = pivot.count { "group"<Int>() != 1 }
147+
val expected = dataFrameOf(
148+
"1" to columnOf(0),
149+
"2" to columnOf(1),
150+
)[0]
151+
counted shouldBe expected
152+
}
153+
154+
@Test
155+
fun `count on Pivot with nulls`() {
156+
val counted = pivotWithNulls.count()
157+
val expected = dataFrameOf(
158+
"1" to columnOf(2),
159+
"2" to columnOf(1),
160+
"null" to columnOf(1),
161+
)[0]
162+
counted shouldBe expected
163+
}
164+
165+
@Test
166+
fun `count on Pivot with nulls and predicate`() {
167+
val counted = pivotWithNulls.count { it["age"] != null }
168+
val expected = dataFrameOf(
169+
"1" to columnOf(2),
170+
"2" to columnOf(1),
171+
"null" to columnOf(0),
172+
)[0]
173+
counted shouldBe expected
174+
}
175+
176+
// PivotGroupBy
177+
178+
@Test
179+
fun `count on PivotGroupBy`() {
180+
val pivotGrouped = pivot.groupBy("age")
181+
val counted = pivotGrouped.count()
182+
val expected = dataFrameOf(
183+
"age" to columnOf(15, 20, 25),
184+
"group" to columnOf(
185+
"1" to columnOf(1, 1, 0),
186+
"2" to columnOf(0, 0, 1),
187+
),
188+
)
189+
counted shouldBe expected
190+
}
191+
192+
@Test
193+
fun `count on PivotGroupBy with predicate`() {
194+
val pivotGrouped = pivot.groupBy("age")
195+
val counted = pivotGrouped.count { "name"<String>() == "Alice" }
196+
val expected = dataFrameOf(
197+
"age" to columnOf(15, 20, 25),
198+
"group" to columnOf(
199+
"1" to columnOf(1, 0, 0),
200+
"2" to columnOf(0, 0, 0),
201+
),
202+
)
203+
counted shouldBe expected
204+
}
205+
206+
@Test
207+
fun `count on PivotGroupBy with nulls`() {
208+
val pivotGrouped = pivotWithNulls.groupBy("age")
209+
val counted = pivotGrouped.count()
210+
val expected = dataFrameOf(
211+
"age" to columnOf(15, 20, 25, null),
212+
"group" to columnOf(
213+
"1" to columnOf(1, 1, 0, 0),
214+
"2" to columnOf(0, 0, 1, 0),
215+
"null" to columnOf(0, 0, 0, 1),
216+
),
217+
)
218+
counted shouldBe expected
219+
}
220+
221+
@Test
222+
fun `count PivotGroupBy with nulls and predicate`() {
223+
val pivotGrouped = pivotWithNulls.groupBy("age")
224+
val counted = pivotGrouped.count { it["age"] != null && "age"<Int>() > 15 }
225+
val expected = dataFrameOf(
226+
"age" to columnOf(15, 20, 25, null),
227+
"group" to columnOf(
228+
"1" to columnOf(0, 1, 0, 0),
229+
"2" to columnOf(0, 0, 1, 0),
230+
"null" to columnOf(0, 0, 0, 0),
231+
),
232+
)
233+
counted shouldBe expected
234+
}
235+
}

0 commit comments

Comments
 (0)