feat(sql): add computeWindowValues and assignWindowValues

This commit is contained in:
dualfroz committed 2026-09-07 16:50:39 +02:00
1 parent de2cc8854d
commit d3a1880792
1 file changed
+150
@@ -3360,3 +3360,153 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
} }
} }
private fun computeWindowValues(
window: WindowExpression,
combined: List<List<Value>>,
evaluateRow: (Expression, List<Value>) -> Value,
declaredCollationOf: (ColumnExpression) -> Collation? = { null }
): List<Value> {
val values = arrayOfNulls<Value>(combined.size)
val partitions = combined.indices.groupBy { index -> window.partitionBy.map { evaluateRow(it, combined[index]) } }
for (memberIndices in partitions.values) {
val ordered = if (window.orderBy.isEmpty()) {
memberIndices
} else {
memberIndices.sortedWith(Comparator { first, second ->
var result = 0
for (key in window.orderBy) {
val comparison = compareOrderKeyValues(
evaluateRow(key.expression, combined[first]),
evaluateRow(key.expression, combined[second]),
key,
declaredCollationOf
)
if (comparison != 0) {
result = comparison
break
}
}
result
})
}
assignWindowValues(window, ordered, combined, values, evaluateRow)
}
return values.map { it ?: error("window value not computed for every row") }
}
private fun assignWindowValues(
window: WindowExpression,
ordered: List<Int>,
combined: List<List<Value>>,
out: Array<Value?>,
evaluateRow: (Expression, List<Value>) -> Value
) {
when (window.kind) {
WindowFunctionKind.RowNumber ->
ordered.forEachIndexed { position, index -> out[index] = Value.IntegerValue((position + 1).toLong()) }
WindowFunctionKind.Rank, WindowFunctionKind.DenseRank -> {
var rank = 0
var denseRank = 0
var previousKey: List<Value>? = null
ordered.forEachIndexed { position, index ->
val key = window.orderBy.map { evaluateRow(it.expression, combined[index]) }
if (previousKey == null || key != previousKey) {
rank = position + 1
denseRank++
previousKey = key
}
out[index] = Value.IntegerValue((if (window.kind == WindowFunctionKind.Rank) rank else denseRank).toLong())
}
}
WindowFunctionKind.Aggregate -> assignAggregateWindowValues(window, ordered, combined, out, evaluateRow)
WindowFunctionKind.Lag -> assignOffsetWindowValues(window, ordered, combined, out, direction = -1, evaluateRow)
WindowFunctionKind.Lead -> assignOffsetWindowValues(window, ordered, combined, out, direction = 1, evaluateRow)
WindowFunctionKind.FirstValue -> {
val frame = window.frame
if (frame == null) {
val firstValue = evaluateRow(window.argument!!, combined[ordered.first()])
ordered.forEach { index -> out[index] = firstValue }
} else {
ordered.indices.forEach { position ->
val members = frameMembers(window, position, ordered, combined, evaluateRow)
out[ordered[position]] =
if (members.isEmpty()) Value.NullValue else evaluateRow(window.argument!!, combined[ordered[members.first()]])
}
}
}
WindowFunctionKind.LastValue -> {
val frame = window.frame
if (frame == null) {
for (range in peerGroupRanges(window.orderBy, ordered, combined, evaluateRow)) {
val value = evaluateRow(window.argument!!, combined[ordered[range.last]])
for (position in range) out[ordered[position]] = value
}
} else {
ordered.indices.forEach { position ->
val members = frameMembers(window, position, ordered, combined, evaluateRow)
out[ordered[position]] =
if (members.isEmpty()) Value.NullValue else evaluateRow(window.argument!!, combined[ordered[members.last()]])
}
}
}
WindowFunctionKind.Ntile -> assignNtileWindowValues(window, ordered, combined, out, evaluateRow)
}
}
/**
* With no ORDER BY in OVER, the whole partition is one frame (every row gets the same value).
* With an ORDER BY, the default SQL frame is `RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT
* ROW`: a running aggregate where rows tied on the order key (a "peer group") share one value -
* mirroring [kestrel.engine.SelectRunner]'s own `assignAggregateWindowValues` exactly.
*/
private fun assignAggregateWindowValues(
window: WindowExpression,
ordered: List<Int>,
combined: List<List<Value>>,
out: Array<Value?>,
evaluateRow: (Expression, List<Value>) -> Value
) {
val frame = window.frame
if (frame == null) {
for (range in peerGroupRanges(window.orderBy, ordered, combined, evaluateRow)) {
val value = aggregateWindowValue(window, (0..range.last).map { combined[ordered[it]] }, evaluateRow)
for (position in range) out[ordered[position]] = value
}
} else {
ordered.indices.forEach { position ->
val members = frameMembers(window, position, ordered, combined, evaluateRow)
out[ordered[position]] = aggregateWindowValue(window, members.map { combined[ordered[it]] }, evaluateRow)
}
}
}
/**
* The frame's base [frameRange], narrowed by its [FrameExclude] mode (default
* [FrameExclude.NoOthers], which leaves the range untouched) - mirrors
* [kestrel.engine.SelectRunner]'s own `frameMembers` exactly.
*/
private fun frameMembers(
window: WindowExpression,
position: Int,
ordered: List<Int>,
combined: List<List<Value>>,
evaluateRow: (Expression, List<Value>) -> Value
): List<Int> {
val range = frameRange(window, position, ordered, combined, evaluateRow)
if (range.isEmpty()) return emptyList()
return when (window.frame!!.exclude) {
FrameExclude.NoOthers -> range.toList()
FrameExclude.CurrentRow -> range.filter { it != position }
FrameExclude.Group -> {
val peer = peerBoundary(window.orderBy, position, ordered, combined, isStart = true, evaluateRow)..
peerBoundary(window.orderBy, position, ordered, combined, isStart = false, evaluateRow)
range.filter { it !in peer }
}
FrameExclude.Ties -> {
val peer = peerBoundary(window.orderBy, position, ordered, combined, isStart = true, evaluateRow)..
peerBoundary(window.orderBy, position, ordered, combined, isStart = false, evaluateRow)
range.filter { it == position || it !in peer }
}
}
}