Skip to content

Commit b07b7a1

Browse files
JolanrensenJolanrensen
committed
fixed edge cases unify numbers for json and fixed tests
1 parent ff833c1 commit b07b7a1

File tree

9 files changed

+171
-66
lines changed

9 files changed

+171
-66
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import org.jetbrains.kotlinx.dataframe.impl.zero
2525
import org.jetbrains.kotlinx.dataframe.math.sum
2626
import org.jetbrains.kotlinx.dataframe.math.sumOf
2727
import kotlin.reflect.KProperty
28+
import kotlin.reflect.full.isSubtypeOf
2829
import kotlin.reflect.typeOf
2930

3031
// region DataColumn
@@ -42,7 +43,11 @@ public inline fun <T, reified R : Number> DataColumn<T>.sumOf(crossinline expres
4243

4344
// region DataRow
4445

45-
public fun AnyRow.rowSum(): Number = Aggregators.sum.aggregateMixed(values().filterIsInstance<Number>()) ?: 0
46+
public fun AnyRow.rowSum(): Number =
47+
Aggregators.sum.aggregateMixed(
48+
values = values().filterIsInstance<Number>(),
49+
types = columnTypes().filter { it.isSubtypeOf(typeOf<Number?>()) }.toSet(),
50+
) ?: 0
4651

4752
public inline fun <reified T : Number> AnyRow.rowSumOf(): T = values().filterIsInstance<T>().sum(typeOf<T>())
4853

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/TypeUtils.kt‎

Lines changed: 24 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -167,35 +167,6 @@ internal fun resolve(actualType: KType, declaredType: KType): Map<KTypeParameter
167167
return map
168168
}
169169

170-
internal val numberTypeExtensions: Map<Pair<KClass<*>, KClass<*>>, KClass<*>> by lazy {
171-
val map = mutableMapOf<Pair<KClass<*>, KClass<*>>, KClass<*>>()
172-
173-
fun add(from: KClass<*>, to: KClass<*>) {
174-
map[from to to] = to
175-
map[to to from] = to
176-
}
177-
178-
val intTypes = listOf(Byte::class, Short::class, Int::class, Long::class)
179-
for (i in intTypes.indices) {
180-
for (j in i + 1 until intTypes.size) {
181-
add(intTypes[i], intTypes[j])
182-
}
183-
add(intTypes[i], Double::class)
184-
}
185-
add(Float::class, Double::class)
186-
map
187-
}
188-
189-
internal fun getCommonNumberType(first: KClass<*>?, second: KClass<*>): KClass<*> =
190-
when {
191-
first == null -> second
192-
first == second -> first
193-
else -> numberTypeExtensions[first to second] ?: error("Can not find common number type for $first and $second")
194-
}
195-
196-
internal fun Iterable<KClass<*>>.commonNumberClass(): KClass<*> =
197-
fold(null as KClass<*>?, ::getCommonNumberType) ?: Number::class
198-
199170
internal fun commonParent(classes: Iterable<KClass<*>>): KClass<*>? = commonParents(classes).withMostSuperclasses()
200171

201172
internal fun commonParent(vararg classes: KClass<*>): KClass<*>? = commonParent(classes.toList())
@@ -648,3 +619,27 @@ internal fun Any.asArrayAsListOrNull(): List<*>? =
648619
}
649620

650621
internal fun Any.isBigNumber(): Boolean = this is BigInteger || this is BigDecimal
622+
623+
/**
624+
* Returns a set containing the [KClass] of each element in the iterable.
625+
*
626+
* This can be a heavy operation!
627+
*
628+
* The [KClass] is determined by retrieving the runtime class of each element.
629+
*
630+
* @return A set of [KClass] objects representing the runtime types of elements in the iterable.
631+
*/
632+
internal fun Iterable<Any>.classes(): Set<KClass<*>> = mapTo(mutableSetOf()) { it::class }
633+
634+
/**
635+
* Returns a set of [KType] objects representing the star-projected types of the runtime classes
636+
* of all unique elements in the iterable.
637+
*
638+
* The method internally relies on the [classes] function to collect the runtime classes of the
639+
* elements in the iterable and then maps each class to its star-projected type.
640+
*
641+
* This can be a heavy operation!
642+
*
643+
* @return A set of [KType] objects corresponding to the star-projected runtime types of elements in the iterable.
644+
*/
645+
internal fun Iterable<Any>.types(): Set<KType> = classes().mapTo(mutableSetOf()) { it.createStarProjectedType(false) }

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/aggregation/aggregators/NumbersAggregator.kt‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,36 @@
11
package org.jetbrains.kotlinx.dataframe.impl.aggregation.aggregators
22

33
import org.jetbrains.kotlinx.dataframe.DataColumn
4-
import org.jetbrains.kotlinx.dataframe.impl.commonNumberClass
5-
import org.jetbrains.kotlinx.dataframe.impl.createStarProjectedType
4+
import org.jetbrains.kotlinx.dataframe.impl.convertToUnifiedNumberType
5+
import org.jetbrains.kotlinx.dataframe.impl.unifiedNumberType
66
import kotlin.reflect.KProperty
77
import kotlin.reflect.KType
88

9-
internal class NumbersAggregator<C : Number>(name: String, aggregate: (Iterable<C>, KType) -> C?) :
10-
AggregatorBase<C, C>(name, aggregate) {
9+
internal class NumbersAggregator(name: String, aggregate: (Iterable<Number>, KType) -> Number?) :
10+
AggregatorBase<Number, Number>(name, aggregate) {
1111

12-
override fun aggregate(columns: Iterable<DataColumn<C?>>): C? = aggregateMixed(columns.mapNotNull { aggregate(it) })
12+
override fun aggregate(columns: Iterable<DataColumn<Number?>>): Number? =
13+
aggregateMixed(
14+
values = columns.mapNotNull { aggregate(it) },
15+
types = columns.map { it.type() }.toSet(),
16+
)
1317

1418
class Factory(private val aggregate: Iterable<Number>.(KType) -> Number?) : AggregatorProvider<Number, Number> {
1519
override fun create(name: String) = NumbersAggregator(name, aggregate)
1620

17-
override operator fun getValue(obj: Any?, property: KProperty<*>): NumbersAggregator<Number> =
18-
create(property.name)
21+
override operator fun getValue(obj: Any?, property: KProperty<*>): NumbersAggregator = create(property.name)
1922
}
2023

21-
fun aggregateMixed(values: Iterable<C>): C? {
22-
val classes = values.map { it.javaClass.kotlin }
23-
return aggregate(values, classes.commonNumberClass().createStarProjectedType(false))
24+
/**
25+
* Can aggregate numbers with different types by first converting them to a compatible type.
26+
*/
27+
@Suppress("UNCHECKED_CAST")
28+
fun aggregateMixed(values: Iterable<Number>, types: Set<KType>): Number? {
29+
val commonType = types.unifiedNumberType()
30+
return aggregate(
31+
values = values.convertToUnifiedNumberType(commonType),
32+
type = commonType,
33+
)
2434
}
2535

2636
override val preservesType = false

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/statistics/sum.kt‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@ import io.kotest.matchers.shouldBe
44
import org.jetbrains.kotlinx.dataframe.DataColumn
55
import org.jetbrains.kotlinx.dataframe.api.columnOf
66
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
7+
import org.jetbrains.kotlinx.dataframe.api.rowSum
78
import org.jetbrains.kotlinx.dataframe.api.sum
89
import org.jetbrains.kotlinx.dataframe.api.sumOf
910
import org.junit.Test
11+
import java.math.BigDecimal
1012

1113
class SumTests {
1214

@@ -61,4 +63,24 @@ class SumTests {
6163
df.sum { value2 } shouldBe expected2
6264
df.sum { value3 } shouldBe expected3
6365
}
66+
67+
/** [Issue #1068](https://github.com/Kotlin/dataframe/issues/1068) */
68+
@Test
69+
fun `rowSum mixed number types`() {
70+
dataFrameOf("a", "b")(1, 2f)[0].rowSum().let {
71+
it shouldBe 3.0
72+
it::class shouldBe Double::class
73+
}
74+
75+
// NOTE! unsigned numbers are not Number, they are skipped for now
76+
dataFrameOf("a", "b")(1, 2u)[0].rowSum().let {
77+
it shouldBe 1
78+
it::class shouldBe Int::class
79+
}
80+
81+
dataFrameOf("a", "b")(1.0, 2L)[0].rowSum().let {
82+
it shouldBe (3.0.toBigDecimal())
83+
it::class shouldBe BigDecimal::class
84+
}
85+
}
6486
}

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/types/UtilTests.kt‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,23 @@ import io.kotest.matchers.shouldBe
44
import org.jetbrains.kotlinx.dataframe.DataColumn
55
import org.jetbrains.kotlinx.dataframe.DataRow
66
import org.jetbrains.kotlinx.dataframe.api.columnOf
7+
import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
78
import org.jetbrains.kotlinx.dataframe.impl.asArrayAsListOrNull
89
import org.jetbrains.kotlinx.dataframe.impl.commonParent
910
import org.jetbrains.kotlinx.dataframe.impl.commonParents
1011
import org.jetbrains.kotlinx.dataframe.impl.commonType
1112
import org.jetbrains.kotlinx.dataframe.impl.commonTypeListifyValues
1213
import org.jetbrains.kotlinx.dataframe.impl.createType
14+
import org.jetbrains.kotlinx.dataframe.impl.getUnifiedNumberClass
1315
import org.jetbrains.kotlinx.dataframe.impl.guessValueType
1416
import org.jetbrains.kotlinx.dataframe.impl.isArray
1517
import org.jetbrains.kotlinx.dataframe.impl.isPrimitiveArray
1618
import org.jetbrains.kotlinx.dataframe.impl.nothingType
1719
import org.jetbrains.kotlinx.dataframe.impl.replaceGenericTypeParametersWithUpperbound
1820
import org.junit.Test
1921
import java.io.Serializable
22+
import java.math.BigDecimal
23+
import java.math.BigInteger
2024
import kotlin.reflect.KClass
2125
import kotlin.reflect.KType
2226
import kotlin.reflect.typeOf
@@ -414,4 +418,63 @@ class UtilTests {
414418
typeOf<Set<Nothing>?>(),
415419
).commonTypeListifyValues() shouldBe typeOf<Collection<out Nothing?>?>()
416420
}
421+
422+
/**
423+
* See [UnifyingNumbers] for more information.
424+
* ```
425+
* BigDecimal
426+
* / \
427+
* BigInteger \
428+
* / \ \
429+
* ULong Long Double
430+
* .. | / | / | \..
431+
* \ | / | / |
432+
* UInt Int Float
433+
* .. | / | / \..
434+
* \ | / | /
435+
* UShort Short
436+
* | / |
437+
* | / |
438+
* UByte Byte
439+
* ```
440+
*/
441+
@Test
442+
fun `common number types`() {
443+
// Same type
444+
getUnifiedNumberClass(Int::class, Int::class) shouldBe Int::class
445+
getUnifiedNumberClass(Double::class, Double::class) shouldBe Double::class
446+
447+
// Direct parent-child relationships
448+
getUnifiedNumberClass(Int::class, UShort::class) shouldBe Int::class
449+
getUnifiedNumberClass(Long::class, UInt::class) shouldBe Long::class
450+
getUnifiedNumberClass(Double::class, Float::class) shouldBe Double::class
451+
getUnifiedNumberClass(UShort::class, Short::class) shouldBe Int::class
452+
getUnifiedNumberClass(UByte::class, Byte::class) shouldBe Short::class
453+
454+
getUnifiedNumberClass(UByte::class, UShort::class) shouldBe UShort::class
455+
456+
// Multi-level relationships
457+
getUnifiedNumberClass(Byte::class, Int::class) shouldBe Int::class
458+
getUnifiedNumberClass(UByte::class, Long::class) shouldBe Long::class
459+
getUnifiedNumberClass(Short::class, Double::class) shouldBe Double::class
460+
getUnifiedNumberClass(UInt::class, Int::class) shouldBe Long::class
461+
462+
// Top-level types
463+
getUnifiedNumberClass(BigDecimal::class, Double::class) shouldBe BigDecimal::class
464+
getUnifiedNumberClass(BigInteger::class, Long::class) shouldBe BigInteger::class
465+
getUnifiedNumberClass(BigDecimal::class, BigInteger::class) shouldBe BigDecimal::class
466+
467+
// Distant relationships
468+
getUnifiedNumberClass(Byte::class, BigDecimal::class) shouldBe BigDecimal::class
469+
getUnifiedNumberClass(UByte::class, Double::class) shouldBe Double::class
470+
471+
// Complex type promotions
472+
getUnifiedNumberClass(Int::class, Float::class) shouldBe Double::class
473+
getUnifiedNumberClass(Long::class, Double::class) shouldBe BigDecimal::class
474+
getUnifiedNumberClass(ULong::class, Double::class) shouldBe BigDecimal::class
475+
getUnifiedNumberClass(BigInteger::class, Double::class) shouldBe BigDecimal::class
476+
477+
// Edge case with null
478+
getUnifiedNumberClass(null, Int::class) shouldBe Int::class
479+
}
417480
}

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -472,10 +472,15 @@ internal fun guessValueType(
472472
!nullsInCollection
473473

474474
if (unifyNumbers) {
475-
val usedNumberClasses = classes.filter { it.isSubclassOf(Number::class) }
476-
val unifiedNumberClass = usedNumberClasses.unifiedNumberClass() as KClass<Number>
477-
classes -= usedNumberClasses
478-
classes += unifiedNumberClass
475+
val nothingClass = Nothing::class
476+
val usedNumberClasses = classes.filter {
477+
it.isSubclassOf(Number::class) && it != nothingClass
478+
}
479+
if (usedNumberClasses.isNotEmpty()) {
480+
val unifiedNumberClass = usedNumberClasses.unifiedNumberClass() as KClass<Number>
481+
classes -= usedNumberClasses
482+
classes += unifiedNumberClass
483+
}
479484
}
480485

481486
return when {

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/columns/constructors.kt‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,13 @@ import org.jetbrains.kotlinx.dataframe.columns.ColumnsResolver
3131
import org.jetbrains.kotlinx.dataframe.columns.TypeSuggestion
3232
import org.jetbrains.kotlinx.dataframe.columns.ValueColumn
3333
import org.jetbrains.kotlinx.dataframe.columns.toColumnsSetOf
34+
import org.jetbrains.kotlinx.dataframe.documentation.UnifyingNumbers
3435
import org.jetbrains.kotlinx.dataframe.impl.DataFrameReceiver
3536
import org.jetbrains.kotlinx.dataframe.impl.DataRowImpl
3637
import org.jetbrains.kotlinx.dataframe.impl.api.createConverter
3738
import org.jetbrains.kotlinx.dataframe.impl.asList
3839
import org.jetbrains.kotlinx.dataframe.impl.guessValueType
40+
import org.jetbrains.kotlinx.dataframe.impl.isNothing
3941
import org.jetbrains.kotlinx.dataframe.impl.replaceGenericTypeParametersWithUpperbound
4042
import org.jetbrains.kotlinx.dataframe.index
4143
import org.jetbrains.kotlinx.dataframe.nrow
@@ -197,7 +199,7 @@ internal fun Array<out String>.toNumberColumns() = toColumnsSetOf<Number>()
197199
* @param allColsMakesColGroup if `true`, then, if all values are non-null same-sized columns,
198200
* a column group will be created instead of a [DataColumn][DataColumn]`<`[AnyCol][AnyCol]`>`.
199201
* @param unifyNumbers if `true`, then all numbers encountered in [values] will be converted to the smallest possible
200-
* number-type that can hold all the values lossless. Unsigned numbers are not supported.
202+
* number-type that can hold all the values lossless. Unsigned numbers are not supported. See [UnifyingNumbers].
201203
* For example, if the values are `[1, 2f, 3.0]`, then all values will be converted to [Double].
202204
*/
203205
@PublishedApi
@@ -257,7 +259,7 @@ internal fun <T> createColumnGuessingType(
257259
to = targetType,
258260
) as (Number) -> Number?
259261

260-
return { value -> if (value is Number) converter(value) else value }
262+
return { value -> if (value != null && value is Number) converter(value) else value }
261263
}
262264

263265
return when (type.classifier!! as KClass<*>) {
@@ -292,7 +294,9 @@ internal fun <T> createColumnGuessingType(
292294
var isListOfRows: Boolean? = null
293295
val subType = type.arguments.first().type!! // List<T> -> T
294296

295-
val needsNumberConversion = unifyNumbers && subType.isSubtypeOf(typeOf<Number?>())
297+
val needsNumberConversion = unifyNumbers &&
298+
subType.isSubtypeOf(typeOf<Number?>()) &&
299+
!subType.isNothing
296300
val numberConverter: (Any?) -> Any? by lazy { getSafeNumberConverter(subType) }
297301

298302
val lists = values.map { value ->
@@ -334,7 +338,9 @@ internal fun <T> createColumnGuessingType(
334338
}
335339

336340
else -> {
337-
val needsNumberConversion = unifyNumbers && type.isSubtypeOf(typeOf<Number?>())
341+
val needsNumberConversion = unifyNumbers &&
342+
type.isSubtypeOf(typeOf<Number?>()) &&
343+
!type.isNothing
338344
val numberConverter by lazy { getSafeNumberConverter(type) }
339345

340346
DataColumn.createValueColumn(

0 commit comments

Comments
 (0)