Skip to content

Commit 8a46ded

Browse files
committed
added other test for sort grouped df to make sure it's fixed
1 parent 49592ba commit 8a46ded

File tree

1 file changed

+72
-0
lines changed

1 file changed

+72
-0
lines changed
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package org.jetbrains.kotlinx.dataframe.api
2+
3+
import io.kotest.matchers.shouldBe
4+
import org.jetbrains.kotlinx.dataframe.DataFrame
5+
import org.jetbrains.kotlinx.dataframe.alsoDebug
6+
import org.jetbrains.kotlinx.dataframe.io.read
7+
import org.junit.Test
8+
9+
class SortGroupedDataframeTests {
10+
11+
@Test
12+
fun `Sorted grouped iris dataset`() {
13+
val irisData = DataFrame.read("src/test/resources/irisDataset.csv")
14+
irisData.alsoDebug()
15+
16+
irisData.groupBy("variety").let {
17+
it.sortBy("petal.length").toString() shouldBe
18+
it.sortBy { it["petal.length"] }.toString()
19+
}
20+
}
21+
22+
enum class State {
23+
Idle, Productive, Maintenance
24+
}
25+
26+
@Test
27+
fun test4() {
28+
class Event(val toolId: String, val state: State, val timestamp: Long)
29+
30+
val tool1 = "tool_1"
31+
val tool2 = "tool_2"
32+
val tool3 = "tool_3"
33+
34+
val events = listOf(
35+
Event(tool1, State.Idle, 0),
36+
Event(tool1, State.Productive, 5),
37+
Event(tool2, State.Idle, 0),
38+
Event(tool2, State.Maintenance, 10),
39+
Event(tool2, State.Idle, 20),
40+
Event(tool3, State.Idle, 0),
41+
Event(tool3, State.Productive, 25),
42+
).toDataFrame()
43+
44+
val lastTimestamp = events.maxOf { getValue<Long>("timestamp") }
45+
val groupBy = events
46+
.groupBy("toolId")
47+
.sortBy("timestamp")
48+
.add("stateDuration") {
49+
(next()?.getValue("timestamp") ?: lastTimestamp) - getValue<Long>("timestamp")
50+
}
51+
52+
groupBy.toDataFrame().alsoDebug()
53+
groupBy.schema().print()
54+
groupBy.keys.print()
55+
groupBy.keys[0].print()
56+
57+
val df1 = groupBy.updateGroups {
58+
val missingValues = State.values().asList().toDataFrame {
59+
"state" from { it }
60+
}
61+
62+
val df = it
63+
.fullJoin(missingValues, "state")
64+
.fillNulls("stateDuration")
65+
.with { 100L }
66+
67+
df.groupBy("state").sumFor("stateDuration")
68+
}
69+
70+
df1.toDataFrame().alsoDebug().isNotEmpty() shouldBe true
71+
}
72+
}

0 commit comments

Comments
 (0)