feat(sql): add rangeLowerBound and rangeUpperBound to DiskExecutor
This commit is contained in:
dualfroz committed
1 parent
cb73b8d3f9
commit
441b6073ee
1 file changed
+166
@@ -3665,3 +3665,169 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private fun rangeLowerBound(
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
key: OrderKey,
|
||||||
|
threshold: BigDecimal,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int {
|
||||||
|
var low = 0
|
||||||
|
var high = ordered.size
|
||||||
|
while (low < high) {
|
||||||
|
val mid = (low + high) / 2
|
||||||
|
if (rangeLogicalValue(key, combined[ordered[mid]], evaluateRow) >= threshold) high = mid else low = mid + 1
|
||||||
|
}
|
||||||
|
return low
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun rangeUpperBound(
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
key: OrderKey,
|
||||||
|
threshold: BigDecimal,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): Int {
|
||||||
|
var low = 0
|
||||||
|
var high = ordered.size
|
||||||
|
while (low < high) {
|
||||||
|
val mid = (low + high) / 2
|
||||||
|
if (rangeLogicalValue(key, combined[ordered[mid]], evaluateRow) <= threshold) low = mid + 1 else high = mid
|
||||||
|
}
|
||||||
|
return low - 1
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun rangeLogicalValue(key: OrderKey, row: List<Value>, evaluateRow: (Expression, List<Value>) -> Value): BigDecimal {
|
||||||
|
val magnitude = rangeMagnitude(evaluateRow(key.expression, row))
|
||||||
|
return if (key.descending) magnitude.negate() else magnitude
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun rangeMagnitude(value: Value): BigDecimal = when (value) {
|
||||||
|
is Value.IntegerValue -> BigDecimal(value.value)
|
||||||
|
is Value.DecimalValue -> value.value
|
||||||
|
is Value.RealValue -> BigDecimal.valueOf(value.value)
|
||||||
|
is Value.DateValue -> BigDecimal(value.value.toEpochDay())
|
||||||
|
else -> throw DiskExecutorException("a RANGE frame requires the ORDER BY column to be numeric or DATE")
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun resolveRangeFrameOffset(bound: FrameBound, row: List<Value>, evaluateRow: (Expression, List<Value>) -> Value): BigDecimal {
|
||||||
|
val offset = when (val value = evaluateRow(bound.offset!!, row)) {
|
||||||
|
is Value.IntegerValue -> BigDecimal(value.value)
|
||||||
|
is Value.DecimalValue -> value.value
|
||||||
|
is Value.RealValue -> BigDecimal.valueOf(value.value)
|
||||||
|
else -> throw DiskExecutorException("RANGE frame offset must be numeric")
|
||||||
|
}
|
||||||
|
if (offset.signum() < 0) throw DiskExecutorException("RANGE frame offset must not be negative")
|
||||||
|
return offset
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Splits [ordered] into contiguous "peer group" ranges of rows tied on the order key - or one
|
||||||
|
* range spanning the whole partition when there is no ORDER BY. Mirrors
|
||||||
|
* [kestrel.engine.SelectRunner]'s own `peerGroupRanges` exactly.
|
||||||
|
*/
|
||||||
|
private fun peerGroupRanges(
|
||||||
|
orderBy: List<OrderKey>,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
): List<IntRange> {
|
||||||
|
if (orderBy.isEmpty()) return listOf(ordered.indices)
|
||||||
|
val keys = ordered.map { index -> orderBy.map { key -> evaluateRow(key.expression, combined[index]) } }
|
||||||
|
val ranges = mutableListOf<IntRange>()
|
||||||
|
var i = 0
|
||||||
|
while (i < ordered.size) {
|
||||||
|
var j = i
|
||||||
|
while (j < ordered.size && keys[j] == keys[i]) j++
|
||||||
|
ranges.add(i until j)
|
||||||
|
i = j
|
||||||
|
}
|
||||||
|
return ranges
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun aggregateWindowValue(window: WindowExpression, rows: List<List<Value>>, evaluateRow: (Expression, List<Value>) -> Value): Value =
|
||||||
|
AggregateFunctions.compute(window.name, window.argument, window.star, rows, evaluateRow)
|
||||||
|
|
||||||
|
/**
|
||||||
|
* LAG (direction -1) / LEAD (direction +1): reads [WindowExpression.argument] from the row
|
||||||
|
* [WindowExpression.offset] positions before/after the current one in [ordered] (default 1).
|
||||||
|
* Past either edge of the partition, falls back to [WindowExpression.defaultValue] (evaluated
|
||||||
|
* against the *current* row) or NULL. Mirrors [kestrel.engine.SelectRunner]'s own
|
||||||
|
* `assignOffsetWindowValues` exactly.
|
||||||
|
*/
|
||||||
|
private fun assignOffsetWindowValues(
|
||||||
|
window: WindowExpression,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
out: Array<Value?>,
|
||||||
|
direction: Int,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
) {
|
||||||
|
ordered.forEachIndexed { position, index ->
|
||||||
|
val offset = resolveOffset(window, combined[index], evaluateRow)
|
||||||
|
val targetPosition = position + direction * offset
|
||||||
|
out[index] = if (targetPosition in ordered.indices) {
|
||||||
|
evaluateRow(window.argument!!, combined[ordered[targetPosition]])
|
||||||
|
} else {
|
||||||
|
window.defaultValue?.let { evaluateRow(it, combined[index]) } ?: Value.NullValue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun resolveOffset(window: WindowExpression, row: List<Value>, evaluateRow: (Expression, List<Value>) -> Value): Int {
|
||||||
|
val expression = window.offset ?: return 1
|
||||||
|
val value = evaluateRow(expression, row)
|
||||||
|
val amount = (value as? Value.IntegerValue)?.value ?: throw DiskExecutorException("${window.name}() offset must be an integer")
|
||||||
|
if (amount < 0) throw DiskExecutorException("${window.name}() offset must not be negative")
|
||||||
|
return amount.toInt()
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* NTILE(n): splits the partition into `n` buckets as evenly as possible - when the partition
|
||||||
|
* size does not divide evenly, the first `size % n` buckets get one extra row. Mirrors
|
||||||
|
* [kestrel.engine.SelectRunner]'s own `assignNtileWindowValues` exactly.
|
||||||
|
*/
|
||||||
|
private fun assignNtileWindowValues(
|
||||||
|
window: WindowExpression,
|
||||||
|
ordered: List<Int>,
|
||||||
|
combined: List<List<Value>>,
|
||||||
|
out: Array<Value?>,
|
||||||
|
evaluateRow: (Expression, List<Value>) -> Value
|
||||||
|
) {
|
||||||
|
val bucketValue = evaluateRow(window.argument!!, combined[ordered.first()])
|
||||||
|
val buckets = (bucketValue as? Value.IntegerValue)?.value ?: throw DiskExecutorException("${window.name}() bucket count must be an integer")
|
||||||
|
if (buckets <= 0) throw DiskExecutorException("${window.name}() bucket count must be positive")
|
||||||
|
val bucketCount = buckets.toInt()
|
||||||
|
val base = ordered.size / bucketCount
|
||||||
|
val remainder = ordered.size % bucketCount
|
||||||
|
var position = 0
|
||||||
|
for (bucket in 1..bucketCount) {
|
||||||
|
val size = base + if (bucket <= remainder) 1 else 0
|
||||||
|
repeat(size) {
|
||||||
|
out[ordered[position]] = Value.IntegerValue(bucket.toLong())
|
||||||
|
position++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Candidate (rid, row) pairs for [where] over [table]: an indexed
|
||||||
|
* point/range/secondary lookup via [DiskIndexAccessPlanner] when the
|
||||||
|
* predicate allows it, a full [DiskTable.scan] otherwise. Callers must
|
||||||
|
* still evaluate the full [where] against whatever this returns - the
|
||||||
|
* chosen plan only narrows candidates, it never claims to satisfy every
|
||||||
|
* conjunct of the predicate (see [DiskAccessPlan]'s own documentation).
|
||||||
|
*/
|
||||||
|
private fun candidateRows(table: DiskTable, txn: Long, where: Expression?, alias: String?): List<Pair<RecordId, Row>> =
|
||||||
|
when (val plan = accessPlanner.plan(table.schema, table.secondaryIndexes, where, alias)) {
|
||||||
|
DiskAccessPlan.FullScan -> table.scan(txn)
|
||||||
|
is DiskAccessPlan.PrimaryKeyPoint -> listOfNotNull(table.findByPrimaryKey(txn, plan.key))
|
||||||
|
is DiskAccessPlan.PrimaryKeyRange -> table.rangeByPrimaryKey(txn, plan.lower, plan.lowerInclusive, plan.upper, plan.upperInclusive)
|
||||||
|
is DiskAccessPlan.CompositePrimaryKeyPoint -> listOfNotNull(table.findByPrimaryKey(txn, plan.keys))
|
||||||
|
is DiskAccessPlan.CompositePrimaryKeyRange ->
|
||||||
|
table.rangeByPrimaryKey(txn, plan.lower, plan.lowerInclusive, plan.upper, plan.upperInclusive, plan.prefix)
|
||||||
|
is DiskAccessPlan.Secondary -> table.findBySecondary(txn, plan.indexName, plan.keys)
|
||||||
|
is DiskAccessPlan.SecondaryRange ->
|
||||||
|
table.rangeBySecondary(txn, plan.indexName, plan.prefix, plan.lower, plan.lowerInclusive, plan.upper, plan.upperInclusive)
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user