From b67ac3386ac5e9a2f0c1dee3a6b596ed213ade34 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Mon, 7 Sep 2026 12:04:22 +0200 Subject: [PATCH] feat(sql): add declaredSetOperationCollation --- .../kotlin/kestrel/store/sql/DiskExecutor.kt | 174 ++++++++++++++++++ 1 file changed, 174 insertions(+) diff --git a/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt b/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt index a655050..cdc7a6c 100644 --- a/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt +++ b/src/main/kotlin/kestrel/store/sql/DiskExecutor.kt @@ -2399,3 +2399,177 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess { return SelectResult(left.columns, rows) } + /** + * Mirrors [kestrel.engine.SelectRunner]'s own `declaredSetOperationCollation` 1:1 (separate + * copy, historical reasons like the rest of this class): the declared `COLLATE` of a set + * operation's combined output column, traced back to a real table column in [statement]'s LEFT + * branch only, resolving only for a plain `SELECT ... FROM realTable [JOIN realTable ...]` with + * no CTE, view, or derived table anywhere in its FROM clause. Everything else - an expression, + * a nested set operation, a subquery source - reports no declared collation, same as before this + * function existed. + */ + private fun declaredSetOperationCollation(statement: SetOperationStatement, columns: List): (ColumnExpression) -> Collation? = + (statement.left as? SelectStatement)?.let { declaredCollationForSelectOutput(it, columns) } ?: { null } + + /** + * Mirrors [kestrel.engine.SelectRunner]'s own `declaredCollationForSelectOutput` 1:1 (separate + * copy, historical reasons like the rest of this class): factored out of + * [declaredSetOperationCollation] so [declaredRecursiveCteCollation] can apply the identical + * "plain SELECT over real tables" resolution to a recursive CTE's own anchor (base) term. + */ + private fun declaredCollationForSelectOutput(select: SelectStatement, columns: List): (ColumnExpression) -> Collation? { + val items = listOf(select.from) + select.joins.map { it.from } + val sources = items.map { item -> + (item as? TableSource)?.let { source -> catalog.tableOrNull(source.name)?.schema?.let { (source.alias ?: source.name) to it } } + } + if (sources.any { it == null }) return { null } + val resolvedSources = sources.filterNotNull() + fun columnCollation(expression: ColumnExpression): Collation? { + val matches = resolvedSources + .filter { (alias, _) -> expression.qualifier == null || alias.equals(expression.qualifier, ignoreCase = true) } + .mapNotNull { (_, schema) -> schema.column(expression.name) } + return matches.singleOrNull()?.collation + } + return when (val projection = select.projection) { + is ExpressionProjection -> outputExpr@{ outputExpr -> + val index = columns.indexOfFirst { it.equals(outputExpr.name, ignoreCase = true) } + val itemExpression = index.takeIf { it in projection.items.indices }?.let { projection.items[it].expression } + (itemExpression as? ColumnExpression)?.let(::columnCollation) + } + is StarProjection -> outputExpr@{ outputExpr -> + val index = columns.indexOfFirst { it.equals(outputExpr.name, ignoreCase = true) } + if (index < 0) return@outputExpr null + var remaining = index + for ((_, schema) in resolvedSources) { + if (remaining < schema.columns.size) return@outputExpr schema.columns[remaining].collation + remaining -= schema.columns.size + } + null + } + } + } + + /** Mirrors [kestrel.engine.SelectRunner]'s own `declaredRecursiveCteCollation` 1:1 (separate copy, historical reasons like the rest of this class). */ + private fun declaredRecursiveCteCollation(base: QueryStatement, columns: List): (ColumnExpression) -> Collation? = + (base as? SelectStatement)?.let { declaredCollationForSelectOutput(it, columns) } ?: { null } + + private fun multisetFilter( + left: List>, + right: List>, + keepOnMatch: Boolean, + keyOf: (List) -> List + ): List> { + val remainingInRight = right.groupingBy(keyOf).eachCount().toMutableMap() + return left.filter { row -> + val key = keyOf(row) + val remaining = remainingInRight[key] ?: 0 + if (remaining > 0) { + remainingInRight[key] = remaining - 1 + keepOnMatch + } 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 + } + + /** Resolves ORDER BY columns of a combined UNION/INTERSECT/EXCEPT result purely by output column name (never qualified). */ + private fun evaluateOutputRow(columns: List, expression: Expression, row: List, txn: Long): Value { + val resolver = ColumnResolver { column -> + if (column.qualifier != null) throw DiskExecutorException("no such column: ${column.qualifier}.${column.name}") + val index = columns.indexOfFirst { it.equals(column.name, ignoreCase = true) } + if (index == -1) throw DiskExecutorException("no such column: ${column.name}") + row[index] + } + return Evaluator(resolver, subqueryRunnerFor(txn)).evaluate(expression) + } + + /** + * Pushes every CTE of [query] onto [activeCtes] (in definition order, so a later CTE may + * reference an earlier one, mirroring [kestrel.engine.SelectRunner.runWith]), runs [query]'s + * body, then pops them back off - a CTE is scoped to the single query that defines it, unlike a + * VIEW, which lives in the catalog. + */ + private fun runWith(query: WithQuery, txn: Long, outer: ColumnResolver?): SelectResult { + var pushed = 0 + try { + for (cte in query.ctes) { + activeCtes.add(buildCteScope(query.recursive, cte, txn)) + pushed++ + } + return runQuery(query.body, txn, outer) + } finally { + repeat(pushed) { activeCtes.removeAt(activeCtes.size - 1) } + } + } + + /** + * Materializes one CTE's rows. A non-recursive CTE just runs its query once. A recursive one + * (`WITH RECURSIVE name AS (base UNION [ALL] recursiveTerm)`, split by [splitRecursiveCte]) + * repeatedly re-runs [RecursiveCtePlan.recursiveTerm] with [CteScope.rows] set to only the + * *previous iteration's new rows* (the standard fixed-point working-set semantics, not the + * whole accumulated set - mirrors [kestrel.engine.SelectRunner.buildCteScope] exactly, including + * the [MAX_RECURSION_ITERATIONS] safety cap against a query that never reaches a fixed point). + */ + private fun buildCteScope(recursive: Boolean, cte: CommonTableExpression, txn: Long): CteScope { + val recursivePlan = if (recursive) splitRecursiveCte(cte) else null + if (recursivePlan == null) { + val result = runQuery(cte.query, txn, outer = null) + return CteScope(cte.name, cteSchema(cte, result), result.rows) + } + val (base, operator, recursiveTerm) = recursivePlan + val baseResult = runQuery(base, txn, outer = null) + val schema = cteSchema(cte, baseResult) + val scope = CteScope(cte.name, schema, baseResult.rows) + 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 DiskExecutorException( + "recursive common table expression '${cte.name}' did not terminate within $MAX_RECURSION_ITERATIONS iterations" + ) + } + scope.rows = working + var newRows = runQuery(recursiveTerm, txn, outer = 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 + } finally { + activeCtes.removeAt(activeCtes.size - 1) + } + return scope + } +