Skip to content

Commit 25762f8

Browse files
koperagenkoperagen
committed
Clarify that schema of property with DataFrame<T>? type is not a FrameColumn
This makes extracted schema represent runtime more accurately
1 parent cf60da2 commit 25762f8

5 files changed

Lines changed: 105 additions & 47 deletions

File tree

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import kotlin.reflect.typeOf
2222
internal fun KType.getFieldKind(): FieldKind =
2323
FieldKind.of(
2424
this,
25-
isDataFrame = { jvmErasure == DataFrame::class },
25+
isDataFrame = { jvmErasure == DataFrame::class && !isMarkedNullable },
2626
isListToFrame = {
2727
jvmErasure == List::class && (arguments[0].type?.jvmErasure?.hasAnnotation<DataSchema>() == true)
2828
},

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/convertTo.kt‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,14 @@ internal fun AnyFrame.convertToImpl(
193193

194194
when (targetSchema.kind) {
195195
ColumnKind.Value ->
196-
convertedColumn ?: originalColumn.convertTo(to)
196+
when {
197+
convertedColumn != null -> convertedColumn
198+
199+
originalColumn.kind == ColumnKind.Frame && to.jvmErasure == DataFrame::class ->
200+
originalColumn
201+
202+
else -> originalColumn.convertTo(to)
203+
}
197204

198205
ColumnKind.Group -> {
199206
val column = when {

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -320,18 +320,27 @@ internal fun convertToDataFrame(
320320
val kClass = returnType.classifier as KClass<*>
321321
val fieldKind = returnType.getFieldKind()
322322

323+
// property type might not be of FrameColumn kind if `AnyFrame?`, but we still have to narrow it to FrameCol
324+
// when no actual nulls are met
325+
val shouldCreateFrameCol = kClass == DataFrame::class && !nullable
326+
323327
val keepSubtree =
324-
maxDepth <= 0 && !fieldKind.shouldBeConvertedToFrameColumn && !fieldKind.shouldBeConvertedToColumnGroup
325-
val shouldCreateValueCol = keepSubtree ||
326-
kClass in preserveClasses ||
327-
property in preserveProperties ||
328-
(
329-
!kClass.canBeUnfolded &&
330-
!fieldKind.shouldBeConvertedToFrameColumn &&
331-
!fieldKind.shouldBeConvertedToColumnGroup
332-
)
328+
maxDepth <= 0 &&
329+
!fieldKind.shouldBeConvertedToFrameColumn &&
330+
!fieldKind.shouldBeConvertedToColumnGroup &&
331+
!shouldCreateFrameCol
332+
333+
val shouldCreateValueCol =
334+
keepSubtree ||
335+
kClass in preserveClasses ||
336+
property in preserveProperties ||
337+
(
338+
!kClass.canBeUnfolded &&
339+
!fieldKind.shouldBeConvertedToFrameColumn &&
340+
!fieldKind.shouldBeConvertedToColumnGroup &&
341+
!shouldCreateFrameCol
342+
)
333343

334-
val shouldCreateFrameCol = kClass == DataFrame::class && !nullable
335344
val shouldCreateColumnGroup = kClass == DataRow::class
336345

337346
if (shouldCreateFrameCol && shouldCreateValueCol) {
@@ -358,7 +367,7 @@ internal fun convertToDataFrame(
358367
shouldCreateColumnGroup ->
359368
DataColumn.createColumnGroup(
360369
name = it.columnName,
361-
df = (values as List<AnyRow>).concat(),
370+
df = (values as List<AnyRow?>).concat(),
362371
)
363372

364373
kClass.isSubclassOf(Iterable::class) ->

‎core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/convertTo.kt‎

Lines changed: 59 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,21 @@ class ConvertToTests {
184184
@DataSchema
185185
data class DataSchemaWithAnyFrame(val dfs: AnyFrame?)
186186

187+
private fun locationsFrame(): DataFrame<Location?> =
188+
listOf(
189+
Location("Home", Gps(0.0, 0.0)),
190+
Location("Away", null),
191+
null,
192+
).toDataFrame()
193+
.alsoDebug("locations:")
194+
195+
private fun gpsFrame(): DataFrame<Gps?> =
196+
listOf(
197+
Gps(0.0, 0.0),
198+
null,
199+
).toDataFrame()
200+
.alsoDebug("gps:")
201+
187202
@Test
188203
fun test() {
189204
val df1 = dataFrameOf("a")(1, 2, 3)
@@ -202,67 +217,70 @@ class ConvertToTests {
202217
}
203218

204219
@Test
205-
fun `convert df with AnyFrame to itself`() {
206-
val locationsList = listOf(
207-
Location("Home", Gps(0.0, 0.0)),
208-
Location("Away", null),
209-
null,
210-
)
211-
val locations = locationsList
212-
.toDataFrame()
213-
.alsoDebug("locations:")
220+
fun `convert df with AnyFrame containing locations to itself`() {
221+
val locations = locationsFrame()
214222

215-
val gpsList = listOf(
216-
Gps(0.0, 0.0),
217-
null,
218-
)
219-
val gps = gpsList
223+
listOf(DataSchemaWithAnyFrame(locations))
220224
.toDataFrame()
221-
.alsoDebug("gps:")
222-
223-
val df1 = listOf(
224-
DataSchemaWithAnyFrame(locations),
225-
).toDataFrame()
226225
.alsoDebug("df1:")
226+
.convertTo<DataSchemaWithAnyFrame>()
227+
}
227228

228-
df1.convertTo<DataSchemaWithAnyFrame>()
229+
@Test
230+
fun `convert df with AnyFrame containing gps to itself`() {
231+
val gps = gpsFrame()
229232

230-
val df2 = listOf(
231-
DataSchemaWithAnyFrame(gps),
232-
).toDataFrame()
233+
listOf(DataSchemaWithAnyFrame(gps))
234+
.toDataFrame()
233235
.alsoDebug("df2:")
236+
.convertTo<DataSchemaWithAnyFrame>()
237+
}
234238

235-
df2.convertTo<DataSchemaWithAnyFrame>()
239+
@Test
240+
fun `convert df with preserved AnyFrame containing null and gps to itself`() {
241+
val gps = gpsFrame()
236242

237-
val df3 = listOf(
243+
listOf(
238244
DataSchemaWithAnyFrame(null),
239245
DataSchemaWithAnyFrame(gps),
240246
).toDataFrame { properties { preserve(DataFrame::class) } }
241247
.alsoDebug("df3 before convert:")
248+
.convertTo<DataSchemaWithAnyFrame>()
249+
}
242250

243-
df3.convertTo<DataSchemaWithAnyFrame>()
244-
245-
val df4 = listOf(
251+
@Test
252+
fun `convert df with preserved null AnyFrame to itself`() {
253+
listOf(
246254
DataSchemaWithAnyFrame(null),
247255
).toDataFrame { properties { preserve(DataFrame::class) } }
248256
.alsoDebug("df4 before convert:")
257+
.convertTo<DataSchemaWithAnyFrame>()
258+
}
249259

250-
df4.convertTo<DataSchemaWithAnyFrame>()
260+
@Test
261+
fun `convert raw df with AnyFrame column to itself`() {
262+
val locations = locationsFrame()
263+
val gps = gpsFrame()
251264

252-
val df5a: DataFrame<*> = dataFrameOf(
265+
val df: DataFrame<*> = dataFrameOf(
253266
columnOf(locations, gps, null).named("dfs"),
254267
).alsoDebug("df5a:")
255268

256-
df5a.convertTo<DataSchemaWithAnyFrame>()
269+
df.convertTo<DataSchemaWithAnyFrame>()
270+
}
257271

258-
val df5 = listOf(
272+
@Test
273+
fun `convert df with preserved mixed AnyFrame values to itself repeatedly`() {
274+
val locations = locationsFrame()
275+
val gps = gpsFrame()
276+
277+
listOf(
259278
DataSchemaWithAnyFrame(null),
260279
DataSchemaWithAnyFrame(locations),
261280
DataSchemaWithAnyFrame(gps),
262281
).toDataFrame { properties { preserve(DataFrame::class) } }
263282
.alsoDebug("df5 before convert:")
264-
265-
df5.convertTo<DataSchemaWithAnyFrame>()
283+
.convertTo<DataSchemaWithAnyFrame>()
266284
.alsoDebug("df5 after convert:")
267285
.convertTo<DataSchemaWithAnyFrame>()
268286
.alsoDebug("df5 after second convert:")
@@ -394,4 +412,11 @@ class ConvertToTests {
394412
converted["a"].type() shouldBe typeOf<SimpleEnum?>()
395413
converted shouldBe dataFrameOf("a")(SimpleEnum.A, SimpleEnum.B, null)
396414
}
415+
416+
// @Test
417+
// fun fff() {
418+
// dataFrameOf("a" to columnOf(1, 2, 3)).convertTo<NullableA>()
419+
// }
420+
//
421+
// data class NullableA(val a: Int?)
397422
}

‎core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ import io.kotest.matchers.shouldBe
44
import org.jetbrains.kotlinx.dataframe.AnyRow
55
import org.jetbrains.kotlinx.dataframe.ColumnsScope
66
import org.jetbrains.kotlinx.dataframe.DataColumn
7+
import org.jetbrains.kotlinx.dataframe.DataFrame
78
import org.jetbrains.kotlinx.dataframe.DataRow
9+
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
810
import org.jetbrains.kotlinx.dataframe.api.columnOf
911
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
1012
import org.jetbrains.kotlinx.dataframe.api.default
@@ -16,16 +18,19 @@ import org.jetbrains.kotlinx.dataframe.api.groupBy
1618
import org.jetbrains.kotlinx.dataframe.api.into
1719
import org.jetbrains.kotlinx.dataframe.api.move
1820
import org.jetbrains.kotlinx.dataframe.api.pathOf
21+
import org.jetbrains.kotlinx.dataframe.api.print
1922
import org.jetbrains.kotlinx.dataframe.api.schema
2023
import org.jetbrains.kotlinx.dataframe.api.toCodeString
2124
import org.jetbrains.kotlinx.dataframe.api.under
2225
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
2326
import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGenerator
2427
import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
2528
import org.jetbrains.kotlinx.dataframe.impl.toCamelCaseByDelimiters
29+
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
2630
import org.jetbrains.kotlinx.dataframe.testSets.person.BaseTest
2731
import org.jetbrains.kotlinx.dataframe.testSets.person.Person
2832
import org.junit.Test
33+
import kotlin.reflect.typeOf
2934
import kotlin.test.assertEquals
3035

3136
class CodeGenerationTests : BaseTest() {
@@ -570,6 +575,18 @@ class CodeGenerationTests : BaseTest() {
570575
assertEquals(expected, df.generateDataClasses().value)
571576
}
572577

578+
@DataSchema
579+
class C(val i: Int)
580+
581+
@DataSchema
582+
class Schema(val df: DataFrame<C>?)
583+
584+
@Test
585+
fun extractNullableDataFrameSchema() {
586+
val schema = MarkersExtractor.get<Schema>().schema
587+
schema.columns["df"] shouldBe ColumnSchema.Value(typeOf<DataFrame<C>?>())
588+
}
589+
573590
// region Tests for generateX functions
574591

575592
@Test

0 commit comments

Comments
 (0)