Skip to content

Commit ba4d321

Browse files
JolanrensenJolanrensen
authored
Merge pull request #401 from Kotlin/jupyter-any-detection
Fix: Jupyter compile-time DF type not recognized
2 parents f1b245e + 7c3f204 commit ba4d321

File tree

8 files changed

+350
-143
lines changed

8 files changed

+350
-143
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/ColumnsSelectionDsl.kt

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1786,12 +1786,10 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
17861786
*
17871787
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { `[colsOf][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.colsOf]`<`[String][String]`>().`[cols][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`() }`
17881788
*
1789-
* @see [all]
1790-
*
17911789
*
17921790
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
17931791
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
1794-
*/
1792+
* @see [all] */
17951793
@Suppress("UNCHECKED_CAST")
17961794
public fun <C> ColumnSet<C>.cols(
17971795
predicate: ColumnFilter<C> = { true },
@@ -1829,12 +1827,10 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
18291827
*
18301828
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { `[colsOf][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.colsOf]`<`[String][String]`>().`[cols][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`() }`
18311829
*
1832-
* @see [all]
1833-
*
18341830
*
18351831
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
18361832
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
1837-
*/
1833+
* @see [all] */
18381834
public operator fun <C> ColumnSet<C>.get(
18391835
predicate: ColumnFilter<C> = { true },
18401836
): TransformableColumnSet<C> = cols(predicate)
@@ -1928,12 +1924,10 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
19281924
*
19291925
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { myColumnGroup`[`[`][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`{ ... }`[`]`][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]` }`
19301926
*
1931-
* @see [all]
1932-
*
19331927
*
19341928
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
19351929
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
1936-
*/
1930+
* @see [all] */
19371931
public fun SingleColumn<*>.cols(
19381932
predicate: ColumnFilter<*> = { true },
19391933
): TransformableColumnSet<*> = colsInternal(predicate)
@@ -1979,12 +1973,11 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
19791973
*
19801974
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { myColumnGroup`[`[`][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`{ ... }`[`]`][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]` }`
19811975
*
1982-
* @see [all]
1983-
*
1984-
*
19851976
*
19861977
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
19871978
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
1979+
* @see [all]
1980+
*
19881981
*/
19891982
public operator fun SingleColumn<*>.get(
19901983
predicate: ColumnFilter<*> = { true },
@@ -2172,12 +2165,10 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
21722165
*
21732166
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { Type::columnGroup.`[cols][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`() }`
21742167
*
2175-
* @see [all]
2176-
*
21772168
*
21782169
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
21792170
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
2180-
*/
2171+
* @see [all] */
21812172
public fun KProperty<*>.cols(
21822173
predicate: ColumnFilter<*> = { true },
21832174
): TransformableColumnSet<*> = colGroup(this).cols(predicate)
@@ -2212,12 +2203,10 @@ public interface ColumnsSelectionDsl<out T> : ColumnSelectionDsl<T>, SingleColum
22122203
*
22132204
* `df.`[select][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.select]` { Type::columnGroup.`[cols][org.jetbrains.kotlinx.dataframe.api.ColumnsSelectionDsl.cols]`() }`
22142205
*
2215-
* @see [all]
2216-
*
22172206
*
22182207
* @param [predicate] A [ColumnFilter function][org.jetbrains.kotlinx.dataframe.ColumnFilter] that takes a [ColumnReference][org.jetbrains.kotlinx.dataframe.columns.ColumnReference] and returns a [Boolean].
22192208
* @return A [ColumnSet][org.jetbrains.kotlinx.dataframe.columns.ColumnSet] containing the columns that match the given [predicate].
2220-
*/
2209+
* @see [all] */
22212210
public operator fun KProperty<*>.get(
22222211
predicate: ColumnFilter<*> = { true },
22232212
): TransformableColumnSet<Any?> = cols(predicate)

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

Lines changed: 129 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,33 @@ import org.jetbrains.dataframe.impl.codeGen.ReplCodeGenerator
55
import org.jetbrains.kotlinx.dataframe.AnyCol
66
import org.jetbrains.kotlinx.dataframe.AnyFrame
77
import org.jetbrains.kotlinx.dataframe.AnyRow
8+
import org.jetbrains.kotlinx.dataframe.DataColumn
9+
import org.jetbrains.kotlinx.dataframe.DataFrame
10+
import org.jetbrains.kotlinx.dataframe.DataRow
811
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
9-
import org.jetbrains.kotlinx.dataframe.api.*
12+
import org.jetbrains.kotlinx.dataframe.api.Convert
13+
import org.jetbrains.kotlinx.dataframe.api.FormattedFrame
14+
import org.jetbrains.kotlinx.dataframe.api.Gather
15+
import org.jetbrains.kotlinx.dataframe.api.GroupBy
16+
import org.jetbrains.kotlinx.dataframe.api.Merge
17+
import org.jetbrains.kotlinx.dataframe.api.Pivot
18+
import org.jetbrains.kotlinx.dataframe.api.PivotGroupBy
19+
import org.jetbrains.kotlinx.dataframe.api.ReducedGroupBy
20+
import org.jetbrains.kotlinx.dataframe.api.ReducedPivot
21+
import org.jetbrains.kotlinx.dataframe.api.ReducedPivotGroupBy
22+
import org.jetbrains.kotlinx.dataframe.api.Split
23+
import org.jetbrains.kotlinx.dataframe.api.SplitWithTransform
24+
import org.jetbrains.kotlinx.dataframe.api.Update
25+
import org.jetbrains.kotlinx.dataframe.api.asColumnGroup
26+
import org.jetbrains.kotlinx.dataframe.api.asDataFrame
27+
import org.jetbrains.kotlinx.dataframe.api.columnsCount
28+
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
29+
import org.jetbrains.kotlinx.dataframe.api.frames
30+
import org.jetbrains.kotlinx.dataframe.api.into
31+
import org.jetbrains.kotlinx.dataframe.api.isColumnGroup
1032
import org.jetbrains.kotlinx.dataframe.api.name
33+
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
34+
import org.jetbrains.kotlinx.dataframe.api.values
1135
import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
1236
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
1337
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
@@ -31,6 +55,7 @@ import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
3155
import org.jetbrains.kotlinx.jupyter.api.libraries.resources
3256
import kotlin.reflect.KClass
3357
import kotlin.reflect.KProperty
58+
import kotlin.reflect.KType
3459
import kotlin.reflect.full.isSubtypeOf
3560

3661
/** Users will get an error if their Kotlin Jupyter kernel is older than this version. */
@@ -45,6 +70,101 @@ internal class Integration(
4570

4671
val version = options["v"]
4772

73+
private fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
74+
val code = codeWithConverter.with(argument)
75+
return if (code.isNotBlank()) {
76+
val result = execute(code)
77+
if (codeWithConverter.hasConverter) {
78+
result.name
79+
} else null
80+
} else null
81+
}
82+
83+
private fun KotlinKernelHost.execute(
84+
codeWithConverter: CodeWithConverter,
85+
property: KProperty<*>,
86+
type: KType,
87+
): VariableName? {
88+
val variableName = "(${property.name}${if (property.returnType.isMarkedNullable) "!!" else ""} as $type)"
89+
return execute(codeWithConverter, variableName)
90+
}
91+
92+
private fun KotlinKernelHost.updateImportDataSchemaVariable(
93+
importDataSchema: ImportDataSchema,
94+
property: KProperty<*>,
95+
): VariableName? {
96+
val formats = supportedFormats.filterIsInstance<SupportedCodeGenerationFormat>()
97+
val name = property.name + "DataSchema"
98+
return when (
99+
val codeGenResult = CodeGenerator.urlCodeGenReader(importDataSchema.url, name, formats, true)
100+
) {
101+
is CodeGenerationReadResult.Success -> {
102+
val readDfMethod = codeGenResult.getReadDfMethod(importDataSchema.url.toExternalForm())
103+
val code = readDfMethod.additionalImports.joinToString("\n") +
104+
"\n" +
105+
codeGenResult.code
106+
107+
execute(code)
108+
execute("""DISPLAY("Data schema successfully imported as ${property.name}: $name")""")
109+
110+
name
111+
}
112+
113+
is CodeGenerationReadResult.Error -> {
114+
execute("""DISPLAY("Failed to read data schema from ${importDataSchema.url}: ${codeGenResult.reason}")""")
115+
null
116+
}
117+
}
118+
}
119+
120+
private fun KotlinKernelHost.updateAnyFrameVariable(
121+
df: AnyFrame,
122+
property: KProperty<*>,
123+
codeGen: ReplCodeGenerator,
124+
): VariableName? = execute(
125+
codeWithConverter = codeGen.process(df, property),
126+
property = property,
127+
type = DataFrame::class.createStarProjectedType(false),
128+
129+
)
130+
131+
private fun KotlinKernelHost.updateAnyRowVariable(
132+
row: AnyRow,
133+
property: KProperty<*>,
134+
codeGen: ReplCodeGenerator,
135+
): VariableName? = execute(
136+
codeWithConverter = codeGen.process(row, property),
137+
property = property,
138+
type = DataRow::class.createStarProjectedType(false),
139+
)
140+
141+
private fun KotlinKernelHost.updateColumnGroupVariable(
142+
col: ColumnGroup<*>,
143+
property: KProperty<*>,
144+
codeGen: ReplCodeGenerator,
145+
): VariableName? = execute(
146+
codeWithConverter = codeGen.process(col.asDataFrame(), property),
147+
property = property,
148+
type = ColumnGroup::class.createStarProjectedType(false),
149+
)
150+
151+
private fun KotlinKernelHost.updateAnyColVariable(
152+
col: AnyCol,
153+
property: KProperty<*>,
154+
codeGen: ReplCodeGenerator,
155+
): VariableName? = if (col.isColumnGroup()) {
156+
val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property).let { c ->
157+
CodeWithConverter(c.declarations) { c.converter("$it.asColumnGroup()") }
158+
}
159+
execute(
160+
codeWithConverter = codeWithConverter,
161+
property = property,
162+
type = DataColumn::class.createStarProjectedType(false),
163+
)
164+
} else {
165+
null
166+
}
167+
48168
override fun Builder.onLoaded() {
49169
if (version != null) {
50170
dependencies(
@@ -152,65 +272,17 @@ internal class Integration(
152272
import("org.jetbrains.kotlinx.dataframe.dataTypes.*")
153273
import("org.jetbrains.kotlinx.dataframe.impl.codeGen.urlCodeGenReader")
154274

155-
fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
156-
val code = codeWithConverter.with(argument)
157-
return if (code.isNotBlank()) {
158-
val result = execute(code)
159-
if (codeWithConverter.hasConverter) {
160-
result.name
161-
} else null
162-
} else null
163-
}
164-
165-
fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, property: KProperty<*>): VariableName? {
166-
val variableName = property.name + if (property.returnType.isMarkedNullable) "!!" else ""
167-
return execute(codeWithConverter, variableName)
168-
}
169-
170-
updateVariable<ImportDataSchema> { importDataSchema, property ->
171-
val formats = supportedFormats.filterIsInstance<SupportedCodeGenerationFormat>()
172-
val name = property.name + "DataSchema"
173-
when (val codeGenResult = CodeGenerator.urlCodeGenReader(importDataSchema.url, name, formats, true)) {
174-
is CodeGenerationReadResult.Success -> {
175-
val readDfMethod = codeGenResult.getReadDfMethod(importDataSchema.url.toExternalForm())
176-
val code = readDfMethod.additionalImports.joinToString("\n") +
177-
"\n" +
178-
codeGenResult.code
179-
180-
execute(code)
181-
execute("""DISPLAY("Data schema successfully imported as ${property.name}: $name")""")
182-
183-
name
184-
}
185-
186-
is CodeGenerationReadResult.Error -> {
187-
execute("""DISPLAY("Failed to read data schema from ${importDataSchema.url}: ${codeGenResult.reason}")""")
188-
null
189-
}
275+
updateVariable<Any> { instance, property ->
276+
when (instance) {
277+
is AnyCol -> updateAnyColVariable(instance, property, codeGen)
278+
is ColumnGroup<*> -> updateColumnGroupVariable(instance, property, codeGen)
279+
is AnyRow -> updateAnyRowVariable(instance, property, codeGen)
280+
is AnyFrame -> updateAnyFrameVariable(instance, property, codeGen)
281+
is ImportDataSchema -> updateImportDataSchemaVariable(instance, property)
282+
else -> null
190283
}
191284
}
192285

193-
updateVariable<AnyFrame> { df, property ->
194-
execute(codeGen.process(df, property), property)
195-
}
196-
197-
updateVariable<AnyRow> { row, property ->
198-
execute(codeGen.process(row, property), property)
199-
}
200-
201-
updateVariable<ColumnGroup<*>> { col, property ->
202-
execute(codeGen.process(col.asDataFrame(), property), property)
203-
}
204-
205-
updateVariable<AnyCol> { col, property ->
206-
if (col.isColumnGroup()) {
207-
val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property).let { c ->
208-
CodeWithConverter(c.declarations) { c.converter("$it.asColumnGroup()") }
209-
}
210-
execute(codeWithConverter, property)
211-
} else null
212-
}
213-
214286
fun KotlinKernelHost.addDataSchemas(classes: List<KClass<*>>) {
215287
val code = classes.joinToString("\n") {
216288
codeGen.process(it)

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/CodeGenerationTests.kt

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.jupyter
22

3+
import org.intellij.lang.annotations.Language
34
import org.jetbrains.kotlinx.jupyter.api.Code
45
import org.junit.Test
56

@@ -11,9 +12,20 @@ class CodeGenerationTests : DataFrameJupyterTest() {
1112
}
1213
}
1314

15+
@Test
16+
fun `Type erased dataframe`() {
17+
@Language("kts")
18+
val a = """
19+
fun create(): Any? = dataFrameOf("a")(1)
20+
val df = create()
21+
df.a
22+
""".checkCompilation()
23+
}
24+
1425
@Test
1526
fun `nullable dataframe`() {
16-
"""
27+
@Language("kts")
28+
val a = """
1729
fun create(): AnyFrame? = dataFrameOf("a")(1)
1830
val df = create()
1931
df.a
@@ -22,7 +34,8 @@ class CodeGenerationTests : DataFrameJupyterTest() {
2234

2335
@Test
2436
fun `nullable columnGroup`() {
25-
"""
37+
@Language("kts")
38+
val a = """
2639
fun create(): AnyCol? = dataFrameOf("a")(1).asColumnGroup().asDataColumn()
2740
val col = create()
2841
col.a
@@ -31,7 +44,8 @@ class CodeGenerationTests : DataFrameJupyterTest() {
3144

3245
@Test
3346
fun `nullable dataRow`() {
34-
"""
47+
@Language("kts")
48+
val a = """
3549
fun create(): AnyRow? = dataFrameOf("a")(1).single()
3650
val row = create()
3751
row.a

0 commit comments

Comments
 (0)