feat(sql): add RecursiveCtePlan and splitRecursiveCte to DiskExecutor
This commit is contained in:
1 parent
b67ac3386a
commit
cf9a017913
1 file changed
+162
@@ -2573,3 +2573,165 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
|
|||||||
return scope
|
return scope
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private data class RecursiveCtePlan(val base: QueryStatement, val operator: SetOperator, val recursiveTerm: QueryStatement)
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Recognizes `base UNION [ALL] recursiveTerm` where exactly one side names [cte] in its own
|
||||||
|
* FROM/JOIN - structurally, on the AST, without running anything - mirroring
|
||||||
|
* [kestrel.engine.SelectRunner.splitRecursiveCte] 1:1. Both sides referencing it, or neither,
|
||||||
|
* are rejected/treated as a plain (non-recursive) compound respectively, exactly like the
|
||||||
|
* in-memory engine.
|
||||||
|
*/
|
||||||
|
private fun splitRecursiveCte(cte: CommonTableExpression): RecursiveCtePlan? {
|
||||||
|
val query = cte.query
|
||||||
|
if (query !is SetOperationStatement) return null
|
||||||
|
if (query.operator != SetOperator.Union && query.operator != SetOperator.UnionAll) return null
|
||||||
|
val leftRefs = referencesName(query.left, cte.name)
|
||||||
|
val rightRefs = referencesName(query.right, cte.name)
|
||||||
|
return when {
|
||||||
|
leftRefs && rightRefs ->
|
||||||
|
throw DiskExecutorException(
|
||||||
|
"recursive common table expression '${cte.name}' must not reference itself on both sides of ${query.operator}"
|
||||||
|
)
|
||||||
|
rightRefs -> RecursiveCtePlan(query.left, query.operator, query.right)
|
||||||
|
leftRefs -> RecursiveCtePlan(query.right, query.operator, query.left)
|
||||||
|
else -> null
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun referencesName(query: QueryStatement, name: String): Boolean = when (query) {
|
||||||
|
is SelectStatement -> referencesName(query.from, name) || query.joins.any { referencesName(it.from, name) }
|
||||||
|
is SetOperationStatement -> referencesName(query.left, name) || referencesName(query.right, name)
|
||||||
|
is WithQuery -> query.ctes.none { it.name.equals(name, ignoreCase = true) } && referencesName(query.body, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun referencesName(item: FromItem, name: String): Boolean = when (item) {
|
||||||
|
is TableSource -> item.name.equals(name, ignoreCase = true)
|
||||||
|
is SubquerySource -> referencesName(item.select, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun cteSchema(cte: CommonTableExpression, result: SelectResult): TableSchema {
|
||||||
|
val labels = cte.columns ?: result.columns
|
||||||
|
if (labels.size != result.columns.size) {
|
||||||
|
throw DiskExecutorException(
|
||||||
|
"common table expression '${cte.name}' expects ${labels.size} columns but its query produces ${result.columns.size}"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
return derivedSchema(cte.name, result, labels)
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* A synthetic schema for a CTE's materialized rows: same idea as
|
||||||
|
* [kestrel.engine.SelectRunner.derivedSchema] - each column's type is inferred from the first
|
||||||
|
* non-NULL value seen in it (defaulting to TEXT when every row is NULL there), since a CTE has
|
||||||
|
* no backing [TableSchema] of its own to borrow one from.
|
||||||
|
*/
|
||||||
|
private fun derivedSchema(name: String, result: SelectResult, labels: List<String> = result.columns): TableSchema {
|
||||||
|
val columns = labels.mapIndexed { index, columnName ->
|
||||||
|
val type = result.rows.asSequence().map { it[index] }.firstOrNull { !it.isNull }?.type ?: ColumnType.Text
|
||||||
|
Column(columnName, type, notNull = false, primaryKey = false)
|
||||||
|
}
|
||||||
|
return TableSchema(name, columns)
|
||||||
|
}
|
||||||
|
|
||||||
|
private class CteScope(val name: String, val schema: TableSchema, var rows: List<List<Value>>)
|
||||||
|
|
||||||
|
/** One FROM/JOIN source's row supply: either a real disk table or a CTE's already-materialized rows. */
|
||||||
|
private sealed interface DiskFromSource {
|
||||||
|
val schema: TableSchema
|
||||||
|
fun scan(txn: Long): List<List<Value>>
|
||||||
|
}
|
||||||
|
|
||||||
|
private class DiskTableSource(val table: DiskTable) : DiskFromSource {
|
||||||
|
override val schema get() = table.schema
|
||||||
|
override fun scan(txn: Long): List<List<Value>> = table.scan(txn).map { it.second.values }
|
||||||
|
}
|
||||||
|
|
||||||
|
private class CteRowSource(override val schema: TableSchema, val rows: List<List<Value>>) : DiskFromSource {
|
||||||
|
override fun scan(txn: Long): List<List<Value>> = rows
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Resolves one FROM/JOIN [name] to its row supply: an active CTE shadows a same-named real
|
||||||
|
* table or view (checked innermost-first via [MutableList.lastOrNull], exactly like
|
||||||
|
* [kestrel.engine.SelectRunner.resolveNamedSource]'s own CTE lookup), then a temp table, then a
|
||||||
|
* real disk table, and only then a view - materialized fresh by running its defining query
|
||||||
|
* inside the same [txn], with [expandingViews] guarding against a view (directly or
|
||||||
|
* transitively) referencing itself, mirroring [kestrel.engine.SelectRunner.resolveNamedSource]
|
||||||
|
* exactly. A temp table's current rows are captured into a [CteRowSource] snapshot exactly like
|
||||||
|
* a CTE's or view's - a plain SELECT never mutates its own FROM source mid-query, so this is
|
||||||
|
* exactly as fresh as scanning a live [DiskTable] would be, and it is what gives a temp table
|
||||||
|
* full JOIN/subquery/aggregate participation for free instead of needing its own parallel query
|
||||||
|
* pipeline. A name that is neither falls through to [DiskCatalog.table], solely for its
|
||||||
|
* "no such table" [DiskCatalogException] - already confirmed absent from
|
||||||
|
* [DiskCatalog.temporaryTableOrNull]/[DiskCatalog.tableOrNull]/[DiskCatalog.viewOrNull], so this
|
||||||
|
* always throws, preserving the exact exception type callers relied on before views existed.
|
||||||
|
*/
|
||||||
|
private fun resolveFromSource(name: String, txn: Long): DiskFromSource {
|
||||||
|
activeCtes.lastOrNull { it.name.equals(name, ignoreCase = true) }?.let { return CteRowSource(it.schema, it.rows) }
|
||||||
|
catalog.temporaryTableOrNull(name)?.let { return CteRowSource(it.schema, it.rows.map { row -> row.values }) }
|
||||||
|
catalog.tableOrNull(name)?.let { return DiskTableSource(it) }
|
||||||
|
val view = catalog.viewOrNull(name) ?: return DiskTableSource(catalog.table(name))
|
||||||
|
val key = name.lowercase()
|
||||||
|
if (!expandingViews.add(key)) {
|
||||||
|
throw DiskExecutorException("recursive view definition: $name")
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
val result = runQuery(view, txn, outer = null)
|
||||||
|
return CteRowSource(derivedSchema(name, result), result.rows)
|
||||||
|
} finally {
|
||||||
|
expandingViews.remove(key)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Resolves one FROM/JOIN [item] to its row supply, mirroring
|
||||||
|
* [kestrel.engine.SelectRunner.resolveSource]/`buildSources`: a [TableSource] falls through to
|
||||||
|
* [resolveFromSource] by name (CTE/table/view, unchanged); a non-LATERAL [SubquerySource] (a
|
||||||
|
* derived table) runs its query eagerly, right here, inside the same [txn] - it cannot correlate
|
||||||
|
* to anything, so its rows are fixed for the whole statement, exactly like a materialized view. A
|
||||||
|
* LATERAL [SubquerySource] is different: its rows depend on whichever row of an earlier FROM/JOIN
|
||||||
|
* source it correlates with, which does not exist yet at this point - only its column *names* are
|
||||||
|
* needed here (via [lateralSchema], computed structurally from the AST without running anything)
|
||||||
|
* for offset bookkeeping; its real rows are produced later, per left row, in [executeLateralJoin].
|
||||||
|
*/
|
||||||
|
private fun resolveFromItem(item: FromItem, txn: Long): DiskFromSource = when (item) {
|
||||||
|
is TableSource -> resolveFromSource(item.name, txn)
|
||||||
|
is SubquerySource ->
|
||||||
|
if (item.lateral) {
|
||||||
|
CteRowSource(lateralSchema(item), emptyList())
|
||||||
|
} else {
|
||||||
|
val result = runQuery(item.select, txn, outer = null)
|
||||||
|
CteRowSource(derivedSchema(item.alias, result), result.rows)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun lateralSchema(item: SubquerySource): TableSchema {
|
||||||
|
val labels = structuralColumnNames(item.select)
|
||||||
|
val columns = labels.map { Column(it, ColumnType.Text, notNull = false, primaryKey = false) }
|
||||||
|
return TableSchema(item.alias, columns)
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The output column names of [query], derived purely from its AST (no execution, no catalog row
|
||||||
|
* data) - needed to give a LATERAL derived table a stable schema before any row of the source it
|
||||||
|
* correlates with is known. Mirrors [kestrel.engine.SelectRunner.structuralColumnNames] 1:1.
|
||||||
|
*/
|
||||||
|
private fun structuralColumnNames(query: QueryStatement, localCtes: Map<String, List<String>> = emptyMap()): List<String> =
|
||||||
|
when (query) {
|
||||||
|
is SelectStatement -> when (val projection = query.projection) {
|
||||||
|
is StarProjection ->
|
||||||
|
structuralSourceColumnNames(query.from, localCtes) +
|
||||||
|
query.joins.flatMap { structuralSourceColumnNames(it.from, localCtes) }
|
||||||
|
is ExpressionProjection -> projection.items.mapIndexed { index, item -> columnLabel(item, index) }
|
||||||
|
}
|
||||||
|
is SetOperationStatement -> structuralColumnNames(query.left, localCtes)
|
||||||
|
is WithQuery -> {
|
||||||
|
val extended = localCtes.toMutableMap()
|
||||||
|
for (cte in query.ctes) {
|
||||||
|
extended[cte.name.lowercase()] = cte.columns ?: structuralColumnNames(cte.query, extended)
|
||||||
|
}
|
||||||
|
structuralColumnNames(query.body, extended)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Reference in new issue
Block a user