diff --git a/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt b/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt index 240703d..0e55298 100644 --- a/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt +++ b/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt @@ -3665,3 +3665,169 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess { } } + private fun rangeLowerBound( + ordered: List, + combined: List>, + key: OrderKey, + threshold: BigDecimal, + evaluateRow: (Expression, List) -> 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, + combined: List>, + key: OrderKey, + threshold: BigDecimal, + evaluateRow: (Expression, List) -> 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, evaluateRow: (Expression, List) -> 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, evaluateRow: (Expression, List) -> 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, + ordered: List, + combined: List>, + evaluateRow: (Expression, List) -> Value + ): List { + if (orderBy.isEmpty()) return listOf(ordered.indices) + val keys = ordered.map { index -> orderBy.map { key -> evaluateRow(key.expression, combined[index]) } } + val ranges = mutableListOf() + 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>, evaluateRow: (Expression, List) -> 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, + combined: List>, + out: Array, + direction: Int, + evaluateRow: (Expression, List) -> 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, evaluateRow: (Expression, List) -> 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, + combined: List>, + out: Array, + evaluateRow: (Expression, List) -> 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> = + 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) + } +