feat(engine): extend SelectRunner with explain and resolveMergeSource

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

@@ -73,7 +73,179 @@ class SelectRunner(private val catalog: Catalog) {
         is WithQuery -> runWith(query, outer)
     }
 
-    private fun buildPlan(statement: SelectStatement): Plan {
+    fun explain(query: QueryStatement): List<String> = when (query) {
+        is SelectStatement -> buildPlan(query, null).describe()
+        is SetOperationStatement -> explainSetOperation(query)
+        is WithQuery -> explainWith(query)
+    }
+
+    /**
+     * Resolves a single [FromItem] to its schema and row snapshot - the same lookup a `FROM`/`JOIN`
+     * source uses ([resolveSource]), exposed for [Executor.executeMerge]'s `USING` clause so a MERGE
+     * source can be either a real table or a derived-table subquery without duplicating view/CTE/
+     * subquery resolution.
+     */
+    fun resolveMergeSource(item: FromItem): Pair<TableSchema, List<Row>> = resolveSource(item)
+
+    /** The name a MERGE `ON`/`WHEN` condition would qualify a column from [item] with - its alias if given, its bare table name otherwise (a subquery source always has one). */
+    fun mergeSourceLabel(item: FromItem): String = sourceLabel(item)
+
+    private fun explainSetOperation(statement: SetOperationStatement): List<String> {
+        val header = setOperatorKeyword(statement.operator)
+        val left = explain(statement.left).map { "  -> $it" }
+        val right = explain(statement.right).map { "  -> $it" }
+        var lines = listOf(header) + left + right
+        statement.orderBy?.let { lines = listOf("Sort") + lines.map { line -> "  $line" } }
+        statement.offset?.let { lines = listOf("Offset ${statement.offset}") + lines.map { line -> "  $line" } }
+        statement.limit?.let { lines = listOf("Limit ${statement.limit}") + lines.map { line -> "  $line" } }
+        return lines
+    }
+
+    private fun explainWith(query: WithQuery): List<String> {
+        var pushed = 0
+        try {
+            for (cte in query.ctes) {
+                activeCtes.add(buildCteScope(query.recursive, cte))
+                pushed++
+            }
+            val cteLines = query.ctes.flatMap { cte ->
+                listOf("CTE ${cte.name}") + explain(cte.query).map { "  -> $it" }
+            }
+            return cteLines + explain(query.body)
+        } finally {
+            repeat(pushed) { activeCtes.removeAt(activeCtes.size - 1) }
+        }
+    }
+
+    private fun setOperatorKeyword(operator: SetOperator): String = when (operator) {
+        SetOperator.Union -> "Union"
+        SetOperator.UnionAll -> "Union All"
+        SetOperator.Intersect -> "Intersect"
+        SetOperator.IntersectAll -> "Intersect All"
+        SetOperator.Except -> "Except"
+        SetOperator.ExceptAll -> "Except All"
+    }
+
+    private fun runSetOperation(statement: SetOperationStatement, outer: ColumnResolver?): SelectResult {
+        val left = runQuery(statement.left, outer)
+        val right = runQuery(statement.right, outer)
+        if (left.columns.size != right.columns.size) {
+            throw KestrelException(
+                "${setOperatorKeyword(statement.operator)} queries must have the same number of columns " +
+                    "but got ${left.columns.size} and ${right.columns.size}"
+            )
+        }
+        val declaredCollationOf = declaredSetOperationCollation(statement, left.columns)
+        val outputExpressions = left.columns.map { ColumnExpression(it) }
+        fun foldedKey(row: List<Value>) = collationFoldedKey(row, outputExpressions, declaredCollationOf)
+        var rows = when (statement.operator) {
+            SetOperator.UnionAll -> left.rows + right.rows
+            SetOperator.Union -> (left.rows + right.rows).distinctBy(::foldedKey)
+            SetOperator.Intersect -> {
+                val rightKeys = right.rows.mapTo(HashSet(), ::foldedKey)
+                left.rows.filter { foldedKey(it) in rightKeys }.distinctBy(::foldedKey)
+            }
+            SetOperator.Except -> {
+                val rightKeys = right.rows.mapTo(HashSet(), ::foldedKey)
+                left.rows.filter { foldedKey(it) !in rightKeys }.distinctBy(::foldedKey)
+            }
+            SetOperator.IntersectAll -> multisetFilter(left.rows, right.rows, keepOnMatch = true, keyOf = ::foldedKey)
+            SetOperator.ExceptAll -> multisetFilter(left.rows, right.rows, keepOnMatch = false, keyOf = ::foldedKey)
+        }
+        val orderKeys = statement.orderBy?.keys.orEmpty()
+        if (orderKeys.isNotEmpty()) {
+            rows = rows.sortedWith(
+                multiKeyComparator(
+                    orderKeys,
+                    evaluate = { expression, row -> evaluateOutputRow(left.columns, expression, row) },
+                    declaredCollationOf = declaredCollationOf
+                )
+            )
+        }
+        statement.offset?.let { if (it < 0) throw KestrelException("OFFSET must not be negative") }
+        statement.limit?.let { if (it < 0) throw KestrelException("LIMIT must not be negative") }
+        if (statement.offset != null || statement.limit != null) {
+            rows = rows.drop(statement.offset ?: 0)
+            statement.limit?.let { rows = rows.take(it) }
+        }
+        return SelectResult(left.columns, rows)
+    }
+
+    /**
+     * The declared `COLLATE` of [statement]'s combined output column named by an ORDER BY key
+     * expression (an unqualified [ColumnExpression] - the only kind [evaluateOutputRow] resolves
+     * for a set operation's own ORDER BY, since it rejects any qualifier), traced back to a real
+     * table column in [statement]'s LEFT branch only - the branch [SelectResult.columns] itself
+     * always takes column identity from (`runSetOperation` returns `SelectResult(left.columns,
+     * ...)`, never consulting the right side's names). Resolves only when that branch is a plain
+     * `SELECT ... FROM realTable [JOIN realTable ...]` (no CTE, view, or derived table anywhere in
+     * its FROM clause) referencing the column bare - going further would mean inheriting collation
+     * through a view/CTE/derived-table boundary, the still-open scope-cut already left by #65/#66
+     * for a plain SELECT's own WHERE/ORDER BY over such a source (see kestrel-stan.md), not
+     * extended here.
+     */
+    private fun declaredSetOperationCollation(statement: SetOperationStatement, columns: List<String>): (ColumnExpression) -> Collation? =
+        (statement.left as? SelectStatement)?.let { declaredCollationForSelectOutput(it, columns) } ?: { null }
+
+    /**
+     * The same "plain `SELECT ... FROM realTable [JOIN realTable ...]`, no CTE/view/derived table"
+     * structural resolution [declaredSetOperationCollation] does for a set operation's LEFT branch,
+     * factored out so [declaredRecursiveCteCollation] can apply the identical rule to a recursive
+     * CTE's own anchor (base) term - the branch that owns [CteScope.schema]/column identity there,
+     * exactly like the LEFT branch owns [SelectResult.columns] for a plain set operation.
+     */
+    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 [declaredSetOperationCollation] but resolves against a recursive CTE's anchor (`base`) term. */
+    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
+            }
         val sources = buildSources(statement)
         val plan = Plan(statement, sources)
         plan.validate()