Skip to content

Commit 6007bbc

Browse files
Allex-NikAllex-Nik
committed
Add tests for the first and firstOrNull functions
1 parent 1a6d6e4 commit 6007bbc

File tree

1 file changed

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

1 file changed

+225
-0
lines changed

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

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
package org.jetbrains.kotlinx.dataframe.api
22

33
import io.kotest.assertions.throwables.shouldThrow
4+
import io.kotest.matchers.shouldBe
5+
import org.jetbrains.kotlinx.dataframe.nrow
6+
import org.jetbrains.kotlinx.dataframe.samples.api.age
47
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
8+
import org.jetbrains.kotlinx.dataframe.samples.api.isHappy
9+
import org.jetbrains.kotlinx.dataframe.samples.api.lastName
510
import org.jetbrains.kotlinx.dataframe.samples.api.name
11+
import org.jetbrains.kotlinx.dataframe.samples.api.weight
612
import org.junit.Test
713

814
class FirstTests : ColumnsSelectionDslTests() {
@@ -36,8 +42,227 @@ class FirstTests : ColumnsSelectionDslTests() {
3642
df.select { "name".firstCol { col -> col.any { it == "Alice" } } },
3743
df.select { Person::name.firstCol { col -> col.any { it == "Alice" } } },
3844
df.select { NonDataSchemaPerson::name.firstCol { col -> col.any { it == "Alice" } } },
45+
df.select { pathOf("name").firstCol() },
3946
df.select { pathOf("name").firstCol { col -> col.any { it == "Alice" } } },
4047
df.select { it["name"].asColumnGroup().firstCol { col -> col.any { it == "Alice" } } },
4148
).shouldAllBeEqual()
4249
}
50+
51+
@Test
52+
fun `first on DataColumn`() {
53+
df.name.lastName.first() shouldBe "Cooper"
54+
df.age.first { it in 18..<40 } shouldBe 20
55+
56+
shouldThrow<IndexOutOfBoundsException> {
57+
df.drop(df.nrow).isHappy.first()
58+
}
59+
}
60+
61+
@Test
62+
fun `firstOrNull on DataColumn`() {
63+
df.name.lastName.firstOrNull() shouldBe "Cooper"
64+
df.drop(2).weight.firstOrNull() shouldBe null
65+
df.drop(df.nrow).age.firstOrNull() shouldBe null
66+
67+
df.age.firstOrNull { it in 21..30 } shouldBe 30
68+
df.age.firstOrNull { it > 50 } shouldBe null
69+
}
70+
71+
@Test
72+
fun `first on DataFrame`() {
73+
df.first().name.lastName shouldBe "Cooper"
74+
df.first { !isHappy }.name.lastName shouldBe "Daniels"
75+
76+
shouldThrow<NoSuchElementException> {
77+
df.drop(df.nrow).first()
78+
}
79+
shouldThrow<NoSuchElementException> {
80+
df.first { age > 50 }
81+
}
82+
shouldThrow<NoSuchElementException> {
83+
df.drop(df.nrow).first { isHappy }
84+
}
85+
}
86+
87+
@Test
88+
fun `firstOrNull on DataFrame`() {
89+
df.firstOrNull()?.name?.lastName shouldBe "Cooper"
90+
df.drop(df.nrow).firstOrNull() shouldBe null
91+
92+
df.firstOrNull { !isHappy }?.name?.lastName shouldBe "Daniels"
93+
df.firstOrNull { age > 50 } shouldBe null
94+
df.drop(df.nrow).firstOrNull { isHappy } shouldBe null
95+
}
96+
97+
@Test
98+
fun `first on GroupBy`() {
99+
val grouped = df.groupBy { isHappy }
100+
val reducedGrouped = grouped.first()
101+
val firstHappy = reducedGrouped.values()[0]
102+
val firstUnhappy = reducedGrouped.values()[1]
103+
104+
firstHappy shouldBe dataFrameOf(
105+
"isHappy" to columnOf(true),
106+
"name" to columnOf(
107+
"firstName" to columnOf("Alice"),
108+
"lastName" to columnOf("Cooper"),
109+
),
110+
"age" to columnOf(15),
111+
"city" to columnOf("London"),
112+
"weight" to columnOf(54),
113+
)[0]
114+
115+
firstUnhappy shouldBe dataFrameOf(
116+
"isHappy" to columnOf(false),
117+
"name" to columnOf(
118+
"firstName" to columnOf("Charlie"),
119+
"lastName" to columnOf("Daniels")
120+
),
121+
"age" to columnOf(20),
122+
"city" to columnOf("Moscow"),
123+
"weight" to columnOf(null),
124+
)[0]
125+
}
126+
127+
@Test
128+
fun `first on GroupBy with predicate`() {
129+
val grouped = df.groupBy { isHappy }
130+
val reducedGrouped = grouped.first{ it["age"] as Int > 17 && it["city"] != "Moscow" }
131+
val firstHappy = reducedGrouped.values()[0]
132+
val firstUnhappy = reducedGrouped.values()[1]
133+
134+
firstHappy shouldBe dataFrameOf(
135+
"isHappy" to columnOf(true),
136+
"name" to columnOf(
137+
"firstName" to columnOf("Bob"),
138+
"lastName" to columnOf("Dylan"),
139+
),
140+
"age" to columnOf(45),
141+
"city" to columnOf("Dubai"),
142+
"weight" to columnOf(87),
143+
)[0]
144+
145+
firstUnhappy shouldBe dataFrameOf(
146+
"isHappy" to columnOf(false),
147+
"name" to columnOf(
148+
"firstName" to columnOf("Alice"),
149+
"lastName" to columnOf("Wolf")
150+
),
151+
"age" to columnOf(20),
152+
"city" to columnOf(null),
153+
"weight" to columnOf(55),
154+
)[0]
155+
}
156+
157+
@Test
158+
fun `first on Pivot`() {
159+
val pivot = df.pivot { isHappy }
160+
val reducedPivot = pivot.first()
161+
val firstHappy = reducedPivot.values()[0]
162+
val firstUnhappy = reducedPivot.values()[1]
163+
firstHappy shouldBe dataFrameOf(
164+
"name" to columnOf(
165+
"firstName" to columnOf("Alice"),
166+
"lastName" to columnOf("Cooper")
167+
),
168+
"age" to columnOf(15),
169+
"city" to columnOf("London"),
170+
"weight" to columnOf(54),
171+
)[0]
172+
173+
firstUnhappy shouldBe dataFrameOf(
174+
"name" to columnOf(
175+
"firstName" to columnOf("Charlie"),
176+
"lastName" to columnOf("Daniels")
177+
),
178+
"age" to columnOf(20),
179+
"city" to columnOf("Moscow"),
180+
"weight" to columnOf(null),
181+
)[0]
182+
}
183+
184+
@Test
185+
fun `first on Pivot with predicate`() {
186+
val pivot = df.pivot { isHappy }
187+
val reducedPivotAdults = pivot.first { age > 17 }
188+
val firstHappyAdult = reducedPivotAdults.values()[0]
189+
val firstUnhappyAdult = reducedPivotAdults.values()[1]
190+
191+
firstHappyAdult shouldBe dataFrameOf(
192+
"name" to columnOf(
193+
"firstName" to columnOf("Bob"),
194+
"lastName" to columnOf("Dylan")
195+
),
196+
"age" to columnOf(45),
197+
"city" to columnOf("Dubai"),
198+
"weight" to columnOf(87),
199+
)[0]
200+
201+
firstUnhappyAdult shouldBe dataFrameOf(
202+
"name" to columnOf(
203+
"firstName" to columnOf("Charlie"),
204+
"lastName" to columnOf("Daniels")
205+
),
206+
"age" to columnOf(20),
207+
"city" to columnOf("Moscow"),
208+
"weight" to columnOf(null),
209+
)[0]
210+
}
211+
212+
@Test
213+
fun `first on PivotGroupBy`() {
214+
val students = dataFrameOf(
215+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
216+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
217+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2)
218+
)
219+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
220+
val studentsPivotGroupedReduced = studentsPivotGrouped.first().values()
221+
val expectedDf = dataFrameOf(
222+
"name" to columnOf("Alice", "Bob"),
223+
"age" to columnOf(
224+
"15" to columnOf(1, 1),
225+
"20" to columnOf(1, 1),
226+
)
227+
)
228+
studentsPivotGroupedReduced shouldBe expectedDf
229+
}
230+
231+
@Test
232+
fun `first on PivotGroupBy with predicate`() {
233+
val students = dataFrameOf(
234+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
235+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
236+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2)
237+
)
238+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
239+
val studentsPivotGroupedReduced = studentsPivotGrouped.first { it["group"] == 2 }.values()
240+
val expected = dataFrameOf(
241+
"name" to columnOf("Alice", "Bob"),
242+
"age" to columnOf(
243+
"15" to columnOf(2, 2),
244+
"20" to columnOf(2, 2),
245+
)
246+
)
247+
studentsPivotGroupedReduced shouldBe expected
248+
}
249+
250+
@Test
251+
fun `first on PivotGroupBy with predicate without match`() {
252+
val students = dataFrameOf(
253+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
254+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
255+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2)
256+
)
257+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
258+
val studentsPivotGroupedReduced = studentsPivotGrouped.first { it["group"] == 3 }.values()
259+
val expected = dataFrameOf(
260+
"name" to columnOf("Alice", "Bob"),
261+
"age" to columnOf(
262+
"15" to columnOf(null, null),
263+
"20" to columnOf(null, null),
264+
)
265+
)
266+
studentsPivotGroupedReduced shouldBe expected
267+
}
43268
}

0 commit comments

Comments
 (0)