diff --git a/src/main/kotlin/kestrel/engine/GroupAggregation.kt b/src/main/kotlin/kestrel/engine/GroupAggregation.kt new file mode 100644 index 0000000..0a72735 --- /dev/null +++ b/src/main/kotlin/kestrel/engine/GroupAggregation.kt @@ -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, + private val label: (SelectItem, Int) -> String +) { + fun aggregate(statement: SelectStatement, orderExpression: Expression?, combined: List>): SelectResult { + val groups = LinkedHashMap, MutableList>>() + for (row in combined) { + val key = statement.groupBy.map { evaluateRow(it, row) } + groups.getOrPut(key) { mutableListOf() }.add(row) + } + var groupRows: List>> = + 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>> { 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>): 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>): 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>): 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 { + 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 { + 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 +} diff --git a/src/main/kotlin/kestrel/engine/IndexAccessPlanner.kt b/src/main/kotlin/kestrel/engine/IndexAccessPlanner.kt new file mode 100644 index 0000000..dacd328 --- /dev/null +++ b/src/main/kotlin/kestrel/engine/IndexAccessPlanner.kt @@ -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) : 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>>() + 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() + 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? { + 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 = + 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 + ) + } +}