Support sorting FrameColumn and column of Lists by their size in notebooks

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/KotlinNotebookPluginUtils.kt

@@ -30,6 +30,8 @@ import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.frames
 import org.jetbrains.kotlinx.dataframe.api.getColumn
 import org.jetbrains.kotlinx.dataframe.api.into
+import org.jetbrains.kotlinx.dataframe.api.isFrameColumn
+import org.jetbrains.kotlinx.dataframe.api.isList
 import org.jetbrains.kotlinx.dataframe.api.sortWith
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.api.values
@@ -111,8 +113,15 @@ public object KotlinNotebookPluginUtils {
     private fun createComparator(sortKeys: List<ColumnPath>, isDesc: List<Boolean>): Comparator<DataRow<*>> {
         return Comparator { row1, row2 ->
             for ((key, desc) in sortKeys.zip(isDesc)) {
-                val comparisonResult = if (row1.df().getColumn(key).valuesAreComparable()) {
+                val column = row1.df().getColumn(key)
+                val comparisonResult = if (column.valuesAreComparable()) {
                     compareComparableValues(row1, row2, key, desc)
+                } else if (column.isFrameColumn()) {
+                    val firstValue = column[row1].rowsCount()
+                    val secondValue = column[row2].rowsCount()
+                    firstValue.compare(secondValue, desc)
+                } else if (column.isList()) {
+                    compareListSizes(row1, row2, key, desc)
                 } else {
                     compareStringValues(row1, row2, key, desc)
                 }
@@ -124,6 +133,21 @@ public object KotlinNotebookPluginUtils {
         }
     }
 
+    private fun compareListSizes(
+        row1: DataRow<*>,
+        row2: DataRow<*>,
+        key: ColumnPath,
+        desc: Boolean,
+    ): Int {
+        val firstValue = (row1.getValueOrNull(key) as? List<*>)?.size ?: 0
+        val secondValue = (row2.getValueOrNull(key) as? List<*>)?.size ?: 0
+        return if (desc) {
+            secondValue.compareTo(firstValue)
+        } else {
+            firstValue.compareTo(secondValue)
+        }
+    }
+
     @Suppress("UNCHECKED_CAST")
     private fun compareComparableValues(
         row1: DataRow<*>,
@@ -159,6 +183,8 @@ public object KotlinNotebookPluginUtils {
         }
     }
 
+    private fun <T : Comparable<T>> T.compare(other: T, desc: Boolean) = if (desc) other.compareTo(this) else this.compareTo(other)
+
     internal fun isDataframeConvertable(dataframeLike: Any?): Boolean =
         when (dataframeLike) {
             is Pivot<*>,

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/KotlinNotebookPluginUtilsTests.kt

@@ -0,0 +1,84 @@
+package org.jetbrains.kotlinx.dataframe
+
+import io.kotest.matchers.shouldBe
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
+import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils
+import org.junit.Test
+import kotlin.random.Random
+
+/**
+ * Other tests are located in Jupyter module:
+ * org.jetbrains.kotlinx.dataframe.jupyter.RenderingTests
+ */
+class KotlinNotebookPluginUtilsTests {
+    @Test
+    fun `sort lists by size descending`() {
+        val random = Random(123)
+        val lists = List(20) { List(random.nextInt(1, 100)) { it } } + null
+        val df = dataFrameOf("listColumn" to lists)
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("listColumn")), desc = listOf(true))
+
+        res["listColumn"].values() shouldBe lists.sortedByDescending { it?.size ?: 0 }
+    }
+
+    @Test
+    fun `sort lists by size ascending`() {
+        val lists = listOf(listOf(1, 2, 3), listOf(1), listOf(1, 2), null)
+        val df = dataFrameOf("listColumn" to lists)
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("listColumn")), desc = listOf(false))
+
+        res["listColumn"].values() shouldBe listOf(null, listOf(1), listOf(1, 2), listOf(1, 2, 3))
+    }
+
+    @Test
+    fun `sort empty lists`() {
+        val lists = listOf(listOf(1, 2), emptyList(), listOf(1), emptyList())
+        val df = dataFrameOf("listColumn" to lists)
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("listColumn")), desc = listOf(true))
+
+        res["listColumn"].values() shouldBe listOf(listOf(1, 2), listOf(1), emptyList(), emptyList())
+    }
+
+    @Test
+    fun `sort lists with equal sizes preserves stability`() {
+        val lists = listOf(listOf("a"), listOf("b"), listOf("c"))
+        val df = dataFrameOf("listColumn" to lists)
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("listColumn")), desc = listOf(true))
+
+        res["listColumn"].values() shouldBe lists
+    }
+
+    @Test
+    fun `sort frame column by row count descending`() {
+        val frames = listOf(
+            dataFrameOf("x" to listOf(1)),
+            dataFrameOf("x" to listOf(1, 2, 3)),
+            dataFrameOf("x" to listOf(1, 2)),
+            DataFrame.empty(),
+        )
+        val df = dataFrameOf("nested" to frames.toColumn())
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("nested")), desc = listOf(true))
+
+        res["nested"].values().map { (it as DataFrame<*>).rowsCount() } shouldBe listOf(3, 2, 1, 0)
+    }
+
+    @Test
+    fun `sort frame column by row count ascending`() {
+        val frames = listOf(
+            dataFrameOf("x" to listOf(1, 2, 3)),
+            dataFrameOf("x" to listOf(1)),
+            DataFrame.empty(),
+        )
+        val df = dataFrameOf("nested" to frames.toColumn())
+
+        val res = KotlinNotebookPluginUtils.sortByColumns(df, listOf(listOf("nested")), desc = listOf(false))
+
+        res["nested"].values().map { (it as DataFrame<*>).rowsCount() } shouldBe listOf(0, 1, 3)
+    }
+}