feat(engine): extend SelectRunner with multiKeyComparator
This commit is contained in:
1 parent
48cb9dbaba
commit
eeaada39b9
1 file changed
+145
-8
@@ -246,8 +246,135 @@ class SelectRunner(private val catalog: Catalog) {
|
|||||||
} else {
|
} else {
|
||||||
!keepOnMatch
|
!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
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun evaluateOutputRow(columns: List<String>, expression: Expression, row: List<Value>): Value {
|
||||||
|
val resolver = kestrel.engine.ColumnResolver { column ->
|
||||||
|
if (column.qualifier != null) throw KestrelException("no such column: ${column.qualifier}.${column.name}")
|
||||||
|
val index = columns.indexOfFirst { it.equals(column.name, ignoreCase = true) }
|
||||||
|
if (index == -1) throw KestrelException("no such column: ${column.name}")
|
||||||
|
row[index]
|
||||||
|
}
|
||||||
|
return Evaluator(resolver, subqueryRunner).evaluate(expression)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun runWith(query: WithQuery, outer: ColumnResolver?): SelectResult {
|
||||||
|
var pushed = 0
|
||||||
|
try {
|
||||||
|
for (cte in query.ctes) {
|
||||||
|
activeCtes.add(buildCteScope(query.recursive, cte))
|
||||||
|
pushed++
|
||||||
|
}
|
||||||
|
return runQuery(query.body, outer)
|
||||||
|
} finally {
|
||||||
|
repeat(pushed) { activeCtes.removeAt(activeCtes.size - 1) }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun buildCteScope(recursive: Boolean, cte: CommonTableExpression): CteScope {
|
||||||
|
val recursivePlan = if (recursive) splitRecursiveCte(cte) else null
|
||||||
|
if (recursivePlan == null) {
|
||||||
|
val result = runQuery(cte.query, null)
|
||||||
|
return CteScope(cte.name, cteSchema(cte, result), result.rows.map { Row(it) })
|
||||||
|
}
|
||||||
|
val (base, operator, recursiveTerm) = recursivePlan
|
||||||
|
val baseResult = runQuery(base, null)
|
||||||
|
val schema = cteSchema(cte, baseResult)
|
||||||
|
val scope = CteScope(cte.name, schema, baseResult.rows.map { Row(it) })
|
||||||
|
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 KestrelException(
|
||||||
|
"recursive common table expression '${cte.name}' did not terminate within $MAX_RECURSION_ITERATIONS iterations"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
scope.rows = working.map { Row(it) }
|
||||||
|
var newRows = runQuery(recursiveTerm, 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.map { Row(it) }
|
||||||
|
} finally {
|
||||||
|
activeCtes.removeAt(activeCtes.size - 1)
|
||||||
|
}
|
||||||
|
return scope
|
||||||
|
}
|
||||||
|
|
||||||
|
private data class RecursiveCtePlan(val base: QueryStatement, val operator: SetOperator, val recursiveTerm: QueryStatement)
|
||||||
|
|
||||||
|
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 KestrelException("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 KestrelException(
|
||||||
|
"common table expression '${cte.name}' expects ${labels.size} columns but its query produces ${result.columns.size}"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
return derivedSchema(cte.name, result, labels)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun buildPlan(statement: SelectStatement, outer: ColumnResolver?): Plan {
|
||||||
val sources = buildSources(statement)
|
val sources = buildSources(statement)
|
||||||
val plan = Plan(statement, sources)
|
val plan = Plan(statement, sources, outer)
|
||||||
plan.validate()
|
plan.validate()
|
||||||
return plan
|
return plan
|
||||||
}
|
}
|
||||||
@@ -255,17 +382,27 @@ class SelectRunner(private val catalog: Catalog) {
|
|||||||
private fun buildSources(statement: SelectStatement): List<Source> {
|
private fun buildSources(statement: SelectStatement): List<Source> {
|
||||||
val sources = mutableListOf<Source>()
|
val sources = mutableListOf<Source>()
|
||||||
var offset = 0
|
var offset = 0
|
||||||
fun add(tableName: String, alias: String?) {
|
fun add(item: FromItem) {
|
||||||
val table = catalog.table(tableName)
|
// A LATERAL derived table's rows depend on the row of an earlier FROM item it
|
||||||
val chosen = alias ?: tableName
|
// correlates with, so they cannot be resolved once up front like every other source -
|
||||||
|
// only its column *names* are needed here (for offset bookkeeping and validating
|
||||||
|
// references to its alias elsewhere in the query), computed structurally from the
|
||||||
|
// subquery's own AST without running it. Real rows are produced per left row in
|
||||||
|
// [executeJoin].
|
||||||
|
val (schema, rows) = if (item is SubquerySource && item.lateral) {
|
||||||
|
lateralSchema(item) to emptyList()
|
||||||
|
} else {
|
||||||
|
resolveSource(item)
|
||||||
|
}
|
||||||
|
val chosen = sourceLabel(item)
|
||||||
if (sources.any { it.alias.equals(chosen, ignoreCase = true) }) {
|
if (sources.any { it.alias.equals(chosen, ignoreCase = true) }) {
|
||||||
throw KestrelException("duplicate table alias: $chosen")
|
throw KestrelException("duplicate table alias: $chosen")
|
||||||
}
|
}
|
||||||
sources.add(Source(chosen, table.schema.name, table.schema, table.rows, offset))
|
sources.add(Source(chosen, schema.name, schema, rows, offset))
|
||||||
offset += table.schema.columns.size
|
offset += schema.columns.size
|
||||||
}
|
}
|
||||||
add(statement.tableName, statement.tableAlias)
|
add(statement.from)
|
||||||
for (join in statement.joins) add(join.tableName, join.alias)
|
for (join in statement.joins) add(join.from)
|
||||||
return sources
|
return sources
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user