Optimized lazy materialization + parallel sort for Kotlin Notebooks

Author: koperagen

core/build.gradle.kts

@@ -17,6 +17,7 @@ plugins {
         alias(kodex)
         alias(buildconfig)
         alias(binary.compatibility.validator)
+        alias(kotlinx.benchmark)
 
         // generates keywords using the :generator module
         alias(keywordGenerator)
@@ -75,6 +76,7 @@ dependencies {
     testImplementation(libs.kotestAssertions) {
         exclude("org.jetbrains.kotlin", "kotlin-stdlib-jdk8")
     }
+    testImplementation(libs.kotlinx.benchmark.runtime)
     testImplementation(libs.kotlin.scriptingJvm)
     testImplementation(libs.jsoup)
     testImplementation(libs.sl4jsimple)
@@ -89,6 +91,17 @@ dependencies {
     testImplementation(projects.dataframeCsv)
 }
 
+benchmark {
+    targets {
+        register("test")
+    }
+    configurations {
+        register("sort") {
+            include("SortingBenchmark")
+        }
+    }
+}
+
 val samplesImplementation by configurations.getting {
     extendsFrom(configurations.testImplementation.get())
 }

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

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.jupyter
 import org.jetbrains.kotlinx.dataframe.AnyCol
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
+import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.annotations.RequiredByIntellijPlugin
 import org.jetbrains.kotlinx.dataframe.api.Convert
@@ -29,15 +30,17 @@ import org.jetbrains.kotlinx.dataframe.api.at
 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.getRows
 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.rows
 import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.api.values
 import org.jetbrains.kotlinx.dataframe.api.valuesAreComparable
 import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
 import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
+import java.util.Arrays
 
 /**
  * A class with utility methods for Kotlin Notebook Plugin integration.
@@ -105,9 +108,34 @@ public object KotlinNotebookPluginUtils {
             ColumnPath(path)
         }
 
+        if (sortKeys.size == 1) {
+            val column = df.getColumn(sortKeys[0])
+
+            // Not sure how to have generic logic that would produce Comparator<Int> and Comparator<DataRow> without overhead
+            // For now Comparator<DataRow> is needed for fallback case of sorting multiple columns. Although it's now impossible in UI
+            // Please make sure to change both this and createColumnComparator
+            val comparator: Comparator<Int> = when {
+                column.valuesAreComparable() -> compareBy(nullsLast()) {
+                    column[it] as Comparable<Any>?
+                }
+
+                column.isFrameColumn() -> compareBy { column[it].rowsCount() }
+
+                column.isList() -> compareBy { (column[it] as? List<*>)?.size ?: 0 }
+
+                else -> compareBy { column[it]?.toString() ?: "" }
+            }
+
+            val finalComparator = if (isDesc[0]) comparator.reversed() else comparator
+
+            val permutation = Array(column.size()) { it }
+            Arrays.parallelSort(permutation, finalComparator)
+            return SortedDataFrameView(df, permutation.asList())
+        }
+
         val comparator = createComparator(df, sortKeys, isDesc)
 
-        return df.sortWith(comparator)
+        return df.sortWithLazy(comparator)
     }
 
     private fun createComparator(
@@ -153,6 +181,31 @@ public object KotlinNotebookPluginUtils {
         return if (desc) comparator.reversed() else comparator
     }
 
+    private fun <T> DataFrame<T>.sortWithLazy(comparator: Comparator<DataRow<T>>): DataFrame<T> {
+        val permutation = rows().sortedWith(comparator).map { it.index() }
+        return SortedDataFrameView(this, permutation)
+    }
+
+    private class SortedDataFrameView<T>(private val source: DataFrame<T>, private val permutation: List<Int>) :
+        DataFrame<T> by source {
+
+        override operator fun get(index: Int): DataRow<T> = source[permutation[index]]
+
+        override operator fun get(range: IntRange): DataFrame<T> {
+            val indices = range.map { permutation[it] }
+            return source.getRows(indices)
+        }
+
+        override operator fun get(indices: Iterable<Int>): DataFrame<T> {
+            val mappedIndices = indices.map { permutation[it] }
+            return source.getRows(mappedIndices)
+        }
+
+        override fun get(columnName: String): AnyCol {
+            return super.get(columnName)[permutation]
+        }
+    }
+
     internal fun isDataframeConvertable(dataframeLike: Any?): Boolean =
         when (dataframeLike) {
             is Pivot<*>,

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

@@ -0,0 +1,59 @@
+@file:Suppress("EmptyRange")
+
+package org.jetbrains.kotlinx.dataframe
+
+import kotlinx.benchmark.Benchmark
+import kotlinx.benchmark.BenchmarkMode
+import kotlinx.benchmark.BenchmarkTimeUnit
+import kotlinx.benchmark.Measurement
+import kotlinx.benchmark.Mode
+import kotlinx.benchmark.OutputTimeUnit
+import kotlinx.benchmark.Param
+import kotlinx.benchmark.Scope
+import kotlinx.benchmark.Setup
+import kotlinx.benchmark.State
+import kotlinx.benchmark.Warmup
+import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
+import org.jetbrains.kotlinx.dataframe.api.toColumn
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils
+import kotlin.random.Random
+
+@State(Scope.Benchmark)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(BenchmarkTimeUnit.MILLISECONDS)
+open class SortingBenchmark {
+
+    @Param("10000", "100000", "1000000")
+    var size: Int = 0
+
+    @Param("int", "string", "double", "category", "list", "frame")
+    lateinit var columnType: String
+
+    private lateinit var df: DataFrame<*>
+    private lateinit var columnPath: List<String>
+
+    @Setup
+    fun setup() {
+        val random = Random(42)
+        df = (0 until size).toDataFrame {
+            "int" from { it }
+            "string" from { "name_${random.nextInt(1000)}" }
+            "double" from { random.nextDouble() }
+            "category" from { listOf("A", "B", "C", "D").random(random) }
+            "list" from { List(random.nextInt(1, 20)) { "tag$it" } }
+            "frame" from {
+                dataFrameOf("x" to List(random.nextInt(1, 50)) { random.nextInt() }.toColumn())
+            }
+        }
+        columnPath = listOf(columnType)
+    }
+
+    @Benchmark
+    fun sort(): DataFrame<*> {
+        val sorted = KotlinNotebookPluginUtils.sortByColumns(df, listOf(columnPath), listOf(false))
+        return KotlinNotebookPluginUtils.getRowsSubsetForRendering(sorted, 0, 20).value
+    }
+}