group by dataset handler test

Kotlin · Apr 28, 2023 · 8c01be2 · 8c01be2
1 parent 42a4e99
commit 8c01be2
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 11 deletions.
diff --git a/kandy-api/src/main/kotlin/org/jetbrains/kotlinx/kandy/dsl/internal/DatasetHandler.kt b/kandy-api/src/main/kotlin/org/jetbrains/kotlinx/kandy/dsl/internal/DatasetHandler.kt
@@ -56,7 +56,9 @@ public class DatasetHandler(
 
     public fun takeColumn(name: String): String {
         return referredColumns[name] ?: run {
-            val columnId = internalAddColumn(initialNamedData.dataFrame.getColumnOrNull(name) ?: error("invalid column id"))
+            val columnId = internalAddColumn(
+                initialNamedData.dataFrame.getColumnOrNull(name) ?: error("invalid column id")
+            )
             referredColumns[name] = columnId
             name
         }

diff --git a/kandy-api/src/main/kotlin/org/jetbrains/kotlinx/kandy/dsl/internal/GroupByPlotContext.kt b/kandy-api/src/main/kotlin/org/jetbrains/kotlinx/kandy/dsl/internal/GroupByPlotContext.kt
@@ -28,11 +28,4 @@ public class GroupByPlotContext<T>(
         DatasetHandler(GroupedData(groupBy))
     )
     override val features: MutableMap<FeatureName, PlotFeature> = mutableMapOf()
-
-/*
-    public fun <C> columns(selector: ColumnsSelector<G, C>): List<DataColumn<C>> = groupBy.groups.concat().get(selector)
-    public fun <C> columns(vararg columns: String): List<AnyCol> = groupBy.groups.concat().getColumns(*columns)
-
-
- */
 }
diff --git a/...y-api/src/test/kotlin/org/jetbrains/kotlinx/kandy/dsl/unit/internal/DatasetHandlerTest.kt b/...y-api/src/test/kotlin/org/jetbrains/kotlinx/kandy/dsl/unit/internal/DatasetHandlerTest.kt
@@ -1,14 +1,15 @@
 package org.jetbrains.kotlinx.kandy.dsl.unit.internal
 
 import org.jetbrains.kotlinx.dataframe.DataFrame
-import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
-import org.jetbrains.kotlinx.dataframe.api.map
-import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.api.*
 import org.jetbrains.kotlinx.kandy.dsl.internal.DatasetHandler
+import org.jetbrains.kotlinx.kandy.ir.data.GroupedData
 import org.jetbrains.kotlinx.kandy.ir.data.NamedData
+import org.jetbrains.kotlinx.kandy.ir.data.TableData
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertIs
 
 internal class DatasetHandlerTest {
 
@@ -21,6 +22,7 @@ internal class DatasetHandlerTest {
 
         val handler = DatasetHandler(NamedData(dataFrame))
         assertEquals(DataFrame.Empty, handler.buffer)
+        assertEquals(NamedData(DataFrame.Empty), handler.data())
 
         // test simple take
         val colTypeIDAfterTake = handler.takeColumn(columnType.name())
@@ -70,4 +72,75 @@ internal class DatasetHandlerTest {
 
     }
 
+
+    private fun assertEqualsContent(expected: GroupBy<*, *>, actual: TableData) {
+        assertIs<GroupedData>(actual)
+        assertEquals(expected.keys.columnNames(), actual.groupBy.keys.columnNames())
+        assertEquals(expected.groups.concat(), actual.groupBy.groups.concat())
+    }
+
+    @Test
+    fun testGroupedData() {
+        val columnNumbers = listOf(12, 34, 56, 78, 90).toColumn("numbers")
+        val columnType = listOf("a", "b", "a", "c", "b").toColumn("TYPE")
+        val columnCond = listOf(true, false, false, true, true).toColumn("cond")
+        val groupKey = "TYPE"
+        val dataFrame = dataFrameOf(columnNumbers, columnType, columnCond).groupBy(groupKey)
+        val groupedData = GroupedData(dataFrame)
+
+        val internalColumnNumbers = groupedData.origin.dataFrame["numbers"]
+        val internalColumnType = groupedData.origin.dataFrame["TYPE"]
+        val internalColumnCond = groupedData.origin.dataFrame["cond"]
+
+        val handler = DatasetHandler(groupedData)
+        assertEquals(dataFrameOf(internalColumnType), handler.buffer, groupKey)
+
+        // test simple take
+        val colTypeIDAfterTake = handler.takeColumn(columnType.name())
+        assertEquals(dataFrameOf(internalColumnType), handler.buffer, groupKey)
+        assertEquals(columnType.name(), colTypeIDAfterTake)
+
+        // test add after take
+        val colTypeIDAfterAdd = handler.addColumn(internalColumnType)
+        assertEquals(dataFrameOf(internalColumnType), handler.buffer, groupKey)
+        assertEquals(columnType.name(), colTypeIDAfterAdd)
+
+        assertEqualsContent(dataFrameOf(internalColumnType).groupBy(groupKey), handler.data())
+
+        // test simple add
+        val colNumbersIDAfterAdd = handler.addColumn(internalColumnNumbers)
+        assertEquals(dataFrameOf(internalColumnType, internalColumnNumbers), handler.buffer)
+        assertEquals(columnNumbers.name(), colNumbersIDAfterAdd)
+
+        // test take after add
+        val colNumbersIDAfterTake = handler.takeColumn(columnNumbers.name())
+        assertEquals(dataFrameOf(internalColumnType, internalColumnNumbers), handler.buffer)
+        assertEquals(columnNumbers.name(), colNumbersIDAfterTake)
+
+        // test add after add
+        val colNumbersIDAfterRepeatAdd = handler.addColumn(internalColumnNumbers)
+        assertEquals(dataFrameOf(internalColumnType, internalColumnNumbers), handler.buffer)
+        assertEquals(columnNumbers.name(), colNumbersIDAfterRepeatAdd)
+
+        assertEqualsContent(dataFrameOf(internalColumnType, internalColumnNumbers).groupBy(groupKey), handler.data())
+
+        // test list add
+        val colCondIDAfterAdd = handler.addColumn(columnCond.values().toList(), columnCond.name())
+        assertEquals(dataFrameOf(internalColumnType, internalColumnNumbers, internalColumnCond), handler.buffer)
+        assertEquals(columnCond.name(), colCondIDAfterAdd)
+
+        val columnNumbers2 = columnNumbers.map { it + 1 }
+
+        // test add same with same name and different values
+        val columnNumbers2IDAfterAdd = handler.addColumn(columnNumbers2)
+        assertEquals(dataFrameOf(internalColumnType, internalColumnNumbers, internalColumnCond, columnNumbers2.rename("numbers*")), handler.buffer)
+        assertEquals("numbers*", columnNumbers2IDAfterAdd)
+
+        // test take
+        assertFailsWith<IllegalStateException>(message = "invalid column id") {
+            handler.takeColumn("column no 6")
+        }
+
+    }
+
 }