feat(sql): add computeWindowValues and assignWindowValues
This commit is contained in:
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 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user