Skip to content

Commit 6a999b0

Browse files
ermolenkodevermolenkodev
committed
Apply DataFrameConvertible encoding only for supported IDE versions
Also refactor JSON encoding to use the unified EncodingOptions interface
1 parent 9507ae1 commit 6a999b0

File tree

9 files changed

+771
-74
lines changed

9 files changed

+771
-74
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt

Lines changed: 38 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.TYPES
3838
import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.VERSION
3939
import org.jetbrains.kotlinx.dataframe.io.ARRAY_COLUMN_NAME
4040
import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions
41+
import org.jetbrains.kotlinx.dataframe.io.DataframeConvertableEncodingOptions
42+
import org.jetbrains.kotlinx.dataframe.io.EncodingOptions
4143
import org.jetbrains.kotlinx.dataframe.io.VALUE_COLUMN_NAME
4244
import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils
4345
import org.jetbrains.kotlinx.dataframe.jupyter.KotlinNotebookPluginUtils.isDataframeConvertable
@@ -111,14 +113,14 @@ internal fun encodeRowWithMetadata(
111113
frame: ColumnsContainer<*>,
112114
index: Int,
113115
rowLimit: Int? = null,
114-
imageEncodingOptions: Base64ImageEncodingOptions? = null,
116+
encodingOptions: List<EncodingOptions>,
115117
): JsonElement? {
116118
val values: List<Pair<String, JsonElement>> = frame.columns().map { col ->
117119
when (col) {
118120
is ColumnGroup<*> -> {
119121
val schema = col.schema()
120122
buildJsonObject {
121-
put(DATA, encodeRowWithMetadata(col, index, rowLimit, imageEncodingOptions) ?: JsonPrimitive(null))
123+
put(DATA, encodeRowWithMetadata(col, index, rowLimit, encodingOptions) ?: JsonPrimitive(null))
122124
putJsonObject(METADATA) {
123125
put(KIND, JsonPrimitive(ColumnKind.Group.toString()))
124126
put(COLUMNS, Json.encodeToJsonElement(schema.columns.keys))
@@ -135,9 +137,9 @@ internal fun encodeRowWithMetadata(
135137

136138
is FrameColumn<*> -> {
137139
val data = if (rowLimit == null) {
138-
encodeFrameWithMetadata(col[index], null, imageEncodingOptions)
140+
encodeFrameWithMetadata(col[index], null, encodingOptions)
139141
} else {
140-
encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit, imageEncodingOptions)
142+
encodeFrameWithMetadata(col[index].take(rowLimit), rowLimit, encodingOptions)
141143
}
142144
val schema = col.schema.value
143145
buildJsonObject {
@@ -158,34 +160,32 @@ internal fun encodeRowWithMetadata(
158160
}
159161
}
160162

161-
else -> encodeValue(col, index, imageEncodingOptions)
163+
else -> encodeValue(col, index, encodingOptions)
162164
}.let { col.name to it }
163165
}
164166
if (values.isEmpty()) return null
165167
return JsonObject(values.toMap())
166168
}
167169

168-
internal fun encodeValue(
169-
col: AnyCol,
170-
index: Int,
171-
imageEncodingOptions: Base64ImageEncodingOptions? = null,
172-
): JsonElement =
170+
internal fun encodeValue(col: AnyCol, index: Int, encodingOptions: List<EncodingOptions>): JsonElement =
173171
when {
174-
isDataframeConvertable(col[index]) -> if (col[index] == null) {
175-
JsonPrimitive(null)
176-
} else {
177-
val data = encodeFrameWithMetadata(
178-
KotlinNotebookPluginUtils.convertToDataFrame(col[index]!!),
179-
null,
180-
imageEncodingOptions,
181-
)
182-
buildJsonObject {
183-
put(DATA, data)
184-
putJsonObject(METADATA) {
185-
put(KIND, JsonPrimitive(CellKind.DataFrameConvertable.toString()))
172+
isDataframeConvertable(col[index]) && encodingOptions.get<DataframeConvertableEncodingOptions>() != null ->
173+
if (col[index] == null) {
174+
JsonPrimitive(null)
175+
} else {
176+
val options = encodingOptions.get<DataframeConvertableEncodingOptions>()!!
177+
val data = encodeFrameWithMetadata(
178+
KotlinNotebookPluginUtils.convertToDataFrame(col[index]!!),
179+
options.rowsLimit,
180+
encodingOptions,
181+
)
182+
buildJsonObject {
183+
put(DATA, data)
184+
putJsonObject(METADATA) {
185+
put(KIND, JsonPrimitive(CellKind.DataFrameConvertable.toString()))
186+
}
186187
}
187188
}
188-
}
189189

190190
col.isList() -> col[index]?.let { list ->
191191
val values = (list as List<*>).map { convert(it) }
@@ -194,14 +194,22 @@ internal fun encodeValue(
194194

195195
col.typeClass in valueTypes -> convert(col[index])
196196

197-
col.typeClass == BufferedImage::class && imageEncodingOptions != null ->
197+
col.typeClass == BufferedImage::class && encodingOptions.get<Base64ImageEncodingOptions>() != null ->
198198
col[index]?.let { image ->
199-
JsonPrimitive(encodeBufferedImageAsBase64(image as BufferedImage, imageEncodingOptions))
199+
JsonPrimitive(
200+
encodeBufferedImageAsBase64(
201+
image as BufferedImage,
202+
encodingOptions.get<Base64ImageEncodingOptions>()!!,
203+
),
204+
)
200205
} ?: JsonPrimitive("")
201206

202207
else -> JsonPrimitive(col[index]?.toString())
203208
}
204209

210+
@Suppress("UNCHECKED_CAST")
211+
private inline fun <reified T : EncodingOptions> List<EncodingOptions>.get(): T? = this.find { it is T } as T?
212+
205213
private fun encodeBufferedImageAsBase64(
206214
image: BufferedImage,
207215
imageEncodingOptions: Base64ImageEncodingOptions = Base64ImageEncodingOptions(),
@@ -236,7 +244,7 @@ private fun createJsonTypeDescriptor(columnSchema: ColumnSchema): JsonObject =
236244
internal fun encodeFrameWithMetadata(
237245
frame: AnyFrame,
238246
rowLimit: Int? = null,
239-
imageEncodingOptions: Base64ImageEncodingOptions? = null,
247+
encodingOptions: List<EncodingOptions>,
240248
): JsonArray {
241249
val valueColumn = frame.extractValueColumn()
242250
val arrayColumn = frame.extractArrayColumn()
@@ -250,13 +258,13 @@ internal fun encodeFrameWithMetadata(
250258
encodeFrameWithMetadata(
251259
it as AnyFrame,
252260
rowLimit,
253-
imageEncodingOptions,
261+
encodingOptions,
254262
)
255263
} else {
256264
null
257265
}
258266
}
259-
?: encodeRowWithMetadata(frame, rowIndex, rowLimit, imageEncodingOptions)
267+
?: encodeRowWithMetadata(frame, rowIndex, rowLimit, encodingOptions)
260268
}
261269

262270
return buildJsonArray { addAll(data.map { convert(it) }) }
@@ -364,7 +372,7 @@ internal fun encodeDataFrameWithMetadata(
364372
frame: AnyFrame,
365373
rowLimit: Int,
366374
nestedRowLimit: Int? = null,
367-
imageEncodingOptions: Base64ImageEncodingOptions? = null,
375+
encodingOptions: List<EncodingOptions>,
368376
): JsonObject =
369377
buildJsonObject {
370378
put(VERSION, JsonPrimitive(SERIALIZATION_VERSION))
@@ -385,7 +393,7 @@ internal fun encodeDataFrameWithMetadata(
385393
encodeFrameWithMetadata(
386394
frame = frame.take(rowLimit),
387395
rowLimit = nestedRowLimit,
388-
imageEncodingOptions = imageEncodingOptions,
396+
encodingOptions = encodingOptions,
389397
),
390398
)
391399
}

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/json.kt

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ public fun AnyFrame.toJsonWithMetadata(
303303
rowLimit: Int,
304304
nestedRowLimit: Int? = null,
305305
prettyPrint: Boolean = false,
306-
imageEncodingOptions: Base64ImageEncodingOptions? = null,
306+
encodingOptions: List<EncodingOptions> = emptyList(),
307307
): String {
308308
val json = Json {
309309
this.prettyPrint = prettyPrint
@@ -312,12 +312,19 @@ public fun AnyFrame.toJsonWithMetadata(
312312
}
313313
return json.encodeToString(
314314
JsonElement.serializer(),
315-
encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit, imageEncodingOptions),
315+
encodeDataFrameWithMetadata(this@toJsonWithMetadata, rowLimit, nestedRowLimit, encodingOptions),
316316
)
317317
}
318318

319319
internal const val DEFAULT_IMG_SIZE = 600
320320

321+
/**
322+
* Interface representing encoding options that can be applied when converting a data structure to JSON format.
323+
* Implementations of this interface can provide specific behaviors for encoding various data types,
324+
* such as images or data frames, when serializing to JSON.
325+
*/
326+
public interface EncodingOptions
327+
321328
/**
322329
* Class representing the options for encoding images.
323330
*
@@ -327,7 +334,7 @@ internal const val DEFAULT_IMG_SIZE = 600
327334
public class Base64ImageEncodingOptions(
328335
public val imageSizeLimit: Int = DEFAULT_IMG_SIZE,
329336
private val options: Int = GZIP_ON or LIMIT_SIZE_ON,
330-
) {
337+
) : EncodingOptions {
331338
public val isGzipOn: Boolean
332339
get() = options and GZIP_ON == GZIP_ON
333340

@@ -341,6 +348,14 @@ public class Base64ImageEncodingOptions(
341348
}
342349
}
343350

351+
/**
352+
* Represents encoding options for converting to JSON objects that can be convertible to DataFrame
353+
*
354+
* @param rowsLimit Optional limit on the number of rows to be included in the JSON output.
355+
* Default is null, meaning no limit is imposed.
356+
*/
357+
public class DataframeConvertableEncodingOptions(public val rowsLimit: Int? = null) : EncodingOptions
358+
344359
public fun AnyRow.toJson(prettyPrint: Boolean = false): String {
345360
val json = Json {
346361
this.prettyPrint = prettyPrint

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterHtmlRenderer.kt

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@ import org.jetbrains.kotlinx.dataframe.impl.io.SerializationKeys.NROW
1313
import org.jetbrains.kotlinx.dataframe.impl.io.encodeFrame
1414
import org.jetbrains.kotlinx.dataframe.io.Base64ImageEncodingOptions
1515
import org.jetbrains.kotlinx.dataframe.io.DataFrameHtmlData
16+
import org.jetbrains.kotlinx.dataframe.io.DataframeConvertableEncodingOptions
1617
import org.jetbrains.kotlinx.dataframe.io.DisplayConfiguration
18+
import org.jetbrains.kotlinx.dataframe.io.EncodingOptions
1719
import org.jetbrains.kotlinx.dataframe.io.toHTML
1820
import org.jetbrains.kotlinx.dataframe.io.toJsonWithMetadata
1921
import org.jetbrains.kotlinx.dataframe.io.toStaticHtml
@@ -34,6 +36,7 @@ import org.jetbrains.kotlinx.jupyter.api.renderHtmlAsIFrameIfNeeded
3436
private const val MIN_KERNEL_VERSION_FOR_NEW_TABLES_UI = "0.11.0.311"
3537
private const val MIN_IDE_VERSION_SUPPORT_JSON_WITH_METADATA = 241
3638
private const val MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER = 242
39+
private const val MIN_IDE_VERSION_SUPPORT_DATAFRAME_CONVERTABLE = 243
3740

3841
internal class JupyterHtmlRenderer(val display: DisplayConfiguration, val builder: JupyterIntegration.Builder)
3942

@@ -85,13 +88,19 @@ internal inline fun <reified T : Any> JupyterHtmlRenderer.render(
8588
}
8689

8790
else -> {
88-
val imageEncodingOptions =
89-
if (ideBuildNumber.supportsImageViewer()) Base64ImageEncodingOptions() else null
91+
val encodingOptions = buildList<EncodingOptions> {
92+
if (ideBuildNumber.supportsDataFrameConvertableValues()) {
93+
add(DataframeConvertableEncodingOptions())
94+
}
95+
if (ideBuildNumber.supportsImageViewer()) {
96+
add(Base64ImageEncodingOptions())
97+
}
98+
}
9099

91100
df.toJsonWithMetadata(
92101
rowLimit = limit,
93102
nestedRowLimit = reifiedDisplayConfiguration.rowsLimit,
94-
imageEncodingOptions = imageEncodingOptions,
103+
encodingOptions = encodingOptions,
95104
)
96105
}
97106
}
@@ -108,6 +117,9 @@ private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsDynamicNestedTable
108117
private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsImageViewer() =
109118
this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_IMAGE_VIEWER
110119

120+
private fun KotlinNotebookPluginUtils.IdeBuildNumber?.supportsDataFrameConvertableValues() =
121+
this != null && majorVersion >= MIN_IDE_VERSION_SUPPORT_DATAFRAME_CONVERTABLE
122+
111123
internal fun Notebook.renderAsIFrameAsNeeded(
112124
data: HtmlData,
113125
staticData: HtmlData,

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/ImageSerializationTests.kt

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,11 @@ class ImageSerializationTests(private val encodingOptions: Base64ImageEncodingOp
6262
encodingOptions: Base64ImageEncodingOptions?,
6363
): JsonObject {
6464
val df = dataFrameOf(listOf("imgs"), images)
65-
val jsonStr = df.toJsonWithMetadata(20, nestedRowLimit = 20, imageEncodingOptions = encodingOptions)
65+
val jsonStr = df.toJsonWithMetadata(
66+
20,
67+
nestedRowLimit = 20,
68+
encodingOptions = if (encodingOptions != null) listOf(encodingOptions) else emptyList(),
69+
)
6670

6771
return parseJsonStr(jsonStr)
6872
}

0 commit comments

Comments
 (0)