Merge pull request #735 from Kotlin/pivot-fix

Jolanrensen · web-flow · commit 45ba9b5ecfde · 2024-06-14T16:00:38.000+02:00
Pivot fix
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
@@ -83,7 +83,11 @@ internal fun <T, G, R> 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
@@ -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 <T> concatImpl(name: String, columns: List<DataColumn<T>?>, 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 }
             }
         }
@@ -63,7 +70,13 @@ internal fun <T> concatImpl(name: String, columns: List<DataColumn<T>?>, 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
@@ -100,17 +100,21 @@ internal fun <T, R> 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
@@ -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<Int>(), typeOf<Boolean>()).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
@@ -83,7 +83,11 @@ internal fun <T, G, R> 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
@@ -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 <T> concatImpl(name: String, columns: List<DataColumn<T>?>, 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 }
             }
         }
@@ -63,7 +70,13 @@ internal fun <T> concatImpl(name: String, columns: List<DataColumn<T>?>, 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
@@ -100,17 +100,21 @@ internal fun <T, R> 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/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<Int>(), typeOf<Boolean>()).commonType()
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,11 @@ internal fun <T, G, R> aggregateGroupBy(`
`83`	`83`	`val result = body(builder, builder)`
`84`	`84`	`if (result != Unit && result !is NamedValue && result !is AggregatedPivot<*>) builder.yield(`
`85`	`85`	`NamedValue.create(`
`86`		`- pathOf(defaultAggregateName), result, null, null, true`
	`86`	`+ path = pathOf(defaultAggregateName),`
	`87`	`+ value = result,`
	`88`	`+ type = null,`
	`89`	`+ defaultValue = null,`
	`90`	`+ guessType = true,`
`87`	`91`	`)`
`88`	`92`	`)`
`89`	`93`	`builder.compute()`