feat(sql): add frameRange and resolveFrameBound to DiskExecutor
This commit is contained in:
dualfroz committed
1 parent
d3a1880792
commit
cb73b8d3f9
1 file changed
+155
@@ -3510,3 +3510,158 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Resolves an explicit `ROWS|RANGE BETWEEN ... AND ...` [WindowFrame] to a range of positions
|
||||||
|
* into [ordered] for the row at [position] - clamped to the partition's edges rather than
|
||||||
|
* erroring, exactly like the implicit default frame already does. Mirrors
|
||||||
|
* [kestrel.engine.SelectRunner]'s own `frameRange` exactly.
|
||||||
|
*/
|
||||||
|
private fun frameRange(
|
||||||
|
window: WindowExpression,
|
||||||
|
position: Int,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): IntRange {
|
||||||
|
val frame = window.frame!!
|
||||||
|
if (frame.unit == FrameUnit.Groups) {
|
||||||
|
return groupsFrameRange(frame, window.orderBy, position, ordered, combined, evaluateRow)
|
||||||
|
}
|
||||||
|
val size = ordered.size
|
||||||
|
val start = resolveFrameBound(frame, frame.start, window, position, ordered, combined, isStart = true, evaluateRow)
|
||||||
|
val end = resolveFrameBound(frame, frame.end, window, position, ordered, combined, isStart = false, evaluateRow)
|
||||||
|
val clampedStart = start.coerceAtLeast(0)
|
||||||
|
val clampedEnd = end.coerceAtMost(size - 1)
|
||||||
|
return if (clampedStart > clampedEnd) IntRange.EMPTY else clampedStart..clampedEnd
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun resolveFrameBound(
|
||||||
|
frame: WindowFrame,
|
||||||
|
bound: FrameBound,
|
||||||
|
window: WindowExpression,
|
||||||
|
position: Int,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
isStart: Boolean,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int = when (bound.type) {
|
||||||
|
FrameBoundType.UnboundedPreceding -> 0
|
||||||
|
FrameBoundType.UnboundedFollowing -> ordered.size - 1
|
||||||
|
FrameBoundType.CurrentRow ->
|
||||||
|
if (frame.unit == FrameUnit.Rows) position else peerBoundary(window.orderBy, position, ordered, combined, isStart, evaluateRow)
|
||||||
|
FrameBoundType.Preceding ->
|
||||||
|
if (frame.unit == FrameUnit.Rows) {
|
||||||
|
position - resolveIntegerFrameOffset(bound, combined[ordered[position]], "ROWS", evaluateRow)
|
||||||
|
} else {
|
||||||
|
rangeBoundary(bound, window.orderBy, position, ordered, combined, isStart, subtract = true, evaluateRow)
|
||||||
|
}
|
||||||
|
FrameBoundType.Following ->
|
||||||
|
if (frame.unit == FrameUnit.Rows) {
|
||||||
|
position + resolveIntegerFrameOffset(bound, combined[ordered[position]], "ROWS", evaluateRow)
|
||||||
|
} else {
|
||||||
|
rangeBoundary(bound, window.orderBy, position, ordered, combined, isStart, subtract = false, evaluateRow)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun resolveIntegerFrameOffset(
|
||||||
|
bound: FrameBound,
|
||||||
|
row: List<Value>,
|
||||||
|
unit: String,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int {
|
||||||
|
val value = evaluateRow(bound.offset!!, row)
|
||||||
|
val amount = (value as? Value.IntegerValue)?.value ?: throw DiskExecutorException("$unit frame offset must be an integer")
|
||||||
|
if (amount < 0) throw DiskExecutorException("$unit frame offset must not be negative")
|
||||||
|
return amount.toInt()
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* `GROUPS BETWEEN <bound> AND <bound>` counts whole peer groups rather than physical rows
|
||||||
|
* ([FrameUnit.Rows]) or distance in a value ([FrameUnit.Range]) - mirrors
|
||||||
|
* [kestrel.engine.SelectRunner]'s own `groupsFrameRange` exactly.
|
||||||
|
*/
|
||||||
|
private fun groupsFrameRange(
|
||||||
|
frame: WindowFrame,
|
||||||
|
orderBy: List<OrderKey>,
|
||||||
|
position: Int,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): IntRange {
|
||||||
|
val groups = peerGroupRanges(orderBy, ordered, combined, evaluateRow)
|
||||||
|
val currentGroup = groups.indexOfFirst { position in it }
|
||||||
|
val row = combined[ordered[position]]
|
||||||
|
val startGroup = resolveGroupBound(frame.start, currentGroup, groups.size, row, evaluateRow).coerceAtLeast(0)
|
||||||
|
val endGroup = resolveGroupBound(frame.end, currentGroup, groups.size, row, evaluateRow).coerceAtMost(groups.size - 1)
|
||||||
|
return if (startGroup > endGroup) IntRange.EMPTY else groups[startGroup].first..groups[endGroup].last
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun resolveGroupBound(
|
||||||
|
bound: FrameBound,
|
||||||
|
currentGroup: Int,
|
||||||
|
groupCount: Int,
|
||||||
|
row: List<Value>,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int = when (bound.type) {
|
||||||
|
FrameBoundType.UnboundedPreceding -> 0
|
||||||
|
FrameBoundType.UnboundedFollowing -> groupCount - 1
|
||||||
|
FrameBoundType.CurrentRow -> currentGroup
|
||||||
|
FrameBoundType.Preceding -> currentGroup - resolveIntegerFrameOffset(bound, row, "GROUPS", evaluateRow)
|
||||||
|
FrameBoundType.Following -> currentGroup + resolveIntegerFrameOffset(bound, row, "GROUPS", evaluateRow)
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* `CURRENT ROW` in a RANGE frame means the boundary of the current row's peer group (every row
|
||||||
|
* tied with it on the OVER's ORDER BY, or the whole partition with no ORDER BY) - mirrors
|
||||||
|
* [kestrel.engine.SelectRunner]'s own `peerBoundary` exactly.
|
||||||
|
*/
|
||||||
|
private fun peerBoundary(
|
||||||
|
orderBy: List<OrderKey>,
|
||||||
|
position: Int,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
isStart: Boolean,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int {
|
||||||
|
if (orderBy.isEmpty()) return if (isStart) 0 else ordered.size - 1
|
||||||
|
fun keyAt(index: Int) = orderBy.map { evaluateRow(it.expression, combined[ordered[index]]) }
|
||||||
|
val key = keyAt(position)
|
||||||
|
var boundary = position
|
||||||
|
if (isStart) {
|
||||||
|
while (boundary > 0 && keyAt(boundary - 1) == key) boundary--
|
||||||
|
} else {
|
||||||
|
while (boundary < ordered.size - 1 && keyAt(boundary + 1) == key) boundary++
|
||||||
|
}
|
||||||
|
return boundary
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* `<value> PRECEDING`/`<value> FOLLOWING` in a RANGE frame measure distance in the value of a
|
||||||
|
* single numeric or DATE ORDER BY column rather than a row count - binary search works because
|
||||||
|
* [ordered] is already sorted on this column. Mirrors [kestrel.engine.SelectRunner]'s own
|
||||||
|
* `rangeBoundary` exactly.
|
||||||
|
*/
|
||||||
|
private fun rangeBoundary(
|
||||||
|
bound: FrameBound,
|
||||||
|
orderBy: List<OrderKey>,
|
||||||
|
position: Int,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
isStart: Boolean,
|
||||||
|
subtract: Boolean,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int {
|
||||||
|
if (orderBy.size != 1) {
|
||||||
|
throw DiskExecutorException("a RANGE frame with a PRECEDING/FOLLOWING bound requires exactly one ORDER BY column")
|
||||||
|
}
|
||||||
|
val key = orderBy.single()
|
||||||
|
val current = rangeLogicalValue(key, combined[ordered[position]], evaluateRow)
|
||||||
|
val offset = resolveRangeFrameOffset(bound, combined[ordered[position]], evaluateRow)
|
||||||
|
val threshold = if (subtract) current - offset else current + offset
|
||||||
|
return if (isStart) {
|
||||||
|
rangeLowerBound(ordered, combined, key, threshold, evaluateRow)
|
||||||
|
} else {
|
||||||
|
rangeUpperBound(ordered, combined, key, threshold, evaluateRow)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user