feat(engine): extend IndexAccessPlanner with ScoredScan and planIndex

src/main/kotlin/kestrel/engine/IndexAccessPlanner.kt

@@ -1,6 +1,7 @@
 package kestrel.engine
 
 import kestrel.catalog.Catalog
+import kestrel.catalog.SortedIndex
 import kestrel.catalog.TableSchema
 import kestrel.sql.ast.BinaryExpression
 import kestrel.sql.ast.BinaryOperator
@@ -14,6 +15,9 @@ sealed class PrimaryAccess {
     class IndexScan(val indexName: String, val description: String, val rowIndices: List<Int>) : PrimaryAccess()
 }
 
+/** An [IndexScan] candidate together with how many of the index's columns it actually pinned down, used to pick the most selective option among several usable indexes. */
+private class ScoredScan(val boundColumns: Int, val scan: PrimaryAccess.IndexScan)
+
 class IndexAccessPlanner(private val catalog: Catalog) {
     fun plan(tableName: String, alias: String, schema: TableSchema, predicate: Expression?): PrimaryAccess {
         if (predicate == null) return PrimaryAccess.FullScan
@@ -22,32 +26,75 @@ class IndexAccessPlanner(private val catalog: Catalog) {
             val simple = asColumnPredicate(conjunct, alias, schema) ?: continue
             byColumn.getOrPut(simple.first) { mutableListOf() }.add(simple.second to simple.third)
         }
-        for ((columnName, predicates) in byColumn) {
+        return catalog.indexesOn(tableName)
-            val index = catalog.indexOn(tableName, columnName) ?: continue
+            .mapNotNull { index -> planIndex(index, byColumn) }
-            val equal = predicates.firstOrNull { it.first == BinaryOperator.Equal }
+            .maxByOrNull { it.boundColumns }
-            if (equal != null) {
+            ?.scan
-                val candidates = index.equalTo(equal.second).sorted()
+            ?: PrimaryAccess.FullScan
-                return PrimaryAccess.IndexScan(index.definition.name, "$columnName = ${literalText(equal.second)}", candidates)
     }
-            var lower: Value? = null
+
+    /**
+     * Tries to serve [byColumn] from a single [index], preferring (in order): an exact match
+     * that binds every column of the index by equality; a range scan (<, <=, >, >=) on the
+     * column right after the longest leading equality run (works even with an empty run, i.e.
+     * a plain range on the index's very first column); or a pure equality prefix scan that
+     * binds a leading run of columns with the rest left entirely free. A composite index whose
+     * leading columns aren't bound at all (e.g. only its second column appears in the WHERE
+     * clause) is skipped - it cannot narrow the scan, since the index is sorted by its first
+     * column first.
+     */
+    private fun planIndex(index: SortedIndex, byColumn: Map<String, List<Pair<BinaryOperator, Value>>>): ScoredScan? {
+        val definition = index.definition
+        val columns = definition.columnNames
+        val equalityValues = mutableListOf<Value>()
+        val descriptionParts = mutableListOf<String>()
+        for (columnName in columns) {
+            val equal = byColumn[columnName]?.firstOrNull { it.first == BinaryOperator.Equal } ?: break
+            equalityValues.add(equal.second)
+            descriptionParts.add("$columnName = ${literalText(equal.second)}")
+        }
+        if (equalityValues.size == columns.size) {
+            val candidates = index.equalTo(equalityValues).sorted()
+            return ScoredScan(columns.size, PrimaryAccess.IndexScan(definition.name, descriptionParts.joinToString(" AND "), candidates))
+        }
+        val nextColumn = columns[equalityValues.size]
+        val rangePredicates = byColumn[nextColumn]?.filter { it.first != BinaryOperator.Equal }.orEmpty()
+        if (rangePredicates.isNotEmpty()) {
+            val scan = rangeScan(index, nextColumn, rangePredicates, prefix = equalityValues.toList(), descriptionParts)
+            if (scan != null) return scan
+        }
+        if (equalityValues.isEmpty()) return null
+        // A pure equality prefix: every bound column matched by "=", nothing further constrained.
+        val candidates = index.rangeOf(equalityValues, true, equalityValues, true).sorted()
+        return ScoredScan(equalityValues.size, PrimaryAccess.IndexScan(definition.name, descriptionParts.joinToString(" AND "), candidates))
+    }
+
+    private fun rangeScan(
+        index: SortedIndex,
+        columnName: String,
+        predicates: List<Pair<BinaryOperator, Value>>,
+        prefix: List<Value>,
+        descriptionPrefix: List<String>
+    ): ScoredScan? {
+        var lowerValue: Value? = null
         var lowerInclusive = false
-            var upper: Value? = null
+        var upperValue: Value? = null
         var upperInclusive = false
-            val parts = mutableListOf<String>()
+        val parts = descriptionPrefix.toMutableList()
         for ((operator, literal) in predicates) {
             when (operator) {
-                    BinaryOperator.Greater -> { lower = literal; lowerInclusive = false; parts.add("$columnName > ${literalText(literal)}") }
+                BinaryOperator.Greater -> { lowerValue = literal; lowerInclusive = false; parts.add("$columnName > ${literalText(literal)}") }
-                    BinaryOperator.GreaterEqual -> { lower = literal; lowerInclusive = true; parts.add("$columnName >= ${literalText(literal)}") }
+                BinaryOperator.GreaterEqual -> { lowerValue = literal; lowerInclusive = true; parts.add("$columnName >= ${literalText(literal)}") }
-                    BinaryOperator.Less -> { upper = literal; upperInclusive = false; parts.add("$columnName < ${literalText(literal)}") }
+                BinaryOperator.Less -> { upperValue = literal; upperInclusive = false; parts.add("$columnName < ${literalText(literal)}") }
-                    BinaryOperator.LessEqual -> { upper = literal; upperInclusive = true; parts.add("$columnName <= ${literalText(literal)}") }
+                BinaryOperator.LessEqual -> { upperValue = literal; upperInclusive = true; parts.add("$columnName <= ${literalText(literal)}") }
                 else -> Unit
             }
         }
-            if (lower == null && upper == null) continue
+        if (lowerValue == null && upperValue == null) return null
+        val lower = lowerValue?.let { prefix + it }
+        val upper = upperValue?.let { prefix + it }
         val candidates = index.rangeOf(lower, lowerInclusive, upper, upperInclusive).sorted()
-            return PrimaryAccess.IndexScan(index.definition.name, parts.joinToString(" AND "), candidates)
+        return ScoredScan(prefix.size + 1, PrimaryAccess.IndexScan(index.definition.name, parts.joinToString(" AND "), candidates))
-        }
-        return PrimaryAccess.FullScan
     }
 
     private fun asColumnPredicate(
@@ -95,6 +142,8 @@ class IndexAccessPlanner(private val catalog: Catalog) {
         is Value.RealValue -> value.value.toString()
         is Value.TextValue -> "'${value.value}'"
         is Value.BooleanValue -> value.value.toString()
+        is Value.DecimalValue -> value.value.toPlainString()
+        is Value.DateValue -> value.value.toString()
         is Value.NullValue -> "NULL"
     }
 

src/main/kotlin/kestrel/engine/OrderKeyOrdering.kt

@@ -0,0 +1,49 @@
+package kestrel.engine
+
+import kestrel.sql.ast.CollateExpression
+import kestrel.sql.ast.ColumnExpression
+import kestrel.sql.ast.OrderKey
+import kestrel.storage.Collation
+import kestrel.storage.Value
+import kestrel.storage.ValueOrdering
+
+/**
+ * Compares two already-evaluated ORDER BY key values, shared by every sort in both the in-memory
+ * and disk engines (top-level ORDER BY, GROUP BY's own ORDER BY, STRING_AGG's inner ORDER BY, and
+ * a window function's `OVER (... ORDER BY ...)`) so `NULLS FIRST`/`NULLS LAST` behaves identically
+ * everywhere an [OrderKey] is honored.
+ *
+ * `NULLS FIRST`/`NULLS LAST` is an explicit override: when either side is NULL, [OrderKey.descending]
+ * plays no part - only [OrderKey.nullsFirst] decides, matching the SQL standard (unlike a plain
+ * value comparison, negating the whole result for DESC would also flip which end NULLs land on,
+ * which is not what an explicit NULLS clause means). When [OrderKey.nullsFirst] is unset (`null`,
+ * meaning the query never wrote `NULLS FIRST`/`LAST`), it defaults to `!descending` - the exact
+ * behavior every ORDER BY in this engine had before this clause existed, since [ValueOrdering]
+ * ranks NULL as the smallest value and the old code just negated that for DESC.
+ *
+ * [key]'s own [OrderKey.expression] also decides which [Collation] two [Value.TextValue]s are
+ * compared with, in the same two-tier precedence [kestrel.engine.Evaluator] uses for a binary
+ * comparison: (1) an explicit `COLLATE` - `ORDER BY name COLLATE NOCASE` reaches here as a plain
+ * [OrderKey] whose expression is a [CollateExpression]; (2) failing that, a bare [ColumnExpression]
+ * key's own declared `COLLATE` (see [kestrel.catalog.Column.collation]), reported by
+ * [declaredCollationOf] - `{ null }` by default so every pre-existing call site (one with no
+ * schema to consult, such as a UNION/INTERSECT/EXCEPT's synthetic output row) keeps behaving
+ * exactly as before; (3) [Collation.Binary] otherwise.
+ */
+fun compareOrderKeyValues(
+    left: Value,
+    right: Value,
+    key: OrderKey,
+    declaredCollationOf: (ColumnExpression) -> Collation? = { null }
+): Int {
+    val leftNull = left is Value.NullValue
+    val rightNull = right is Value.NullValue
+    if (leftNull || rightNull) {
+        if (leftNull && rightNull) return 0
+        val nullsFirst = key.nullsFirst ?: !key.descending
+        return if (leftNull == nullsFirst) -1 else 1
+    }
+    val collation = resolveCollation(key.expression, declaredCollationOf)
+    val comparison = ValueOrdering.compare(left, right, collation)
+    return if (key.descending) -comparison else comparison
+}