feat(sql): add executeMerge and executeMergeInner to DiskExecutor

This commit is contained in:
dualfroz committed 2026-09-07 17:29:14 +02:00
1 parent aee0fc649f
commit 0ee81dcf82
1 file changed
+158
@@ -3990,3 +3990,161 @@ class DiskExecutor(private val catalog: DiskCatalog) : SequenceAccess {
}
}
private fun executeMerge(statement: MergeStatement): StatementResult {
val table = catalog.table(statement.targetTable)
val schema = table.schema
return withStatementTransaction { txn ->
withTriggerAtomicity(schema.name, txn, null) {
executeMergeInner(statement, table, schema, txn)
}
}
}
/**
* Ported 1:1 from [kestrel.engine.Executor.executeMergeInner] - see that method's own doc for
* why [MergeStatement.onCondition] is matched against [originalTargetRows] (a snapshot of the
* target table taken once, before any mutation) rather than the live table, and for the "MERGE
* command cannot affect a target row a second time" guarantee. The one difference: "already
* affected" and "this pair's target row" are tracked by [RecordId] here rather than [Row]
* object identity (`===`) - a disk row is decoded fresh from its page on every read, so no two
* reads of "the same" row are ever the same object, but each row in [originalTargetRows] keeps
* the one [RecordId] it had at snapshot time for the rest of this statement (this MERGE is the
* only writer touching [table] during its own execution, and every [RecordId] this MERGE itself
* mutates is used at most once, guarded by [affectedRids] below, so it is never looked up again
* afterward under a since-changed identity).
*/
private fun executeMergeInner(statement: MergeStatement, table: DiskTable, schema: TableSchema, txn: Long): StatementResult {
val targetLabel = statement.targetAlias ?: statement.targetTable
val (sourceSchema, sourceRows) = resolveMergeSource(statement.source, txn)
val sourceLabel = sourceLabel(statement.source)
if (targetLabel.equals(sourceLabel, ignoreCase = true)) {
throw DiskExecutorException("MERGE target and source cannot share the same name or alias: $targetLabel")
}
checkMergeColumns(statement, schema, targetLabel, sourceSchema, sourceLabel)
val subqueries = subqueryRunnerFor(txn)
val originalTargetRows = table.scan(txn)
val affectedRids = mutableSetOf<RecordId>()
val returnedRows = mutableListOf<Pair<Row, String>>()
var inserted = 0
var updated = 0
var deleted = 0
for (sourceRow in sourceRows) {
val matchedTargets = originalTargetRows.filter { (_, targetRow) ->
val resolver = mergeResolver(schema, targetLabel, targetRow, sourceSchema, sourceLabel, sourceRow)
isTrue(Evaluator(resolver, subqueries).evaluate(statement.onCondition))
}
if (matchedTargets.isNotEmpty()) {
for ((targetRid, targetRow) in matchedTargets) {
val resolver = mergeResolver(schema, targetLabel, targetRow, sourceSchema, sourceLabel, sourceRow)
val clause = statement.whenClauses.firstOrNull { candidate ->
(candidate is MergeMatchedUpdate || candidate is MergeMatchedDelete) &&
(candidate.condition == null || isTrue(Evaluator(resolver, subqueries).evaluate(candidate.condition!!)))
} ?: continue
if (!affectedRids.add(targetRid)) {
throw DiskExecutorException("MERGE command cannot affect a target row a second time")
}
when (clause) {
is MergeMatchedUpdate -> {
val updatedRow = applyMergeUpdate(schema, table, txn, targetRid, targetRow, clause, resolver, subqueries)
returnedRows.add(updatedRow to "UPDATE")
updated++
}
is MergeMatchedDelete -> {
val deletedRow = applyMergeDelete(schema, table, txn, targetRid, targetRow)
returnedRows.add(deletedRow to "DELETE")
deleted++
}
else -> error("unreachable: filtered to MATCHED clauses above")
}
}
continue
}
val resolver = notMatchedResolver(sourceSchema, sourceLabel, sourceRow, targetLabel)
val clause = statement.whenClauses.firstOrNull { candidate ->
candidate is MergeNotMatchedInsert &&
(candidate.condition == null || isTrue(Evaluator(resolver, subqueries).evaluate(candidate.condition!!)))
} as? MergeNotMatchedInsert ?: continue
val insertedRow = applyMergeInsert(schema, table, txn, clause, resolver, subqueries)
returnedRows.add(insertedRow to "INSERT")
inserted++
}
val returning = statement.returning?.let { evaluateMergeReturning(schema, returnedRows, it, subqueries) }
return MergeResult(inserted, updated, deleted, returning)
}
/** Ported 1:1 from [kestrel.engine.Executor.evaluateMergeReturning] - see that method's own doc. */
private fun evaluateMergeReturning(
schema: TableSchema,
actionedRows: List<Pair<Row, String>>,
projection: Projection,
subqueries: SubqueryRunner
): SelectResult {
val columns = when (projection) {
is StarProjection -> schema.columns.map { it.name }
is ExpressionProjection -> projection.items.mapIndexed { index, item ->
checkReturningColumns(item.expression, schema)
columnLabel(item, index)
}
}
val resultRows = actionedRows.map { (row, action) ->
when (projection) {
is StarProjection -> row.values.toList()
is ExpressionProjection -> {
val evaluator = Evaluator(rowResolver(schema, row), subqueries, mergeAction = action)
projection.items.map { evaluator.evaluate(it.expression) }
}
}
}
return SelectResult(columns, resultRows)
}
private fun applyMergeUpdate(
schema: TableSchema,
table: DiskTable,
txn: Long,
targetRid: RecordId,
targetRow: Row,
clause: MergeMatchedUpdate,
resolver: ColumnResolver,
subqueries: SubqueryRunner
): Row {
val assignmentEvaluator = Evaluator(resolver, subqueries, this)
val values = targetRow.values.toMutableList()
for (assignment in clause.assignments) {
val index = schema.columnIndex(assignment.columnName)
?: throw DiskExecutorException("no such column: ${assignment.columnName}")
val column = schema.columns[index]
rejectGeneratedAssignment(column, explicitValueGiven = assignment.value != null)
values[index] = if (assignment.value != null) {
coerceForColumn(assignmentEvaluator.evaluate(assignment.value), column)
} else {
column.default?.let { coerceForColumn(assignmentEvaluator.evaluate(it), column) } ?: Value.NullValue
}
}
applyGeneratedColumns(schema, values, subqueries)
val updatedRow = Row(values)
fireRowTriggers(schema, TriggerTiming.Before, TriggerEvent.Update, listOf(targetRow to updatedRow), txn)
DiskForeignKeyEnforcer.validateRow(catalog, txn, schema, updatedRow.values, uncommittedSiblings = emptyList())
enforceChecks(schema, updatedRow, subqueries)
for (group in schema.enforcedUniqueKeyGroups()) {
val columnIndices = group.map { schema.columnIndex(it)!! }
val oldKey = columnIndices.map { targetRow.values[it] }
val newKey = columnIndices.map { updatedRow.values[it] }
DiskForeignKeyEnforcer.cascadeUpdate(catalog, txn, schema.name, group, listOf(oldKey to newKey))
}
table.update(txn, targetRid, updatedRow)
fireRowTriggers(schema, TriggerTiming.After, TriggerEvent.Update, listOf(targetRow to updatedRow), txn)
return updatedRow
}
/** Returns [targetRow] itself - the pre-deletion image, matching [applyMergeDelete]'s in-memory counterpart. */
private fun applyMergeDelete(schema: TableSchema, table: DiskTable, txn: Long, targetRid: RecordId, targetRow: Row): Row {
fireRowTriggers(schema, TriggerTiming.Before, TriggerEvent.Delete, listOf(targetRow to null), txn)
DiskForeignKeyEnforcer.cascadeDelete(catalog, txn, schema.name, listOf(targetRow))
table.delete(txn, targetRid)
fireRowTriggers(schema, TriggerTiming.After, TriggerEvent.Delete, listOf(targetRow to null), txn)
return targetRow
}