feat: add index access planning and group aggregation

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>) -> 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
+}

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<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
+        )
+    }
+}