diff --git a/src/main/kotlin/kestrel/engine/SelectRunner.kt b/src/main/kotlin/kestrel/engine/SelectRunner.kt index 39ead65..d947784 100644 --- a/src/main/kotlin/kestrel/engine/SelectRunner.kt +++ b/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 = 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> = 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 { + 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 { + 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) = 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): (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): (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): (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 + } val sources = buildSources(statement) val plan = Plan(statement, sources) plan.validate()