Skip to content

Commit 5468953

Browse files
Allex-NikAllex-Nik
committed
Add tests for the last and lastOrNull functions
1 parent 35ddc71 commit 5468953

File tree

1 file changed

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

1 file changed

+307
-0
lines changed

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

Lines changed: 307 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,37 @@
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
7+
import org.jetbrains.kotlinx.dataframe.samples.api.city
48
import org.jetbrains.kotlinx.dataframe.samples.api.firstName
59
import org.jetbrains.kotlinx.dataframe.samples.api.isHappy
10+
import org.jetbrains.kotlinx.dataframe.samples.api.lastName
611
import org.jetbrains.kotlinx.dataframe.samples.api.name
12+
import org.jetbrains.kotlinx.dataframe.samples.api.weight
713
import org.junit.Test
814

15+
/**
16+
* This class tests the behavior of the `last` (`lastCol`) and `lastOrNull` functions, including:
17+
* - **ColumnsSelectionDsl**: selecting the last column or last column matching a condition, with invocations
18+
* on illegal types, and in case when no column matches the condition.
19+
* - **DataColumn**: getting the last value or the last value matching a predicate,
20+
* verifying behavior on empty columns, columns with `null` values, and columns without values matching the predicate.
21+
* - **DataFrame**: getting the last row or last matching row,
22+
* verifying behavior on empty DataFrames, DataFrames with `null` values, and
23+
* DataFrames without rows matching the predicate.
24+
* - **GroupBy**: reducing each group to its last row or the last row matching a predicate,
25+
* with handling groups that contain no matching rows.
26+
* - **Pivot**: reducing each group in the pivot to its last row or the last row matching a predicate,
27+
* with handling groups that contain no matching rows.
28+
* - **PivotGroupBy**: reducing each combined [pivot] + [groupBy] group to its last row
29+
* or the last row matching a predicate, with handling [pivot] + [groupBy] combinations that contain no matching rows.
30+
*/
931
class LastTests : ColumnsSelectionDslTests() {
32+
33+
// region ColumnsSelectionDsl
34+
1035
@Test
1136
fun `ColumnsSelectionDsl last`() {
1237
shouldThrow<IllegalArgumentException> {
@@ -36,8 +61,290 @@ class LastTests : ColumnsSelectionDslTests() {
3661
df.select { "name".lastCol { col -> col.any { it == "Alice" } } },
3762
df.select { Person::name.lastCol { col -> col.any { it == "Alice" } } },
3863
df.select { NonDataSchemaPerson::name.lastCol { col -> col.any { it == "Alice" } } },
64+
df.remove { name.lastName }.select { pathOf("name").lastCol() },
3965
df.select { pathOf("name").lastCol { col -> col.any { it == "Alice" } } },
4066
df.select { it["name"].asColumnGroup().lastCol { col -> col.any { it == "Alice" } } },
4167
).shouldAllBeEqual()
4268
}
69+
70+
// endregion
71+
72+
// region DataColumn
73+
74+
@Test
75+
fun `last on DataColumn`() {
76+
df.name.lastName.last() shouldBe "Byrd"
77+
df.age.last { it > 30 } shouldBe 40
78+
shouldThrow<IndexOutOfBoundsException> {
79+
df.drop(df.nrow).isHappy.last()
80+
}
81+
}
82+
83+
@Test
84+
fun `lastOrNull on DataColumn`() {
85+
df.name.lastName.lastOrNull() shouldBe "Byrd"
86+
df.take(4).weight.lastOrNull() shouldBe null
87+
df.drop(df.nrow).age.lastOrNull() shouldBe null
88+
df.age.lastOrNull { it > 30 } shouldBe 40
89+
df.age.lastOrNull { it > 50 } shouldBe null
90+
}
91+
92+
// endregion
93+
94+
// region DataFrame
95+
96+
@Test
97+
fun `last on DataFrame`() {
98+
df.last().name.lastName shouldBe "Byrd"
99+
df.last { !isHappy }.name.lastName shouldBe "Wolf"
100+
shouldThrow<NoSuchElementException> {
101+
df.drop(df.nrow).last()
102+
}
103+
shouldThrow<NoSuchElementException> {
104+
df.last { age > 50 }
105+
}
106+
shouldThrow<NoSuchElementException> {
107+
df.drop(df.nrow).last { isHappy }
108+
}
109+
}
110+
111+
@Test
112+
fun `lastOrNull on DataFrame`() {
113+
df.lastOrNull()?.name?.lastName shouldBe "Byrd"
114+
df.take(6).lastOrNull()?.city shouldBe null
115+
df.drop(df.nrow).lastOrNull() shouldBe null
116+
df.lastOrNull { !isHappy }?.name?.lastName shouldBe "Wolf"
117+
df.lastOrNull { age > 50 } shouldBe null
118+
df.drop(df.nrow).lastOrNull { isHappy } shouldBe null
119+
}
120+
121+
// endregion
122+
123+
// region GroupBy
124+
125+
@Test
126+
fun `last on GroupBy`() {
127+
val grouped = df.groupBy { isHappy }
128+
val reducedGrouped = grouped.last()
129+
val lastHappy = reducedGrouped.values()[0]
130+
val lastUnhappy = reducedGrouped.values()[1]
131+
lastHappy shouldBe dataFrameOf(
132+
"isHappy" to columnOf(true),
133+
"name" to columnOf(
134+
"firstName" to columnOf("Charlie"),
135+
"lastName" to columnOf("Byrd"),
136+
),
137+
"age" to columnOf(30),
138+
"city" to columnOf("Moscow"),
139+
"weight" to columnOf(90),
140+
)[0]
141+
lastUnhappy shouldBe dataFrameOf(
142+
"isHappy" to columnOf(false),
143+
"name" to columnOf(
144+
"firstName" to columnOf("Alice"),
145+
"lastName" to columnOf("Wolf"),
146+
),
147+
"age" to columnOf(20),
148+
"city" to columnOf(null),
149+
"weight" to columnOf(55),
150+
)[0]
151+
}
152+
153+
@Test
154+
fun `last on GroupBy with predicate`() {
155+
val grouped = df.groupBy { isHappy }
156+
val reducedGrouped = grouped.last { "age"<Int>() < 21 && it["city"] != "Moscow" }
157+
val lastHappy = reducedGrouped.values()[0]
158+
val lastUnhappy = reducedGrouped.values()[1]
159+
lastHappy shouldBe dataFrameOf(
160+
"isHappy" to columnOf(true),
161+
"name" to columnOf(
162+
"firstName" to columnOf("Alice"),
163+
"lastName" to columnOf("Cooper"),
164+
),
165+
"age" to columnOf(15),
166+
"city" to columnOf("London"),
167+
"weight" to columnOf(54),
168+
)[0]
169+
lastUnhappy shouldBe dataFrameOf(
170+
"isHappy" to columnOf(false),
171+
"name" to columnOf(
172+
"firstName" to columnOf("Alice"),
173+
"lastName" to columnOf("Wolf"),
174+
),
175+
"age" to columnOf(20),
176+
"city" to columnOf(null),
177+
"weight" to columnOf(55),
178+
)[0]
179+
}
180+
181+
@Test
182+
fun `last on GroupBy with predicate without match`() {
183+
val grouped = df.groupBy { isHappy }
184+
val reducedGrouped = grouped.last { it["city"] == "London" }
185+
val lastHappy = reducedGrouped.values()[0]
186+
val lastUnhappy = reducedGrouped.values()[1]
187+
lastHappy shouldBe dataFrameOf(
188+
"isHappy" to columnOf(true),
189+
"name" to columnOf(
190+
"firstName" to columnOf("Alice"),
191+
"lastName" to columnOf("Cooper"),
192+
),
193+
"age" to columnOf(15),
194+
"city" to columnOf("London"),
195+
"weight" to columnOf(54),
196+
)[0]
197+
lastUnhappy shouldBe dataFrameOf(
198+
"isHappy" to columnOf(false),
199+
"name" to columnOf(
200+
"firstName" to columnOf(null),
201+
"lastName" to columnOf(null),
202+
),
203+
"age" to columnOf(null),
204+
"city" to columnOf(null),
205+
"weight" to columnOf(null),
206+
)[0]
207+
}
208+
209+
// endregion
210+
211+
// region Pivot
212+
213+
@Test
214+
fun `last on Pivot`() {
215+
val pivot = df.pivot { isHappy }
216+
val reducedPivot = pivot.last()
217+
val lastHappy = reducedPivot.values()[0]
218+
val lastUnhappy = reducedPivot.values()[1]
219+
lastHappy shouldBe dataFrameOf(
220+
"name" to columnOf(
221+
"firstName" to columnOf("Charlie"),
222+
"lastName" to columnOf("Byrd"),
223+
),
224+
"age" to columnOf(30),
225+
"city" to columnOf("Moscow"),
226+
"weight" to columnOf(90),
227+
)[0]
228+
lastUnhappy shouldBe dataFrameOf(
229+
"name" to columnOf(
230+
"firstName" to columnOf("Alice"),
231+
"lastName" to columnOf("Wolf"),
232+
),
233+
"age" to columnOf(20),
234+
"city" to columnOf(null),
235+
"weight" to columnOf(55),
236+
)[0]
237+
}
238+
239+
@Test
240+
fun `last on Pivot with predicate`() {
241+
val pivot = df.pivot { isHappy }
242+
val reducedPivot = pivot.last { "age"<Int>() < 21 && it["city"] != "Moscow" }
243+
val lastHappy = reducedPivot.values()[0]
244+
val lastUnhappy = reducedPivot.values()[1]
245+
lastHappy shouldBe dataFrameOf(
246+
"name" to columnOf(
247+
"firstName" to columnOf("Alice"),
248+
"lastName" to columnOf("Cooper"),
249+
),
250+
"age" to columnOf(15),
251+
"city" to columnOf("London"),
252+
"weight" to columnOf(54),
253+
)[0]
254+
lastUnhappy shouldBe dataFrameOf(
255+
"name" to columnOf(
256+
"firstName" to columnOf("Alice"),
257+
"lastName" to columnOf("Wolf"),
258+
),
259+
"age" to columnOf(20),
260+
"city" to columnOf(null),
261+
"weight" to columnOf(55),
262+
)[0]
263+
}
264+
265+
@Test
266+
fun `last on Pivot with predicate without match`() {
267+
val pivot = df.pivot { isHappy }
268+
val reducedPivot = pivot.last { it["city"] == "London" }
269+
val lastHappy = reducedPivot.values()[0]
270+
val lastUnhappy = reducedPivot.values()[1]
271+
lastHappy shouldBe dataFrameOf(
272+
"name" to columnOf(
273+
"firstName" to columnOf("Alice"),
274+
"lastName" to columnOf("Cooper"),
275+
),
276+
"age" to columnOf(15),
277+
"city" to columnOf("London"),
278+
"weight" to columnOf(54),
279+
)[0]
280+
lastUnhappy shouldBe dataFrameOf(
281+
"name" to columnOf(null),
282+
"age" to columnOf(null),
283+
"city" to columnOf(null),
284+
"weight" to columnOf(null),
285+
)[0]
286+
}
287+
288+
// endregion
289+
290+
// region PivotGroupBy
291+
292+
@Test
293+
fun `last on PivotGroupBy`() {
294+
val students = dataFrameOf(
295+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
296+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
297+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2),
298+
)
299+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
300+
val studentsPivotGroupedReduced = studentsPivotGrouped.last().values()
301+
val expectedDf = dataFrameOf(
302+
"name" to columnOf("Alice", "Bob"),
303+
"age" to columnOf(
304+
"15" to columnOf(2, 2),
305+
"20" to columnOf(2, 2),
306+
),
307+
)
308+
studentsPivotGroupedReduced shouldBe expectedDf
309+
}
310+
311+
@Test
312+
fun `last on PivotGroupBy with predicate`() {
313+
val students = dataFrameOf(
314+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
315+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
316+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2),
317+
)
318+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
319+
val studentsPivotGroupedReduced = studentsPivotGrouped.last { it["group"] == 1 }.values()
320+
val expected = dataFrameOf(
321+
"name" to columnOf("Alice", "Bob"),
322+
"age" to columnOf(
323+
"15" to columnOf(1, 1),
324+
"20" to columnOf(1, 1),
325+
),
326+
)
327+
studentsPivotGroupedReduced shouldBe expected
328+
}
329+
330+
@Test
331+
fun `last on PivotGroupBy with predicate without match`() {
332+
val students = dataFrameOf(
333+
"name" to columnOf("Alice", "Alice", "Alice", "Alice", "Bob", "Bob", "Bob", "Bob"),
334+
"age" to columnOf(15, 15, 20, 20, 15, 15, 20, 20),
335+
"group" to columnOf(1, 2, 1, 2, 1, 2, 1, 2),
336+
)
337+
val studentsPivotGrouped = students.pivot("age").groupBy("name")
338+
val studentsPivotGroupedReduced = studentsPivotGrouped.last { it["group"] == 3 }.values()
339+
val expected = dataFrameOf(
340+
"name" to columnOf("Alice", "Bob"),
341+
"age" to columnOf(
342+
"15" to columnOf(null, null),
343+
"20" to columnOf(null, null),
344+
),
345+
)
346+
studentsPivotGroupedReduced shouldBe expected
347+
}
348+
349+
// endregion
43350
}

0 commit comments

Comments
 (0)