Skip to content

Pivot fix #735

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jun 14, 2024
Merged

Pivot fix #735

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ba43d21
Merge branch 'refs/heads/is-comparable-fix' into pivot-fix
Jolanrensen Jun 6, 2024
9e372d3
Merge branch 'master' into pivot-fix
Jolanrensen Jun 12, 2024
9de6462
adds breaking pivot test
Jolanrensen Jun 12, 2024
c269e08
fixes concatImpl to add defaultValue to types when it's used.
Jolanrensen Jun 12, 2024
df4c997
small linting I did during debugging to make the code more understand…
Jolanrensen Jun 12, 2024
e1668f8
Merge branch 'master' into pivot-fix
Jolanrensen Jun 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
17 changes: 15 additions & 2 deletions core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 }
}
}
Expand All @@ -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()
}
}

Expand Down
16 changes: 10 additions & 6 deletions core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 -> {
Expand Down
17 changes: 17 additions & 0 deletions core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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()
}
}
6 changes: 5 additions & 1 deletion core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/GroupByImpl.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
17 changes: 15 additions & 2 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/concat.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 }
}
}
Expand All @@ -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()
}
}

Expand Down
16 changes: 10 additions & 6 deletions core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/pivot.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 -> {
Expand Down
17 changes: 17 additions & 0 deletions core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/pivot.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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()
}
}