Skip to content

Commit 12deade

Browse files
ermolenkodevermolenkodev
committed
Add unit tests for sortByColumns in RenderingTests.kt.
1 parent 675cd69 commit 12deade

File tree

2 files changed

+354
-0
lines changed

2 files changed

+354
-0
lines changed

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/RenderingTests.kt

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,183 @@ class RenderingTests : JupyterReplTestCase() {
406406
}
407407
}
408408

409+
@Test
410+
fun `test sortByColumns by int column`() {
411+
val json = executeScriptAndParseDataframeResult(
412+
"""
413+
val df = dataFrameOf("nums")(5, 4, 3, 2, 1)
414+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("nums")), listOf(false))
415+
KotlinNotebookPluginUtils.convertToDataFrame(res)
416+
""".trimIndent(),
417+
)
418+
419+
val rows = json[KOTLIN_DATAFRAME]!!.jsonArray
420+
json.extractColumn<Int>(0, "nums") shouldBe 1
421+
json.extractColumn<Int>(rows.size - 1, "nums") shouldBe 5
422+
}
423+
424+
internal inline fun <reified T> JsonObject.extractColumn(index: Int, fieldName: String): T {
425+
val element = this[KOTLIN_DATAFRAME]!!.jsonArray[index].jsonObject[fieldName]!!.jsonPrimitive
426+
return when (T::class) {
427+
String::class -> element.content as T
428+
Int::class -> element.int as T
429+
else -> throw IllegalArgumentException("Unsupported type")
430+
}
431+
}
432+
433+
@Test
434+
fun `test sortByColumns by multiple int columns`() {
435+
val json = executeScriptAndParseDataframeResult(
436+
"""
437+
data class Row(val a: Int, val b: Int)
438+
val df = listOf(Row(1, 1), Row(1, 2), Row(2, 3), Row(2, 4), Row(3, 5), Row(3, 6)).toDataFrame()
439+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("a"), listOf("b")), listOf(true, false))
440+
KotlinNotebookPluginUtils.convertToDataFrame(res)
441+
""".trimIndent(),
442+
)
443+
444+
json.extractColumn<Int>(0, "a") shouldBe 3
445+
json.extractColumn<Int>(0, "b") shouldBe 5
446+
json.extractColumn<Int>(5, "a") shouldBe 1
447+
json.extractColumn<Int>(5, "b") shouldBe 2
448+
}
449+
450+
@Test
451+
fun `test sortByColumns by single string column`() {
452+
val json = executeScriptAndParseDataframeResult(
453+
"""
454+
val df = dataFrameOf("letters")("e", "d", "c", "b", "a")
455+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("letters")), listOf(true))
456+
KotlinNotebookPluginUtils.convertToDataFrame(res)
457+
""".trimIndent(),
458+
)
459+
460+
json.extractColumn<String>(0, "letters") shouldBe "e"
461+
json.extractColumn<String>(4, "letters") shouldBe "a"
462+
}
463+
464+
@Test
465+
fun `test sortByColumns by multiple string columns`() {
466+
val json = executeScriptAndParseDataframeResult(
467+
"""
468+
data class Row(val first: String, val second: String)
469+
val df = listOf(Row("a", "b"), Row("a", "a"), Row("b", "b"), Row("b", "a")).toDataFrame()
470+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("first"), listOf("second")), listOf(false, true))
471+
KotlinNotebookPluginUtils.convertToDataFrame(res)
472+
""".trimIndent(),
473+
)
474+
475+
json.extractColumn<String>(0, "first") shouldBe "a"
476+
json.extractColumn<String>(0, "second") shouldBe "b"
477+
json.extractColumn<String>(3, "first") shouldBe "b"
478+
json.extractColumn<String>(3, "second") shouldBe "a"
479+
}
480+
481+
@Test
482+
fun `test sortByColumns by mix of int and string columns`() {
483+
val json = executeScriptAndParseDataframeResult(
484+
"""
485+
data class Row(val num: Int, val letter: String)
486+
val df = listOf(Row(1, "a"), Row(1, "b"), Row(2, "a"), Row(2, "b"), Row(3, "a")).toDataFrame()
487+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("num"), listOf("letter")), listOf(true, false))
488+
KotlinNotebookPluginUtils.convertToDataFrame(res)
489+
""".trimIndent(),
490+
)
491+
492+
json.extractColumn<Int>(0, "num") shouldBe 3
493+
json.extractColumn<String>(0, "letter") shouldBe "a"
494+
json.extractColumn<Int>(4, "num") shouldBe 1
495+
json.extractColumn<String>(4, "letter") shouldBe "b"
496+
}
497+
498+
@Test
499+
fun `test sortByColumns by multiple non-comparable column`() {
500+
val json = executeScriptAndParseDataframeResult(
501+
"""
502+
data class Person(val name: String, val age: Int) {
503+
override fun toString(): String {
504+
return age.toString()
505+
}
506+
}
507+
val df = dataFrameOf("urls", "person")(
508+
URL("https://example.com/a"), Person("Alice", 10),
509+
URL("https://example.com/b"), Person("Bob", 11),
510+
URL("https://example.com/a"), Person("Nick", 12),
511+
URL("https://example.com/b"), Person("Guy", 13),
512+
)
513+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("urls"), listOf("person")), listOf(false, true))
514+
KotlinNotebookPluginUtils.convertToDataFrame(res)
515+
""".trimIndent(),
516+
)
517+
518+
json.extractColumn<Int>(0, "person") shouldBe 12
519+
json.extractColumn<Int>(3, "person") shouldBe 11
520+
}
521+
522+
@Test
523+
fun `test sortByColumns by mix of comparable and non-comparable columns`() {
524+
val json = executeScriptAndParseDataframeResult(
525+
"""
526+
val df = dataFrameOf("urls", "id")(
527+
URL("https://example.com/a"), 1,
528+
URL("https://example.com/b"), 2,
529+
URL("https://example.com/a"), 2,
530+
URL("https://example.com/b"), 1,
531+
)
532+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("id"), listOf("urls")), listOf(true, true))
533+
KotlinNotebookPluginUtils.convertToDataFrame(res)
534+
""".trimIndent(),
535+
)
536+
537+
json.extractColumn<String>(0, "urls") shouldBe "https://example.com/b"
538+
json.extractColumn<Int>(0, "id") shouldBe 2
539+
json.extractColumn<String>(3, "urls") shouldBe "https://example.com/a"
540+
json.extractColumn<Int>(3, "id") shouldBe 1
541+
}
542+
543+
@Test
544+
fun `test sortByColumns by url column`() {
545+
val json = executeScriptAndParseDataframeResult(
546+
"""
547+
val df = dataFrameOf("urls")(
548+
URL("https://example.com/a"),
549+
URL("https://example.com/c"),
550+
URL("https://example.com/b"),
551+
URL("https://example.com/d")
552+
)
553+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("urls")), listOf(false))
554+
KotlinNotebookPluginUtils.convertToDataFrame(res)
555+
""".trimIndent(),
556+
)
557+
558+
json.extractColumn<String>(0, "urls") shouldBe "https://example.com/a"
559+
json.extractColumn<String>(1, "urls") shouldBe "https://example.com/b"
560+
json.extractColumn<String>(2, "urls") shouldBe "https://example.com/c"
561+
json.extractColumn<String>(3, "urls") shouldBe "https://example.com/d"
562+
}
563+
564+
@Test
565+
fun `test sortByColumns by column group children`() {
566+
val json = executeScriptAndParseDataframeResult(
567+
"""
568+
val df = dataFrameOf(
569+
"a" to listOf(5, 4, 3, 2, 1),
570+
"b" to listOf(1, 2, 3, 4, 5)
571+
)
572+
val res = KotlinNotebookPluginUtils.sortByColumns(df.group("a", "b").into("c"), listOf(listOf("c", "a")), listOf(false))
573+
KotlinNotebookPluginUtils.convertToDataFrame(res)
574+
""".trimIndent(),
575+
)
576+
577+
fun JsonObject.extractBFields(): List<Int> {
578+
val dataframe = this[KOTLIN_DATAFRAME]!!.jsonArray
579+
return dataframe.map { it.jsonObject["c"]!!.jsonObject["data"]!!.jsonObject["b"]!!.jsonPrimitive.int }
580+
}
581+
582+
val bFields = json.extractBFields()
583+
bFields shouldBe listOf(5, 4, 3, 2, 1)
584+
}
585+
409586
companion object {
410587
/**
411588
* Set the system property for the IDE version needed for specific serialization testing purposes.

‎core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/RenderingTests.kt

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,183 @@ class RenderingTests : JupyterReplTestCase() {
406406
}
407407
}
408408

409+
@Test
410+
fun `test sortByColumns by int column`() {
411+
val json = executeScriptAndParseDataframeResult(
412+
"""
413+
val df = dataFrameOf("nums")(5, 4, 3, 2, 1)
414+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("nums")), listOf(false))
415+
KotlinNotebookPluginUtils.convertToDataFrame(res)
416+
""".trimIndent(),
417+
)
418+
419+
val rows = json[KOTLIN_DATAFRAME]!!.jsonArray
420+
json.extractColumn<Int>(0, "nums") shouldBe 1
421+
json.extractColumn<Int>(rows.size - 1, "nums") shouldBe 5
422+
}
423+
424+
internal inline fun <reified T> JsonObject.extractColumn(index: Int, fieldName: String): T {
425+
val element = this[KOTLIN_DATAFRAME]!!.jsonArray[index].jsonObject[fieldName]!!.jsonPrimitive
426+
return when (T::class) {
427+
String::class -> element.content as T
428+
Int::class -> element.int as T
429+
else -> throw IllegalArgumentException("Unsupported type")
430+
}
431+
}
432+
433+
@Test
434+
fun `test sortByColumns by multiple int columns`() {
435+
val json = executeScriptAndParseDataframeResult(
436+
"""
437+
data class Row(val a: Int, val b: Int)
438+
val df = listOf(Row(1, 1), Row(1, 2), Row(2, 3), Row(2, 4), Row(3, 5), Row(3, 6)).toDataFrame()
439+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("a"), listOf("b")), listOf(true, false))
440+
KotlinNotebookPluginUtils.convertToDataFrame(res)
441+
""".trimIndent(),
442+
)
443+
444+
json.extractColumn<Int>(0, "a") shouldBe 3
445+
json.extractColumn<Int>(0, "b") shouldBe 5
446+
json.extractColumn<Int>(5, "a") shouldBe 1
447+
json.extractColumn<Int>(5, "b") shouldBe 2
448+
}
449+
450+
@Test
451+
fun `test sortByColumns by single string column`() {
452+
val json = executeScriptAndParseDataframeResult(
453+
"""
454+
val df = dataFrameOf("letters")("e", "d", "c", "b", "a")
455+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("letters")), listOf(true))
456+
KotlinNotebookPluginUtils.convertToDataFrame(res)
457+
""".trimIndent(),
458+
)
459+
460+
json.extractColumn<String>(0, "letters") shouldBe "e"
461+
json.extractColumn<String>(4, "letters") shouldBe "a"
462+
}
463+
464+
@Test
465+
fun `test sortByColumns by multiple string columns`() {
466+
val json = executeScriptAndParseDataframeResult(
467+
"""
468+
data class Row(val first: String, val second: String)
469+
val df = listOf(Row("a", "b"), Row("a", "a"), Row("b", "b"), Row("b", "a")).toDataFrame()
470+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("first"), listOf("second")), listOf(false, true))
471+
KotlinNotebookPluginUtils.convertToDataFrame(res)
472+
""".trimIndent(),
473+
)
474+
475+
json.extractColumn<String>(0, "first") shouldBe "a"
476+
json.extractColumn<String>(0, "second") shouldBe "b"
477+
json.extractColumn<String>(3, "first") shouldBe "b"
478+
json.extractColumn<String>(3, "second") shouldBe "a"
479+
}
480+
481+
@Test
482+
fun `test sortByColumns by mix of int and string columns`() {
483+
val json = executeScriptAndParseDataframeResult(
484+
"""
485+
data class Row(val num: Int, val letter: String)
486+
val df = listOf(Row(1, "a"), Row(1, "b"), Row(2, "a"), Row(2, "b"), Row(3, "a")).toDataFrame()
487+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("num"), listOf("letter")), listOf(true, false))
488+
KotlinNotebookPluginUtils.convertToDataFrame(res)
489+
""".trimIndent(),
490+
)
491+
492+
json.extractColumn<Int>(0, "num") shouldBe 3
493+
json.extractColumn<String>(0, "letter") shouldBe "a"
494+
json.extractColumn<Int>(4, "num") shouldBe 1
495+
json.extractColumn<String>(4, "letter") shouldBe "b"
496+
}
497+
498+
@Test
499+
fun `test sortByColumns by multiple non-comparable column`() {
500+
val json = executeScriptAndParseDataframeResult(
501+
"""
502+
data class Person(val name: String, val age: Int) {
503+
override fun toString(): String {
504+
return age.toString()
505+
}
506+
}
507+
val df = dataFrameOf("urls", "person")(
508+
URL("https://example.com/a"), Person("Alice", 10),
509+
URL("https://example.com/b"), Person("Bob", 11),
510+
URL("https://example.com/a"), Person("Nick", 12),
511+
URL("https://example.com/b"), Person("Guy", 13),
512+
)
513+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("urls"), listOf("person")), listOf(false, true))
514+
KotlinNotebookPluginUtils.convertToDataFrame(res)
515+
""".trimIndent(),
516+
)
517+
518+
json.extractColumn<Int>(0, "person") shouldBe 12
519+
json.extractColumn<Int>(3, "person") shouldBe 11
520+
}
521+
522+
@Test
523+
fun `test sortByColumns by mix of comparable and non-comparable columns`() {
524+
val json = executeScriptAndParseDataframeResult(
525+
"""
526+
val df = dataFrameOf("urls", "id")(
527+
URL("https://example.com/a"), 1,
528+
URL("https://example.com/b"), 2,
529+
URL("https://example.com/a"), 2,
530+
URL("https://example.com/b"), 1,
531+
)
532+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("id"), listOf("urls")), listOf(true, true))
533+
KotlinNotebookPluginUtils.convertToDataFrame(res)
534+
""".trimIndent(),
535+
)
536+
537+
json.extractColumn<String>(0, "urls") shouldBe "https://example.com/b"
538+
json.extractColumn<Int>(0, "id") shouldBe 2
539+
json.extractColumn<String>(3, "urls") shouldBe "https://example.com/a"
540+
json.extractColumn<Int>(3, "id") shouldBe 1
541+
}
542+
543+
@Test
544+
fun `test sortByColumns by url column`() {
545+
val json = executeScriptAndParseDataframeResult(
546+
"""
547+
val df = dataFrameOf("urls")(
548+
URL("https://example.com/a"),
549+
URL("https://example.com/c"),
550+
URL("https://example.com/b"),
551+
URL("https://example.com/d")
552+
)
553+
val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("urls")), listOf(false))
554+
KotlinNotebookPluginUtils.convertToDataFrame(res)
555+
""".trimIndent(),
556+
)
557+
558+
json.extractColumn<String>(0, "urls") shouldBe "https://example.com/a"
559+
json.extractColumn<String>(1, "urls") shouldBe "https://example.com/b"
560+
json.extractColumn<String>(2, "urls") shouldBe "https://example.com/c"
561+
json.extractColumn<String>(3, "urls") shouldBe "https://example.com/d"
562+
}
563+
564+
@Test
565+
fun `test sortByColumns by column group children`() {
566+
val json = executeScriptAndParseDataframeResult(
567+
"""
568+
val df = dataFrameOf(
569+
"a" to listOf(5, 4, 3, 2, 1),
570+
"b" to listOf(1, 2, 3, 4, 5)
571+
)
572+
val res = KotlinNotebookPluginUtils.sortByColumns(df.group("a", "b").into("c"), listOf(listOf("c", "a")), listOf(false))
573+
KotlinNotebookPluginUtils.convertToDataFrame(res)
574+
""".trimIndent(),
575+
)
576+
577+
fun JsonObject.extractBFields(): List<Int> {
578+
val dataframe = this[KOTLIN_DATAFRAME]!!.jsonArray
579+
return dataframe.map { it.jsonObject["c"]!!.jsonObject["data"]!!.jsonObject["b"]!!.jsonPrimitive.int }
580+
}
581+
582+
val bFields = json.extractBFields()
583+
bFields shouldBe listOf(5, 4, 3, 2, 1)
584+
}
585+
409586
companion object {
410587
/**
411588
* Set the system property for the IDE version needed for specific serialization testing purposes.

0 commit comments

Comments
 (0)