feat: add index access planning and group aggregation
This commit is contained in:
dualfroz committed
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
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user