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
    }