feat: add index access planning and group aggregation

This commit is contained in:
dualfroz committed 2026-08-25 23:10:08 +02:00
1 parent 9fc6338bbd
commit def1ceb55d
2 files changed
+212

No files matched your search

@@ -0,0 +1,105 @@
package kestrel.engine
import kestrel.KestrelException
import kestrel.sql.ast.BinaryExpression
import kestrel.sql.ast.Expression
import kestrel.sql.ast.ExpressionProjection
import kestrel.sql.ast.FunctionExpression
import kestrel.sql.ast.LiteralExpression
import kestrel.sql.ast.SelectItem
import kestrel.sql.ast.SelectStatement
import kestrel.sql.ast.UnaryExpression
import kestrel.storage.Value
import kestrel.storage.ValueOrdering
import kestrel.storage.typeName
class GroupAggregation(
private val resolveIndex: (kestrel.sql.ast.ColumnExpression) -> Int,
private val evaluateRow: (Expression, List<Value>) -> Value,
private val label: (SelectItem, Int) -> String
) {
fun aggregate(statement: SelectStatement, orderExpression: Expression?, combined: List<List<Value>>): SelectResult {
val groups = LinkedHashMap<List<Value>, MutableList<List<Value>>>()
for (row in combined) {
val key = statement.groupBy.map { evaluateRow(it, row) }
groups.getOrPut(key) { mutableListOf() }.add(row)
}
var groupRows: List<List<List<Value>>> =
if (statement.groupBy.isEmpty() && groups.isEmpty()) listOf(emptyList()) else groups.values.toList()
statement.having?.let { having ->
groupRows = groupRows.filter { isTrue(evaluateInGroup(having, it)) }
}
orderExpression?.let { expression ->
val comparator = Comparator<List<List<Value>>> { first, second ->
ValueOrdering.compare(evaluateInGroup(expression, first), evaluateInGroup(expression, second))
}
groupRows = groupRows.sortedWith(if (statement.orderBy!!.descending) comparator.reversed() else comparator)
}
statement.limit?.let {
if (it < 0) throw KestrelException("LIMIT must not be negative")
groupRows = groupRows.take(it)
}
val projection = statement.projection as ExpressionProjection
val labels = projection.items.mapIndexed { index, item -> label(item, index) }
val projected = groupRows.map { group -> projection.items.map { evaluateInGroup(it.expression, group) } }
return SelectResult(labels, projected)
}
private fun evaluateInGroup(expression: Expression, group: List<List<Value>>): Value {
val substituted = substituteAggregates(expression, group)
val resolver = ColumnResolver { columnExpression -> group.first()[resolveIndex(columnExpression)] }
return Evaluator(resolver).evaluate(substituted)
}
private fun substituteAggregates(expression: Expression, group: List<List<Value>>): Expression = when (expression) {
is FunctionExpression -> LiteralExpression(computeAggregate(expression, group))
is BinaryExpression -> BinaryExpression(
expression.operator,
substituteAggregates(expression.left, group),
substituteAggregates(expression.right, group)
)
is UnaryExpression -> UnaryExpression(expression.operator, substituteAggregates(expression.operand, group))
else -> expression
}
private fun computeAggregate(function: FunctionExpression, group: List<List<Value>>): Value {
val name = function.name.uppercase()
if (name == "COUNT") {
if (function.star) return Value.IntegerValue(group.size.toLong())
val argument = function.argument ?: return Value.IntegerValue(group.size.toLong())
val count = group.count { evaluateRow(argument, it) !is Value.NullValue }
return Value.IntegerValue(count.toLong())
}
val argument = function.argument
?: throw KestrelException("aggregate function ${function.name} requires an argument")
val values = group.map { evaluateRow(argument, it) }.filter { it !is Value.NullValue }
return when (name) {
"SUM" -> sum(values)
"AVG" -> average(values)
"MIN" -> values.minWithOrNull(ValueOrdering) ?: Value.NullValue
"MAX" -> values.maxWithOrNull(ValueOrdering) ?: Value.NullValue
else -> throw KestrelException("unknown function: ${function.name}")
}
}
private fun sum(values: List<Value>): Value {
if (values.isEmpty()) return Value.NullValue
if (values.all { it is Value.IntegerValue }) {
return Value.IntegerValue(values.sumOf { (it as Value.IntegerValue).value })
}
return Value.RealValue(values.sumOf { numeric(it) })
}
private fun average(values: List<Value>): Value {
if (values.isEmpty()) return Value.NullValue
return Value.RealValue(values.sumOf { numeric(it) } / values.size)
}
private fun numeric(value: Value): Double = when (value) {
is Value.IntegerValue -> value.value.toDouble()
is Value.RealValue -> value.value
else -> throw KestrelException("aggregate requires numeric operands but got ${value.typeName}")
}
private fun isTrue(value: Value): Boolean = value is Value.BooleanValue && value.value
}
@@ -0,0 +1,107 @@
package kestrel.engine
import kestrel.catalog.Catalog
import kestrel.catalog.TableSchema
import kestrel.sql.ast.BinaryExpression
import kestrel.sql.ast.BinaryOperator
import kestrel.sql.ast.ColumnExpression
import kestrel.sql.ast.Expression
import kestrel.sql.ast.LiteralExpression
import kestrel.storage.Value
sealed class PrimaryAccess {
data object FullScan : PrimaryAccess()
class IndexScan(val indexName: String, val description: String, val rowIndices: List<Int>) : PrimaryAccess()
}
class IndexAccessPlanner(private val catalog: Catalog) {
fun plan(tableName: String, alias: String, schema: TableSchema, predicate: Expression?): PrimaryAccess {
if (predicate == null) return PrimaryAccess.FullScan
val byColumn = LinkedHashMap<String, MutableList<Pair<BinaryOperator, Value>>>()
for (conjunct in conjuncts(predicate)) {
val simple = asColumnPredicate(conjunct, alias, schema) ?: continue
byColumn.getOrPut(simple.first) { mutableListOf() }.add(simple.second to simple.third)
}
for ((columnName, predicates) in byColumn) {
val index = catalog.indexOn(tableName, columnName) ?: continue
val equal = predicates.firstOrNull { it.first == BinaryOperator.Equal }
if (equal != null) {
val candidates = index.equalTo(equal.second).sorted()
return PrimaryAccess.IndexScan(index.definition.name, "$columnName = ${literalText(equal.second)}", candidates)
}
var lower: Value? = null
var lowerInclusive = false
var upper: Value? = null
var upperInclusive = false
val parts = mutableListOf<String>()
for ((operator, literal) in predicates) {
when (operator) {
BinaryOperator.Greater -> { lower = literal; lowerInclusive = false; parts.add("$columnName > ${literalText(literal)}") }
BinaryOperator.GreaterEqual -> { lower = literal; lowerInclusive = true; parts.add("$columnName >= ${literalText(literal)}") }
BinaryOperator.Less -> { upper = literal; upperInclusive = false; parts.add("$columnName < ${literalText(literal)}") }
BinaryOperator.LessEqual -> { upper = literal; upperInclusive = true; parts.add("$columnName <= ${literalText(literal)}") }
else -> Unit
}
}
if (lower == null && upper == null) continue
val candidates = index.rangeOf(lower, lowerInclusive, upper, upperInclusive).sorted()
return PrimaryAccess.IndexScan(index.definition.name, parts.joinToString(" AND "), candidates)
}
return PrimaryAccess.FullScan
}
private fun asColumnPredicate(
expression: Expression,
alias: String,
schema: TableSchema
): Triple<String, BinaryOperator, Value>? {
if (expression !is BinaryExpression) return null
val operator = expression.operator
if (operator !in RANGE_OPERATORS) return null
val left = expression.left
val right = expression.right
fun matches(column: ColumnExpression): String? {
if (column.qualifier != null && !column.qualifier.equals(alias, ignoreCase = true)) return null
return if (schema.columnIndex(column.name) != null) column.name else null
}
if (left is ColumnExpression && right is LiteralExpression && right.value !is Value.NullValue) {
val name = matches(left) ?: return null
return Triple(name, operator, right.value)
}
if (right is ColumnExpression && left is LiteralExpression && left.value !is Value.NullValue) {
val name = matches(right) ?: return null
return Triple(name, flip(operator), left.value)
}
return null
}
private fun conjuncts(expression: Expression): List<Expression> =
if (expression is BinaryExpression && expression.operator == BinaryOperator.And) {
conjuncts(expression.left) + conjuncts(expression.right)
} else {
listOf(expression)
}
private fun flip(operator: BinaryOperator): BinaryOperator = when (operator) {
BinaryOperator.Less -> BinaryOperator.Greater
BinaryOperator.LessEqual -> BinaryOperator.GreaterEqual
BinaryOperator.Greater -> BinaryOperator.Less
BinaryOperator.GreaterEqual -> BinaryOperator.LessEqual
else -> operator
}
private fun literalText(value: Value): String = when (value) {
is Value.IntegerValue -> value.value.toString()
is Value.RealValue -> value.value.toString()
is Value.TextValue -> "'${value.value}'"
is Value.BooleanValue -> value.value.toString()
is Value.NullValue -> "NULL"
}
private companion object {
val RANGE_OPERATORS = setOf(
BinaryOperator.Equal, BinaryOperator.Less, BinaryOperator.LessEqual,
BinaryOperator.Greater, BinaryOperator.GreaterEqual
)
}
}