Skip to content

Commit

Permalink
Merge pull request #501 from ermolenkodev/KTNB-424
Browse files Browse the repository at this point in the history
Fixes related to Kotlin Notebook plugin integration
  • Loading branch information
ermolenkodev authored Nov 13, 2023
2 parents bf7e469 + d15ecda commit e47bc83
Show file tree
Hide file tree
Showing 8 changed files with 562 additions and 116 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,9 @@ 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.at
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.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
Expand Down Expand Up @@ -340,35 +334,3 @@ 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 Split<*, *> -> dataframeLike.toDataFrame()
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
is MoveClause<*, *> -> dataframeLike.df
is RenameClause<*, *> -> dataframeLike.df
is ReplaceClause<*, *> -> dataframeLike.df
is GroupClause<*, *> -> dataframeLike.into("untitled")
is InsertClause<*> -> dataframeLike.at(0)
is FormatClause<*, *> -> dataframeLike.df
else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,21 @@ package org.jetbrains.kotlinx.dataframe.jupyter
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.DataFrameHtmlData
import org.jetbrains.kotlinx.dataframe.io.DisplayConfiguration
import org.jetbrains.kotlinx.dataframe.io.encodeFrame
import org.jetbrains.kotlinx.dataframe.io.toHTML
import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils.convertToDataFrame
import org.jetbrains.kotlinx.dataframe.nrow
import org.jetbrains.kotlinx.dataframe.size
import org.jetbrains.kotlinx.jupyter.api.*
import org.jetbrains.kotlinx.jupyter.api.HtmlData
import org.jetbrains.kotlinx.jupyter.api.JupyterClientType
import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
import org.jetbrains.kotlinx.jupyter.api.MimeTypedResult
import org.jetbrains.kotlinx.jupyter.api.Notebook
import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
import org.jetbrains.kotlinx.jupyter.api.mimeResult
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"
Expand Down
Original file line number Diff line number Diff line change
@@ -1,14 +1,45 @@
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.api.Convert
import org.jetbrains.kotlinx.dataframe.api.FormatClause
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.GroupClause
import org.jetbrains.kotlinx.dataframe.api.InsertClause
import org.jetbrains.kotlinx.dataframe.api.Merge
import org.jetbrains.kotlinx.dataframe.api.MoveClause
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.RenameClause
import org.jetbrains.kotlinx.dataframe.api.ReplaceClause
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.at
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
import org.jetbrains.kotlinx.dataframe.api.filter
import org.jetbrains.kotlinx.dataframe.api.frames
import org.jetbrains.kotlinx.dataframe.api.into
import org.jetbrains.kotlinx.dataframe.api.sortBy
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
import org.jetbrains.kotlinx.dataframe.api.values
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator

/**
* 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
* Kotlin Notebook Plugin acts as a client of Kotlin Jupyter kernel and uses 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), "")
* The plugin sends the following code to the kernel to evaluate:
* DISPLAY(KotlinNotebooksPluginUtils.getRowsSubsetForRendering(Out[...], 0, 20), "")
*/
public object KotlinNotebookPluginUtils {
/**
Expand All @@ -30,4 +61,110 @@ public object KotlinNotebookPluginUtils {
*/
public fun getRowsSubsetForRendering(df: AnyFrame, startIdx: Int, endIdx: Int): DisableRowsLimitWrapper =
DisableRowsLimitWrapper(df.filter { it.index() in startIdx until endIdx })

/**
* Sorts a dataframe-like object by multiple columns.
*
* @param dataFrameLike The dataframe-like object to sort.
* @param columnPaths The list of columns to sort by. Each element in the list represents a column path
* @param desc The list of booleans indicating whether each column should be sorted in descending order.
* The size of this list should be the same as the size of the `columns` list.
*
* @throws IllegalArgumentException if `dataFrameLike` is `null`.
*
* @return The sorted dataframe.
*/
public fun sortByColumns(
dataFrameLike: Any?,
columnPaths: List<List<String>>,
desc: List<Boolean>
): AnyFrame = when (dataFrameLike) {
null -> throw IllegalArgumentException("Dataframe is null")
else -> sortByColumns(convertToDataFrame(dataFrameLike), columnPaths, desc)
}

/**
* Sorts the given data frame by the specified columns.
*
* @param df The data frame to be sorted.
* @param columnPaths The paths of the columns to be sorted. Each path is represented as a list of strings.
* @param isDesc A list of booleans indicating whether each column should be sorted in descending order.
* The size of this list must be equal to the size of the columnPaths list.
* @return The sorted data frame.
*/
public fun sortByColumns(df: AnyFrame, columnPaths: List<List<String>>, isDesc: List<Boolean>): AnyFrame =
df.sortBy {
require(columnPaths.all { it.isNotEmpty() })
require(columnPaths.size == isDesc.size)

val sortKeys = columnPaths.map { path ->
ColumnPath(path)
}

(sortKeys zip isDesc).map { (key, desc) ->
if (desc) key.desc() else key
}.toColumnSet()
}

/**
* 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.
*/
public fun convertToDataFrame(dataframeLike: Any): AnyFrame =
when (dataframeLike) {
is Pivot<*> -> dataframeLike.frames().toDataFrame()
is ReducedGroupBy<*, *> -> dataframeLike.values()
is ReducedPivot<*> -> dataframeLike.values().toDataFrame()
is PivotGroupBy<*> -> dataframeLike.frames()
is ReducedPivotGroupBy<*> -> dataframeLike.values()
is SplitWithTransform<*, *, *> -> dataframeLike.into()
is Split<*, *> -> dataframeLike.toDataFrame()
is Merge<*, *, *> -> dataframeLike.into(
generateRandomVariationOfColumnName(
"merged",
dataframeLike.df.columnNames()
)
)

is Gather<*, *, *, *> -> dataframeLike.into(
generateRandomVariationOfColumnName("key", dataframeLike.df.columnNames()),
generateRandomVariationOfColumnName("value", dataframeLike.df.columnNames())
)

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
is MoveClause<*, *> -> dataframeLike.df
is RenameClause<*, *> -> dataframeLike.df
is ReplaceClause<*, *> -> dataframeLike.df
is GroupClause<*, *> -> dataframeLike.into(
generateRandomVariationOfColumnName(
"untitled",
dataframeLike.df.columnNames()
)
)

is InsertClause<*> -> dataframeLike.at(0)
is FormatClause<*, *> -> dataframeLike.df
else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
}

/**
* Generates a random variation of a column name that is unique among the provided used names.
*
* @param preferredName The preferred name for the column.
* @param usedNames The list of already used column names.
* @return A unique random variation of the preferred name.
*/
public fun generateRandomVariationOfColumnName(
preferredName: String,
usedNames: List<String> = emptyList()
): String =
ColumnNameGenerator(usedNames).addUnique(preferredName)
}
Loading

0 comments on commit e47bc83

Please sign in to comment.