Add Compiler Plugin support for statistics on GroupBy

[zaleslaw]

PR contains:

1. test scenarious for sum*/min/max/std/mean/median family on GroupBy
2. ported test scenarious for Compiler Plugin
3. Changes in Compiler Plugin

Tests for mean/std with results in Double doesn't work well (regarding incorrect resulting types), but please review the logic and keep test muted for a while

Closes #1061 #1074

[zaleslaw]

@koperagen here, I am trying to check if types if comparable to emulate the interComparable() internal function behaviour for median/min/max functions

[Jolanrensen]

You mean this one?

dataframe/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/DataColumnType.kt

Line 74 in d8c6c4b

public fun AnyCol.valuesAreComparable(): Boolean =

It's indeed a tad trickier, because Comparable has a contravariant type argument. They way it's checked at runtime is by first checking whether the type of the column is Comparable, then we get the common Comparable type (this is the tricky part). If the result is Comparable<Nothing> or Comparable<*> we know the values are not comparable to each other.

This works because Comparable<in Pet> + Comparable<in Dog> -> Comparable<in Dog>, but Comparable<in Dog> + Comparable<in Cat> -> Comparable<in Nothing> (aka, you can compare dogs with other pets as long as they're also dogs, but you cannot compare dogs and cats).

So, you'd need to check whether you can do the same with the types at the compiler plugin level, or if you can get a KType, extract the logic to KTypes from :core :).

[koperagen]

valuesAreComparable should be impemented like this, to follow public fun <T : Comparable<T>> Iterable<T>.maxOrNull(): T? signature.

public fun AnyCol.valuesAreComparable(): Boolean {
    val other = Comparable::class.createType(listOf(KTypeProjection(KVariance.IN, type)))
    return type.isSubtypeOf(other)
}

Now it's possible to break it with this code:

val i = 1
    val i1: Comparable<Int> = object : Comparable<Int> {
        override fun compareTo(other: Int): Int {
            return other
        }
    }
    val col by columnOf(i, i1)
    dataFrameOf(col).max().print()

Because indeed, Comparable<Int> does not implement Comparable<Comparable<Int>>
It's not possible to call listOf(i, i1).maxOrNull()

[koperagen]

So for plugin:

fun isInterComparable(col: SimpleDataColumn): Boolean {
      val comparable = StandardClassIds.Comparable.constructClassLikeType(arrayOf(col.type.type))
      return col.type.type.isSubtypeOf(comparable, session)
  }

[Jolanrensen]

@koperagen I created a branch with this suggested fix, but it doesn't produce the expected results either https://github.com/Kotlin/dataframe/tree/values-are-comparable-fix

[Jolanrensen]

@koperagen actually, looking better at it I think you're right and your fix is too (rewritten to this to account for nullability:

public fun AnyCol.valuesAreComparable(): Boolean =
    isSubtypeOf(
        Comparable::class.createType(
            arguments = listOf(KTypeProjection(KVariance.IN, type.withNullability(false))),
            nullable = hasNulls(),
        ),
    )

)
There's some tests that need to be rewritten.

[Jolanrensen]

In theory, this should be the plugin version then:

fun isIntraComparable(col: SimpleDataColumn): Boolean {
    val comparable = StandardClassIds.Comparable.constructClassLikeType(
        typeArguments = arrayOf(col.type.type.withNullability(ConeNullability.NOT_NULL, session.typeContext)),
        isNullable = col.type.type.isNullable,
    )
    return col.type.type.isSubtypeOf(comparable, session)
}

(didn't test this yet)

[zaleslaw]

@Jolanrensen @koperagen please have a look at the supported scenarious, I want to merge it before #1078 changes

[koperagen]

Missing a test case here:
personsDf.groupBy { city }.mean { age and yearsToRetirement }.compareSchemas()

Runtime:
city: String
mean: Double

Compile:
age: Double
yearsToRetirement: Double
city: String

[zaleslaw]

It seems that we agreed that support of multiple columns is under the question?

[Jolanrensen]

True, at least until after the statistics types are set in stone. That said, for mean, the result will always be Double :)

[koperagen]

Let's call it GroupByMeanFor* to distinguish from GroupBy.mean

[zaleslaw]

Sorry, don't understand, GroupByMeanFor will be under mean/meanFor and GroupByMean0 will be under meanFor at that way.

[koperagen]

I suggest GroupByMeanDefault for GroupBy.mean(), GroupByMeanFor for GroupBy.meanFor and GroupByMean for mean

[zaleslaw]

Ok, agree, will do some renaming/refactoring in next PR in that area

[Jolanrensen]

Merged a failing build