From 9de64626105a060369fbd2b607f9de9cc7a389cf Mon Sep 17 00:00:00 2001 From: Jolan Rensen Date: Wed, 12 Jun 2024 13:09:02 +0200 Subject: [PATCH 1/3] adds breaking pivot test --- .../jetbrains/kotlinx/dataframe/api/pivot.kt | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt index ba9025b218..5792ff2067 100644 --- a/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt +++ b/core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt @@ -1,6 +1,7 @@ package org.jetbrains.kotlinx.dataframe.api import io.kotest.matchers.shouldBe +import org.jetbrains.kotlinx.dataframe.impl.commonType import org.junit.Test import kotlin.reflect.typeOf @@ -168,4 +169,20 @@ class PivotTests { } df1 shouldBe df.pivot { "category2" then "category3" }.groupBy("category1").count() } + + @Test + fun `pivot with default of other type`() { + val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")( + "Alice", "Cooper", 15, "London", 54, true, + "Bob", "Dylan", 45, "Dubai", 87, true, + "Charlie", "Daniels", 20, "Moscow", null, false, + "Charlie", "Chaplin", 40, "Milan", null, true, + "Bob", "Marley", 30, "Tokyo", 68, true, + "Alice", "Wolf", 20, null, 55, false, + "Charlie", "Byrd", 30, "Moscow", 90, true + ).group("firstName", "lastName").into("name") + + val pivoted = df.pivot("city").groupBy("name").default(0).min() + pivoted["city"]["London"]["isHappy"].type() shouldBe listOf(typeOf(), typeOf()).commonType() + } } From c269e08049b2dc0563093303a136ff9d1123a85e Mon Sep 17 00:00:00 2001 From: Jolan Rensen Date: Wed, 12 Jun 2024 13:10:11 +0200 Subject: [PATCH 2/3] fixes concatImpl to add defaultValue to types when it's used. --- .../org/jetbrains/kotlinx/dataframe/impl/api/concat.kt | 9 ++++++++- .../org/jetbrains/kotlinx/dataframe/impl/api/concat.kt | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt index 591cbe6033..c5879c1ba8 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt @@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.hasNulls import org.jetbrains.kotlinx.dataframe.impl.columns.guessColumnType import org.jetbrains.kotlinx.dataframe.impl.commonType import org.jetbrains.kotlinx.dataframe.impl.getListType +import org.jetbrains.kotlinx.dataframe.impl.guessValueType import org.jetbrains.kotlinx.dataframe.nrow import kotlin.reflect.KType import kotlin.reflect.full.withNullability @@ -54,7 +55,13 @@ internal fun concatImpl(name: String, columns: List?>, columnS col.toList() } else { val nrow = columnSizes[index] - if (!nulls && nrow > 0 && defaultValue == null) nulls = true + if (!nulls && nrow > 0 && defaultValue == null) { + nulls = true + } else if (defaultValue != null) { + types.add( + guessValueType(sequenceOf(defaultValue)) + ) + } List(nrow) { defaultValue } } } diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt index 591cbe6033..c5879c1ba8 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt @@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.hasNulls import org.jetbrains.kotlinx.dataframe.impl.columns.guessColumnType import org.jetbrains.kotlinx.dataframe.impl.commonType import org.jetbrains.kotlinx.dataframe.impl.getListType +import org.jetbrains.kotlinx.dataframe.impl.guessValueType import org.jetbrains.kotlinx.dataframe.nrow import kotlin.reflect.KType import kotlin.reflect.full.withNullability @@ -54,7 +55,13 @@ internal fun concatImpl(name: String, columns: List?>, columnS col.toList() } else { val nrow = columnSizes[index] - if (!nulls && nrow > 0 && defaultValue == null) nulls = true + if (!nulls && nrow > 0 && defaultValue == null) { + nulls = true + } else if (defaultValue != null) { + types.add( + guessValueType(sequenceOf(defaultValue)) + ) + } List(nrow) { defaultValue } } } From df4c9970ff45cad6a69b9bfd4a63ec3611f0c435 Mon Sep 17 00:00:00 2001 From: Jolan Rensen Date: Wed, 12 Jun 2024 13:11:27 +0200 Subject: [PATCH 3/3] small linting I did during debugging to make the code more understandable --- .../kotlinx/dataframe/impl/GroupByImpl.kt | 6 +++++- .../kotlinx/dataframe/impl/api/concat.kt | 8 +++++++- .../kotlinx/dataframe/impl/api/pivot.kt | 16 ++++++++++------ .../jetbrains/kotlinx/dataframe/api/pivot.kt | 17 +++++++++++++++++ .../kotlinx/dataframe/impl/GroupByImpl.kt | 6 +++++- .../kotlinx/dataframe/impl/api/concat.kt | 8 +++++++- .../kotlinx/dataframe/impl/api/pivot.kt | 16 ++++++++++------ 7 files changed, 61 insertions(+), 16 deletions(-) diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt index bf58b3a662..59747c0c79 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt @@ -87,7 +87,11 @@ internal fun aggregateGroupBy( val result = body(builder, builder) if (result != Unit && result !is NamedValue && result !is AggregatedPivot<*>) builder.yield( NamedValue.create( - pathOf(defaultAggregateName), result, null, null, true + path = pathOf(defaultAggregateName), + value = result, + type = null, + defaultValue = null, + guessType = true, ) ) builder.compute() diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt index c5879c1ba8..41b4d0cf74 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt @@ -70,7 +70,13 @@ internal fun concatImpl(name: String, columns: List?>, columnS val baseType = types.commonType() val tartypeOf = if (guessType || !hasList) baseType.withNullability(nulls) else getListType(baseType.withNullability(listOfNullable)) - return guessColumnType(name, list, tartypeOf, guessType, defaultValue).cast() + return guessColumnType( + name = name, + values = list, + suggestedType = tartypeOf, + suggestedTypeIsUpperBound = guessType, + defaultValue = defaultValue, + ).cast() } } diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt index 209b8f2564..f9edd61622 100644 --- a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt +++ b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt @@ -100,17 +100,21 @@ internal fun aggregatePivot( val hasResult = result != null && result != Unit fun NamedValue.apply(path: ColumnPath) = - copy(path = path, value = this.value ?: default ?: globalDefault, default = default ?: globalDefault) + copy( + path = path, + value = this.value ?: default ?: globalDefault, + default = default ?: globalDefault, + ) val values = builder.values when { values.size == 1 && values[0].path.isEmpty() -> aggregator.yield(values[0].apply(path)) values.isEmpty() -> aggregator.yield( - path, - if (hasResult) result else globalDefault, - null, - globalDefault, - true + path = path, + value = if (hasResult) result else globalDefault, + type = null, + default = globalDefault, + guessType = true, ) else -> { diff --git a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt index ba9025b218..5792ff2067 100644 --- a/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt +++ b/core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt @@ -1,6 +1,7 @@ package org.jetbrains.kotlinx.dataframe.api import io.kotest.matchers.shouldBe +import org.jetbrains.kotlinx.dataframe.impl.commonType import org.junit.Test import kotlin.reflect.typeOf @@ -168,4 +169,20 @@ class PivotTests { } df1 shouldBe df.pivot { "category2" then "category3" }.groupBy("category1").count() } + + @Test + fun `pivot with default of other type`() { + val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")( + "Alice", "Cooper", 15, "London", 54, true, + "Bob", "Dylan", 45, "Dubai", 87, true, + "Charlie", "Daniels", 20, "Moscow", null, false, + "Charlie", "Chaplin", 40, "Milan", null, true, + "Bob", "Marley", 30, "Tokyo", 68, true, + "Alice", "Wolf", 20, null, 55, false, + "Charlie", "Byrd", 30, "Moscow", 90, true + ).group("firstName", "lastName").into("name") + + val pivoted = df.pivot("city").groupBy("name").default(0).min() + pivoted["city"]["London"]["isHappy"].type() shouldBe listOf(typeOf(), typeOf()).commonType() + } } diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt index bf58b3a662..59747c0c79 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt @@ -87,7 +87,11 @@ internal fun aggregateGroupBy( val result = body(builder, builder) if (result != Unit && result !is NamedValue && result !is AggregatedPivot<*>) builder.yield( NamedValue.create( - pathOf(defaultAggregateName), result, null, null, true + path = pathOf(defaultAggregateName), + value = result, + type = null, + defaultValue = null, + guessType = true, ) ) builder.compute() diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt index c5879c1ba8..41b4d0cf74 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt @@ -70,7 +70,13 @@ internal fun concatImpl(name: String, columns: List?>, columnS val baseType = types.commonType() val tartypeOf = if (guessType || !hasList) baseType.withNullability(nulls) else getListType(baseType.withNullability(listOfNullable)) - return guessColumnType(name, list, tartypeOf, guessType, defaultValue).cast() + return guessColumnType( + name = name, + values = list, + suggestedType = tartypeOf, + suggestedTypeIsUpperBound = guessType, + defaultValue = defaultValue, + ).cast() } } diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt index 209b8f2564..f9edd61622 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt @@ -100,17 +100,21 @@ internal fun aggregatePivot( val hasResult = result != null && result != Unit fun NamedValue.apply(path: ColumnPath) = - copy(path = path, value = this.value ?: default ?: globalDefault, default = default ?: globalDefault) + copy( + path = path, + value = this.value ?: default ?: globalDefault, + default = default ?: globalDefault, + ) val values = builder.values when { values.size == 1 && values[0].path.isEmpty() -> aggregator.yield(values[0].apply(path)) values.isEmpty() -> aggregator.yield( - path, - if (hasResult) result else globalDefault, - null, - globalDefault, - true + path = path, + value = if (hasResult) result else globalDefault, + type = null, + default = globalDefault, + guessType = true, ) else -> {