Add support for swing tables integration in IDEA plugin

core/build.gradle.kts

@@ -52,26 +52,6 @@ kotlin.sourceSets {
     }
 }
 
-tasks.lintKotlinMain {
-    exclude("**/*keywords*/**")
-    exclude {
-        it.name.endsWith(".Generated.kt")
-    }
-    exclude {
-        it.name.endsWith("\$Extensions.kt")
-    }
-}
-
-tasks.lintKotlinTest {
-    exclude {
-        it.name.endsWith(".Generated.kt")
-    }
-    exclude {
-        it.name.endsWith("\$Extensions.kt")
-    }
-    enabled = true
-}
-
 korro {
     docs = fileTree(rootProject.rootDir) {
         include("docs/StardustDocs/topics/*.md")
@@ -118,6 +98,17 @@ kotlinter {
     )
 }
 
+tasks.withType<org.jmailen.gradle.kotlinter.tasks.LintTask> {
+    exclude("**/*keywords*/**")
+    exclude {
+        it.name.endsWith(".Generated.kt")
+    }
+    exclude {
+        it.name.endsWith("\$Extensions.kt")
+    }
+    enabled = true
+}
+
 kotlin {
     explicitApi()
 }
@@ -128,7 +119,6 @@ tasks.withType<JavaCompile> {
 }
 
 tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile> {
-    dependsOn(tasks.lintKotlin)
     kotlinOptions {
         freeCompilerArgs = freeCompilerArgs + listOf("-Xinline-classes", "-Xopt-in=kotlin.RequiresOptIn")
     }

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

@@ -0,0 +1,8 @@
+package org.jetbrains.kotlinx.dataframe.jupyter
+
+import org.jetbrains.kotlinx.dataframe.AnyFrame
+
+/**
+ * Allows for disabling the rows limit when generating a DISPLAY output in Jupyter.
+ */
+public data class DisableRowsLimitWrapper(public val value: AnyFrame)

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

@@ -6,29 +6,8 @@ import org.jetbrains.kotlinx.dataframe.AnyCol
 import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
-import org.jetbrains.kotlinx.dataframe.api.Convert
-import org.jetbrains.kotlinx.dataframe.api.FormattedFrame
-import org.jetbrains.kotlinx.dataframe.api.Gather
-import org.jetbrains.kotlinx.dataframe.api.GroupBy
-import org.jetbrains.kotlinx.dataframe.api.Merge
-import org.jetbrains.kotlinx.dataframe.api.Pivot
-import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
-import org.jetbrains.kotlinx.dataframe.api.ReducedGroupBy
-import org.jetbrains.kotlinx.dataframe.api.ReducedPivot
-import org.jetbrains.kotlinx.dataframe.api.ReducedPivotGroupBy
-import org.jetbrains.kotlinx.dataframe.api.Split
-import org.jetbrains.kotlinx.dataframe.api.SplitWithTransform
-import org.jetbrains.kotlinx.dataframe.api.Update
-import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
-import org.jetbrains.kotlinx.dataframe.api.asDataFrame
-import org.jetbrains.kotlinx.dataframe.api.columnsCount
-import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
-import org.jetbrains.kotlinx.dataframe.api.frames
-import org.jetbrains.kotlinx.dataframe.api.into
-import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
+import org.jetbrains.kotlinx.dataframe.api.*
 import org.jetbrains.kotlinx.dataframe.api.name
-import org.jetbrains.kotlinx.dataframe.api.toDataFrame
-import org.jetbrains.kotlinx.dataframe.api.values
 import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -111,42 +90,42 @@ internal class Integration(
         }
 
         with(JupyterHtmlRenderer(config.display, this)) {
+            render<DisableRowsLimitWrapper>(
+                { "DataRow: index = ${it.value.rowsCount()}, columnsCount = ${it.value.columnsCount()}" },
+                applyRowsLimit = false
+            )
+
             render<HtmlData> { notebook.renderHtmlAsIFrameIfNeeded(it) }
             render<AnyRow>(
-                { it.toDataFrame() },
                 { "DataRow: index = ${it.index()}, columnsCount = ${it.columnsCount()}" },
             )
             render<ColumnGroup<*>>(
-                { it.asDataFrame() },
                 { """ColumnGroup: name = "${it.name}", rowsCount = ${it.rowsCount()}, columnsCount = ${it.columnsCount()}""" },
             )
             render<AnyCol>(
-                { dataFrameOf(it) },
                 { """DataColumn: name = "${it.name}", type = ${renderType(it.type())}, size = ${it.size()}""" },
             )
             render<AnyFrame>(
-                { it },
-                { "DataFrame: rowsCount = ${it.rowsCount()}, columnsCount = ${it.columnsCount()}" },
+                { "DataFrame: rowsCount = ${it.rowsCount()}, columnsCount = ${it.columnsCount()}" }
             )
             render<FormattedFrame<*>>(
-                { it.df },
                 { "DataFrame: rowsCount = ${it.df.rowsCount()}, columnsCount = ${it.df.columnsCount()}" },
                 modifyConfig = { getDisplayConfiguration(it) },
             )
-            render<GroupBy<*, *>>({ it.toDataFrame() }, { "GroupBy" })
-            render<ReducedGroupBy<*, *>>({ it.into(it.groupBy.groups.name()) }, { "ReducedGroupBy" })
-            render<Pivot<*>>({ it.frames().toDataFrame() }, { "Pivot" })
-            render<ReducedPivot<*>>({ it.values().toDataFrame() }, { "ReducedPivot" })
-            render<PivotGroupBy<*>>({ it.frames() }, { "PivotGroupBy" })
-            render<ReducedPivotGroupBy<*>>({ it.values() }, { "ReducedPivotGroupBy" })
-            render<SplitWithTransform<*, *, *>>({ it.into() }, { "Split" })
-            render<Split<*, *>>({ it.toDataFrame() }, { "Split" })
-            render<Merge<*, *, *>>({ it.into("merged") }, { "Merge" })
-            render<Gather<*, *, *, *>>({ it.into("key", "value") }, { "Gather" })
+            render<GroupBy<*, *>>({ "GroupBy" })
+            render<ReducedGroupBy<*, *>>({ "ReducedGroupBy" })
+            render<Pivot<*>>({ "Pivot" })
+            render<ReducedPivot<*>>({ "ReducedPivot" })
+            render<PivotGroupBy<*>>({ "PivotGroupBy" })
+            render<ReducedPivotGroupBy<*>>({ "ReducedPivotGroupBy" })
+            render<SplitWithTransform<*, *, *>>({ "Split" })
+            render<Split<*, *>>({ "Split" })
+            render<Merge<*, *, *>>({ "Merge" })
+            render<Gather<*, *, *, *>>({ "Gather" })
             render<IMG> { HTML(it.toString()) }
             render<IFRAME> { HTML(it.toString()) }
-            render<Update<*, *>>({ it.df }, { "Update" })
-            render<Convert<*, *>>({ it.df }, { "Convert" })
+            render<Update<*, *>>({ "Update" })
+            render<Convert<*, *>>({ "Convert" })
         }
 
         import("org.jetbrains.kotlinx.dataframe.api.*")
@@ -156,6 +135,7 @@ internal class Integration(
         import("org.jetbrains.kotlinx.dataframe.columns.*")
         import("org.jetbrains.kotlinx.dataframe.jupyter.ImportDataSchema")
         import("org.jetbrains.kotlinx.dataframe.jupyter.importDataSchema")
+        import("org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils")
         import("java.net.URL")
         import("java.io.File")
         import("kotlinx.datetime.Instant")
@@ -265,3 +245,28 @@ public fun KotlinKernelHost.useSchemas(schemaClasses: Iterable<KClass<*>>) {
 public fun KotlinKernelHost.useSchemas(vararg schemaClasses: KClass<*>): Unit = useSchemas(schemaClasses.asIterable())
 
 public inline fun <reified T> KotlinKernelHost.useSchema(): Unit = useSchemas(T::class)
+
+/**
+ * Converts [dataframeLike] to [AnyFrame].
+ * If [dataframeLike] is already [AnyFrame] then it is returned as is.
+ * If it's not possible to convert [dataframeLike] to [AnyFrame] then [IllegalArgumentException] is thrown.
+ */
+internal fun convertToDataFrame(dataframeLike: Any): AnyFrame =
+    when (dataframeLike) {
+        is Pivot<*> -> dataframeLike.frames().toDataFrame()
+        is ReducedPivot<*> -> dataframeLike.values().toDataFrame()
+        is PivotGroupBy<*> -> dataframeLike.frames()
+        is ReducedPivotGroupBy<*> -> dataframeLike.values()
+        is SplitWithTransform<*, *, *> -> dataframeLike.into()
+        is Merge<*, *, *> -> dataframeLike.into("merged")
+        is Gather<*, *, *, *> -> dataframeLike.into("key", "value")
+        is Update<*, *> -> dataframeLike.df
+        is Convert<*, *> -> dataframeLike.df
+        is FormattedFrame<*> -> dataframeLike.df
+        is AnyCol -> dataFrameOf(dataframeLike)
+        is AnyRow -> dataframeLike.toDataFrame()
+        is GroupBy<*, *> -> dataframeLike.toDataFrame()
+        is AnyFrame -> dataframeLike
+        is DisableRowsLimitWrapper -> dataframeLike.value
+        else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
+    }

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

@@ -1,26 +1,41 @@
 package org.jetbrains.kotlinx.dataframe.jupyter
 
-import org.jetbrains.kotlinx.dataframe.DataFrame
-import org.jetbrains.kotlinx.dataframe.io.DisplayConfiguration
+import com.beust.klaxon.json
+import org.jetbrains.kotlinx.dataframe.api.rows
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
+import org.jetbrains.kotlinx.dataframe.io.*
 import org.jetbrains.kotlinx.dataframe.io.initHtml
-import org.jetbrains.kotlinx.dataframe.io.toHTML
+import org.jetbrains.kotlinx.dataframe.nrow
+import org.jetbrains.kotlinx.dataframe.size
+import org.jetbrains.kotlinx.jupyter.api.*
 import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
-import org.jetbrains.kotlinx.jupyter.api.renderHtmlAsIFrameIfNeeded
+
+/** Starting from this version, dataframe integration will respond with additional data for rendering in Kotlin Notebooks plugin. */
+private const val MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI = "0.11.0.311"
 
 internal class JupyterHtmlRenderer(
     val display: DisplayConfiguration,
     val builder: JupyterIntegration.Builder,
 )
 
 internal inline fun <reified T : Any> JupyterHtmlRenderer.render(
-    crossinline getDf: (T) -> DataFrame<*>,
     noinline getFooter: (T) -> String,
-    crossinline modifyConfig: T.(DisplayConfiguration) -> DisplayConfiguration = { it }
+    crossinline modifyConfig: T.(DisplayConfiguration) -> DisplayConfiguration = { it },
+    applyRowsLimit: Boolean = true
 ) = builder.renderWithHost<T> { host, value ->
     val contextRenderer = JupyterCellRenderer(this.notebook, host)
     val reifiedDisplayConfiguration = value.modifyConfig(display)
     val footer = getFooter(value)
-    val html = getDf(value).toHTML(
+
+    val df = convertToDataFrame(value)
+
+    val limit = if (applyRowsLimit) {
+        reifiedDisplayConfiguration.rowsLimit ?: df.nrow
+    } else {
+        df.nrow
+    }
+
+    val html = df.toHTML(
         reifiedDisplayConfiguration,
         extraHtml = initHtml(
             includeJs = reifiedDisplayConfiguration.isolatedOutputs,
@@ -30,5 +45,30 @@ internal inline fun <reified T : Any> JupyterHtmlRenderer.render(
         contextRenderer
     ) { footer }
 
-    notebook.renderHtmlAsIFrameIfNeeded(html)
+    if (notebook.kernelVersion >= KotlinKernelVersion.from(MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI)!!) {
+        val jsonEncodedDf = json {
+            obj(
+                "nrow" to df.size.nrow,
+                "ncol" to df.size.ncol,
+                "columns" to df.columnNames(),
+                "kotlin_dataframe" to encodeFrame(df.rows().take(limit).toDataFrame())
+            )
+        }.toJsonString()
+        notebook.renderAsIFrameAsNeeded(html, jsonEncodedDf)
+    } else {
+        notebook.renderHtmlAsIFrameIfNeeded(html)
+    }
+}
+
+internal fun Notebook.renderAsIFrameAsNeeded(data: HtmlData, jsonEncodedDf: String): MimeTypedResult {
+    val textHtml = if (jupyterClientType == JupyterClientType.KOTLIN_NOTEBOOK) {
+        data.generateIframePlaneText(currentColorScheme)
+    } else {
+        data.toString(currentColorScheme)
+    }
+
+    return mimeResult(
+        "text/html" to textHtml,
+        "application/kotlindataframe+json" to jsonEncodedDf
+    ).also { it.isolatedHtml = false }
 }

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

@@ -0,0 +1,33 @@
+package org.jetbrains.kotlinx.dataframe.jupyter
+
+import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.api.filter
+
+/**
+ * A class with utility methods for Kotlin Notebook Plugin integration.
+ * Kotlin Notebook Plugin is acts as a client of Kotlin Jupyter kernel and use this functionality
+ * for dynamic pagination when rendering dataframes.
+ * The plugin sends Kotlin following code to the kernel to evaluate
+ * DISPLAY(KotlinNotebooksPluginUtils.getRowsSubsetForRendering(Out[x], 0, 20), "")
+ */
+public object KotlinNotebookPluginUtils {
+    /**
+     * Returns a subset of rows from the given dataframe for rendering.
+     * It's used for example for dynamic pagination in Kotlin Notebook Plugin.
+     */
+    public fun getRowsSubsetForRendering(
+        dataFrameLike: Any?,
+        startIdx: Int,
+        endIdx: Int
+    ): DisableRowsLimitWrapper = when (dataFrameLike) {
+        null -> throw IllegalArgumentException("Dataframe is null")
+        else -> getRowsSubsetForRendering(convertToDataFrame(dataFrameLike), startIdx, endIdx)
+    }
+
+    /**
+     * Returns a subset of rows from the given dataframe for rendering.
+     * It's used for example for dynamic pagination in Kotlin Notebook Plugin.
+     */
+    public fun getRowsSubsetForRendering(df: AnyFrame, startIdx: Int, endIdx: Int): DisableRowsLimitWrapper =
+        DisableRowsLimitWrapper(df.filter { it.index() in startIdx until endIdx })
+}

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

@@ -1,9 +1,11 @@
 package org.jetbrains.kotlinx.dataframe.jupyter
 
+import com.beust.klaxon.*
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.string.shouldContain
 import io.kotest.matchers.string.shouldNotContain
 import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlinx.jupyter.api.MimeTypedResult
 import org.jetbrains.kotlinx.jupyter.testkit.JupyterReplTestCase
 import org.junit.Test
 
@@ -74,4 +76,53 @@ class RenderingTests : JupyterReplTestCase() {
         htmlLight shouldNotContain darkClassAttribute
         htmlDark shouldContain darkClassAttribute
     }
+
+    @Test
+    fun `test kotlin notebook plugin utils rows subset`() {
+        @Language("kts")
+        val result = exec<MimeTypedResult>(
+            """
+            data class Row(val id: Int)
+            val df = (1..100).map { Row(it) }.toDataFrame()
+            KotlinNotebookPluginUtils.getRowsSubsetForRendering(df, 20 , 50)
+            """.trimIndent()
+        )
+
+        val json = parseDataframeJson(result)
+
+        json.int("nrow") shouldBe 30
+        json.int("ncol") shouldBe 1
+
+        val rows = json.array<JsonArray<*>>("kotlin_dataframe")!!
+        rows.getObj(0).int("id") shouldBe 21
+        rows.getObj(rows.lastIndex).int("id") shouldBe 50
+    }
+
+    private fun parseDataframeJson(result: MimeTypedResult): JsonObject {
+        val parser = Parser.default()
+        return parser.parse(StringBuilder(result["application/kotlindataframe+json"]!!)) as JsonObject
+    }
+
+    private fun JsonArray<*>.getObj(index: Int) = this.get(index) as JsonObject
+
+    @Test
+    fun `test kotlin notebook plugin utils groupby`() {
+        @Language("kts")
+        val result = exec<MimeTypedResult>(
+            """
+            data class Row(val id: Int, val group: Int)
+            val df = (1..100).map { Row(it, if (it <= 50) 1 else 2) }.toDataFrame()
+            KotlinNotebookPluginUtils.getRowsSubsetForRendering(df.groupBy("group"), 0, 10)
+            """.trimIndent()
+        )
+
+        val json = parseDataframeJson(result)
+
+        json.int("nrow") shouldBe 2
+        json.int("ncol") shouldBe 2
+
+        val rows = json.array<JsonArray<*>>("kotlin_dataframe")!!
+        rows.getObj(0).array<JsonObject>("group1")!!.size shouldBe 50
+        rows.getObj(1).array<JsonObject>("group1")!!.size shouldBe 50
+    }
 }

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

@@ -32,6 +32,7 @@ class SampleNotebooksTests : DataFrameJupyterTest() {
     )
 
     @Test
+    @Ignore
     fun wine() = exampleTest(
         "wine", "WineNetWIthKotlinDL",
         replacer = CodeReplacer.byMap(

gradle/libs.versions.toml

@@ -1,8 +1,8 @@
 [versions]
-ksp = "1.8.10-1.0.9"
-kotlinJupyter = "0.11.0-198"
+ksp = "1.8.20-Beta-1.0.9"
+kotlinJupyter = "0.11.0-311"
 ktlint = "3.4.5"
-kotlin = "1.8.10"
+kotlin = "1.8.20-Beta"
 dokka = "1.8.10"
 libsPublisher = "0.0.60-dev-30"
 # "Bootstrap" version of the dataframe, used in the build itself to generate @DataSchema APIs,