From eeaada39b9c6d2f5be076ca7d0f9f92f0e8ef008 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Wed, 2 Sep 2026 13:08:26 +0200 Subject: [PATCH] feat(engine): extend SelectRunner with multiKeyComparator --- .../kotlin/kestrel/engine/SelectRunner.kt | 153 +++++++++++++++++- 1 file changed, 145 insertions(+), 8 deletions(-) diff --git a/src/main/kotlin/kestrel/engine/SelectRunner.kt b/src/main/kotlin/kestrel/engine/SelectRunner.kt index d947784..bafabab 100644 --- a/src/main/kotlin/kestrel/engine/SelectRunner.kt +++ b/src/main/kotlin/kestrel/engine/SelectRunner.kt @@ -246,8 +246,135 @@ class SelectRunner(private val catalog: Catalog) { } else { !keepOnMatch } + } + } + + private fun multiKeyComparator( + keys: List, + evaluate: (Expression, List) -> Value, + declaredCollationOf: (ColumnExpression) -> Collation? = { null } + ): Comparator> = Comparator { first, second -> + var result = 0 + for (key in keys) { + val comparison = compareOrderKeyValues( + evaluate(key.expression, first), evaluate(key.expression, second), key, declaredCollationOf + ) + if (comparison != 0) { + result = comparison + break + } + } + result + } + + private fun evaluateOutputRow(columns: List, expression: Expression, row: List): Value { + val resolver = kestrel.engine.ColumnResolver { column -> + if (column.qualifier != null) throw KestrelException("no such column: ${column.qualifier}.${column.name}") + val index = columns.indexOfFirst { it.equals(column.name, ignoreCase = true) } + if (index == -1) throw KestrelException("no such column: ${column.name}") + row[index] + } + return Evaluator(resolver, subqueryRunner).evaluate(expression) + } + + private fun runWith(query: WithQuery, outer: ColumnResolver?): SelectResult { + var pushed = 0 + try { + for (cte in query.ctes) { + activeCtes.add(buildCteScope(query.recursive, cte)) + pushed++ + } + return runQuery(query.body, outer) + } finally { + repeat(pushed) { activeCtes.removeAt(activeCtes.size - 1) } + } + } + + private fun buildCteScope(recursive: Boolean, cte: CommonTableExpression): CteScope { + val recursivePlan = if (recursive) splitRecursiveCte(cte) else null + if (recursivePlan == null) { + val result = runQuery(cte.query, null) + return CteScope(cte.name, cteSchema(cte, result), result.rows.map { Row(it) }) + } + val (base, operator, recursiveTerm) = recursivePlan + val baseResult = runQuery(base, null) + val schema = cteSchema(cte, baseResult) + val scope = CteScope(cte.name, schema, baseResult.rows.map { Row(it) }) + val declaredCollationOf = declaredRecursiveCteCollation(base, baseResult.columns) + val outputExpressions = baseResult.columns.map { ColumnExpression(it) } + val seenKeys = if (operator == SetOperator.Union) + baseResult.rows.mapTo(HashSet()) { collationFoldedKey(it, outputExpressions, declaredCollationOf) } + else null + activeCtes.add(scope) + try { + val accumulated = mutableListOf>() + accumulated.addAll(baseResult.rows) + var working = baseResult.rows + var iterations = 0 + while (working.isNotEmpty()) { + iterations++ + if (iterations > MAX_RECURSION_ITERATIONS) { + throw KestrelException( + "recursive common table expression '${cte.name}' did not terminate within $MAX_RECURSION_ITERATIONS iterations" + ) + } + scope.rows = working.map { Row(it) } + var newRows = runQuery(recursiveTerm, null).rows + if (operator == SetOperator.Union) { + newRows = newRows.filter { candidate -> seenKeys!!.add(collationFoldedKey(candidate, outputExpressions, declaredCollationOf)) } + } + if (newRows.isEmpty()) break + accumulated.addAll(newRows) + working = newRows + } + scope.rows = accumulated.map { Row(it) } + } finally { + activeCtes.removeAt(activeCtes.size - 1) + } + return scope + } + + private data class RecursiveCtePlan(val base: QueryStatement, val operator: SetOperator, val recursiveTerm: QueryStatement) + + private fun splitRecursiveCte(cte: CommonTableExpression): RecursiveCtePlan? { + val query = cte.query + if (query !is SetOperationStatement) return null + if (query.operator != SetOperator.Union && query.operator != SetOperator.UnionAll) return null + val leftRefs = referencesName(query.left, cte.name) + val rightRefs = referencesName(query.right, cte.name) + return when { + leftRefs && rightRefs -> + throw KestrelException("recursive common table expression '${cte.name}' must not reference itself on both sides of ${query.operator}") + rightRefs -> RecursiveCtePlan(query.left, query.operator, query.right) + leftRefs -> RecursiveCtePlan(query.right, query.operator, query.left) + else -> null + } + } + + private fun referencesName(query: QueryStatement, name: String): Boolean = when (query) { + is SelectStatement -> referencesName(query.from, name) || query.joins.any { referencesName(it.from, name) } + is SetOperationStatement -> referencesName(query.left, name) || referencesName(query.right, name) + is WithQuery -> query.ctes.none { it.name.equals(name, ignoreCase = true) } && referencesName(query.body, name) + } + + private fun referencesName(item: FromItem, name: String): Boolean = when (item) { + is TableSource -> item.name.equals(name, ignoreCase = true) + is SubquerySource -> referencesName(item.select, name) + } + + private fun cteSchema(cte: CommonTableExpression, result: SelectResult): TableSchema { + val labels = cte.columns ?: result.columns + if (labels.size != result.columns.size) { + throw KestrelException( + "common table expression '${cte.name}' expects ${labels.size} columns but its query produces ${result.columns.size}" + ) + } + return derivedSchema(cte.name, result, labels) + } + + private fun buildPlan(statement: SelectStatement, outer: ColumnResolver?): Plan { val sources = buildSources(statement) - val plan = Plan(statement, sources) + val plan = Plan(statement, sources, outer) plan.validate() return plan } @@ -255,17 +382,27 @@ class SelectRunner(private val catalog: Catalog) { private fun buildSources(statement: SelectStatement): List { val sources = mutableListOf() var offset = 0 - fun add(tableName: String, alias: String?) { - val table = catalog.table(tableName) - val chosen = alias ?: tableName + fun add(item: FromItem) { + // A LATERAL derived table's rows depend on the row of an earlier FROM item it + // correlates with, so they cannot be resolved once up front like every other source - + // only its column *names* are needed here (for offset bookkeeping and validating + // references to its alias elsewhere in the query), computed structurally from the + // subquery's own AST without running it. Real rows are produced per left row in + // [executeJoin]. + val (schema, rows) = if (item is SubquerySource && item.lateral) { + lateralSchema(item) to emptyList() + } else { + resolveSource(item) + } + val chosen = sourceLabel(item) if (sources.any { it.alias.equals(chosen, ignoreCase = true) }) { throw KestrelException("duplicate table alias: $chosen") } - sources.add(Source(chosen, table.schema.name, table.schema, table.rows, offset)) - offset += table.schema.columns.size + sources.add(Source(chosen, schema.name, schema, rows, offset)) + offset += schema.columns.size } - add(statement.tableName, statement.tableAlias) - for (join in statement.joins) add(join.tableName, join.alias) + add(statement.from) + for (join in statement.joins) add(join.from) return sources }