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