feat(sql): add declaredSetOperationCollation

2026-09-07 12:04:22 +02:00
1 parent e012659bde
commit b67ac3386a
@@ -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<String>): (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<String>): (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<String>): (ColumnExpression) -> Collation? =
+        (base as? SelectStatement)?.let { declaredCollationForSelectOutput(it, columns) } ?: { null }
+
+    private fun multisetFilter(
+        left: List<List<Value>>,
+        right: List<List<Value>>,
+        keepOnMatch: Boolean,
+        keyOf: (List<Value>) -> List<Value>
+    ): List<List<Value>> {
+        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<OrderKey>,
+        evaluate: (Expression, List<Value>) -> Value,
+        declaredCollationOf: (ColumnExpression) -> Collation? = { null }
+    ): Comparator<List<Value>> =
+        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<String>, expression: Expression, row: List<Value>, 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<List<Value>>()
+            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
+    }
+