@quereus/quereus 0.1.0 → 0.2.0

package/src/runtime/emit/project.ts

@@ -5,6 +5,7 @@ import { type Row } from '../../common/types.js';
 import { type OutputValue } from '../../common/types.js';
 import type { EmissionContext } from '../emission-context.js';
 import { buildRowDescriptor } from '../../util/row-descriptor.js';
+import { withAsyncRowContext, withRowContextGenerator } from '../context-helpers.js';
 
 export function emitProject(plan: ProjectNode, ctx: EmissionContext): Instruction {
 	const sourceInstruction = emitPlanNode(plan.source, ctx);
@@ -12,23 +13,28 @@ export function emitProject(plan: ProjectNode, ctx: EmissionContext): Instructio
 		return emitCallFromPlan(projection.node, ctx);
 	});
 
-	// Create row descriptor for source attributes
+	// Row descriptors
 	const sourceRowDescriptor = buildRowDescriptor(plan.source.getAttributes());
+	const outputRowDescriptor = buildRowDescriptor(plan.getAttributes());
 
 	async function* run(rctx: RuntimeContext, source: AsyncIterable<Row>, ...projectionFunctions: Array<(ctx: RuntimeContext) => OutputValue>): AsyncIterable<Row> {
 		for await (const sourceRow of source) {
-			// Set up context for this row using row descriptor
-			rctx.context.set(sourceRowDescriptor, () => sourceRow);
+			// Evaluate projections using the source row context
+			const outputs = await withAsyncRowContext(rctx, sourceRowDescriptor, () => sourceRow, async () => {
+				return Promise.all(projectionFunctions.map(fn => fn(rctx)));
+			});
 
-			try {
-				const outputs = projectionFunctions.map(func => func(rctx));
-				const resolved = await Promise.all(outputs);
-				// Assume we have ensured that these are all scalar values
-				yield resolved as Row;
-			} finally {
-				// Clean up context for this row
-				rctx.context.delete(sourceRowDescriptor);
-			}
+			// Push the output row descriptor for downstream consumers
+			yield* withRowContextGenerator(
+				rctx,
+				outputRowDescriptor,
+				(async function* () {
+					yield outputs as Row;
+				})(),
+				async function* (row) {
+					yield row;
+				}
+			);
 		}
 	}
 

package/src/runtime/emit/recursive-cte.ts

@@ -11,6 +11,7 @@ import { DEFAULT_TUNING } from '../../planner/optimizer-tuning.js';
 import { quereusError } from '../../common/errors.js';
 import { StatusCode } from '../../common/types.js';
 import { buildRowDescriptor } from '../../util/row-descriptor.js';
+import { withRowContext } from '../context-helpers.js';
 
 const log = createLogger('runtime:emit:recursive-cte');
 
@@ -42,12 +43,7 @@ export function emitRecursiveCTE(plan: RecursiveCTENode, ctx: EmissionContext):
 
 			if (shouldYield) {
 				// Yield immediately (streaming)
-				rctx.context.set(rowDescriptor, () => row);
-				try {
-					yield row;
-				} finally {
-					rctx.context.delete(rowDescriptor);
-				}
+				yield withRowContext(rctx, rowDescriptor, () => row, () => row);
 
 				// Add to delta for recursive processing (deep copy to avoid reference issues)
 				deltaRows.push([...row] as Row);
@@ -75,12 +71,7 @@ export function emitRecursiveCTE(plan: RecursiveCTENode, ctx: EmissionContext):
 
 					if (shouldYield) {
 						// Stream the row immediately
-						rctx.context.set(rowDescriptor, () => row);
-						try {
-							yield row;
-						} finally {
-							rctx.context.delete(rowDescriptor);
-						}
+						yield withRowContext(rctx, rowDescriptor, () => row, () => row);
 
 						// Add to next iteration's delta (deep copy to avoid reference issues)
 						newDeltaRows.push([...row] as Row);

package/src/runtime/emit/returning.ts

@@ -4,6 +4,7 @@ import type { Row } from '../../common/types.js';
 import type { EmissionContext } from '../emission-context.js';
 import { emitPlanNode, emitCallFromPlan } from '../emitters.js';
 import { buildRowDescriptor } from '../../util/row-descriptor.js';
+import { withRowContextGenerator } from '../context-helpers.js';
 
 export function emitReturning(plan: ReturningNode, ctx: EmissionContext): Instruction {
 	// Use the executor's attributes to build the row descriptor
@@ -21,20 +22,12 @@ export function emitReturning(plan: ReturningNode, ctx: EmissionContext): Instru
 		...projectionCallbacks: Array<(ctx: RuntimeContext) => any>
 	): AsyncIterable<Row> {
 		// Project the results from the executor rows
-		for await (const sourceRow of executorRows) {
-			// Set up the source context - the row is already in the correct flat OLD/NEW format
-			rctx.context.set(sourceRowDescriptor, () => sourceRow);
-
-			try {
-				// Evaluate projection expressions in the context of this row
-				const outputs = projectionCallbacks.map(func => func(rctx));
-				const resolved = await Promise.all(outputs);
-				yield resolved as Row;
-			} finally {
-				// Clean up source context
-				rctx.context.delete(sourceRowDescriptor);
-			}
-		}
+		yield* withRowContextGenerator(rctx, sourceRowDescriptor, executorRows, async function* (_sourceRow) {
+			// Evaluate projection expressions in the context of this row
+			const outputs = projectionCallbacks.map(func => func(rctx));
+			const resolved = await Promise.all(outputs);
+			yield resolved as Row;
+		});
 	}
 
 	// Emit the executor (now always produces rows)

package/src/runtime/emit/scan.ts

@@ -1,5 +1,5 @@
 import { StatusCode, type Row } from "../../common/types.js";
-import type { TableScanNode } from "../../planner/nodes/scan.js";
+import { SeqScanNode, IndexScanNode, IndexSeekNode } from "../../planner/nodes/table-access-nodes.js";
 import { QuereusError } from "../../common/errors.js";
 import type { VirtualTable } from "../../vtab/table.js";
 import type { BaseModuleConfig } from "../../vtab/module.js";
@@ -8,32 +8,38 @@ import type { EmissionContext } from "../emission-context.js";
 import { createValidatedInstruction } from "../emitters.js";
 import { disconnectVTable } from "../utils.js";
 import { buildRowDescriptor } from "../../util/row-descriptor.js";
+import { createRowSlot } from "../context-helpers.js";
 
-export function emitSeqScan(plan: TableScanNode, ctx: EmissionContext): Instruction {
-	const tableSchema = plan.source.tableSchema;
+/**
+ * Emits instructions for physical table access nodes (SeqScan, IndexScan, IndexSeek)
+ */
+export function emitSeqScan(plan: SeqScanNode | IndexScanNode | IndexSeekNode, ctx: EmissionContext): Instruction {
+	// Handle physical access nodes
+	const source = plan.source;
+	const schema = source.tableSchema;
 
 	// Create row descriptor mapping attribute IDs to column indices
 	const rowDescriptor = buildRowDescriptor(plan.getAttributes());
 
 	// Look up the virtual table module during emission and record the dependency
-	const moduleInfo = ctx.getVtabModule(tableSchema.vtabModuleName);
+	const moduleInfo = ctx.getVtabModule(schema.vtabModuleName);
 	if (!moduleInfo) {
-		throw new QuereusError(`Virtual table module '${tableSchema.vtabModuleName}' not found`, StatusCode.ERROR);
+		throw new QuereusError(`Virtual table module '${schema.vtabModuleName}' not found`, StatusCode.ERROR);
 	}
 
 	// Capture the module info key for runtime retrieval
-	const moduleKey = `vtab_module:${tableSchema.vtabModuleName}`;
+	const moduleKey = `vtab_module:${schema.vtabModuleName}`;
 
 	async function* run(runtimeCtx: RuntimeContext): AsyncIterable<Row> {
 		// Use the captured module info instead of doing a fresh lookup
 		const capturedModuleInfo = ctx.getCapturedSchemaObject<{ module: any, auxData?: unknown }>(moduleKey);
 		if (!capturedModuleInfo) {
-			throw new QuereusError(`Virtual table module '${tableSchema.vtabModuleName}' was not captured during emission`, StatusCode.INTERNAL);
+			throw new QuereusError(`Virtual table module '${schema.vtabModuleName}' was not captured during emission`, StatusCode.INTERNAL);
 		}
 
 		const module = capturedModuleInfo.module;
 		if (typeof module.xConnect !== 'function') {
-			throw new QuereusError(`Virtual table module '${tableSchema.vtabModuleName}' does not implement xConnect`, StatusCode.MISUSE);
+			throw new QuereusError(`Virtual table module '${schema.vtabModuleName}' does not implement xConnect`, StatusCode.MISUSE);
 		}
 
 		let vtabInstance: VirtualTable;
@@ -42,38 +48,34 @@ export function emitSeqScan(plan: TableScanNode, ctx: EmissionContext): Instruct
 			vtabInstance = module.xConnect(
 				runtimeCtx.db,
 				capturedModuleInfo.auxData,
-				tableSchema.vtabModuleName,
-				tableSchema.schemaName,
-				tableSchema.name,
+				schema.vtabModuleName,
+				schema.schemaName,
+				schema.name,
 				options
 			);
 		} catch (e: any) {
 			const message = e instanceof Error ? e.message : String(e);
-			throw new QuereusError(`Module '${tableSchema.vtabModuleName}' xConnect failed for table '${tableSchema.name}': ${message}`, e instanceof QuereusError ? e.code : StatusCode.ERROR, e instanceof Error ? e : undefined);
+			throw new QuereusError(`Module '${schema.vtabModuleName}' xConnect failed for table '${schema.name}': ${message}`, e instanceof QuereusError ? e.code : StatusCode.ERROR, e instanceof Error ? e : undefined);
 		}
 
 		if (typeof vtabInstance.xQuery !== 'function') {
 			// Fallback or error if xQuery is not available. For now, throwing an error.
 			// Later, we could implement the xOpen/xFilter/xNext loop here as a fallback.
-			throw new QuereusError(`Virtual table '${tableSchema.name}' does not support xQuery.`, StatusCode.UNSUPPORTED);
+			throw new QuereusError(`Virtual table '${schema.name}' does not support xQuery.`, StatusCode.UNSUPPORTED);
 		}
 
+		const rowSlot = createRowSlot(runtimeCtx, rowDescriptor);
 		try {
-			// Put cursor row into context using row descriptor
-			let row: Row;
-			runtimeCtx.context.set(rowDescriptor, () => row);
-
 			const asyncRowIterable = vtabInstance.xQuery(plan.filterInfo);
-			for await (row of asyncRowIterable) {
+			for await (const row of asyncRowIterable) {
+				rowSlot.set(row);
 				yield row;
 			}
-
-			// Remove cursor row from context
-			runtimeCtx.context.delete(rowDescriptor);
 		} catch (e: any) {
 			const message = e instanceof Error ? e.message : String(e);
-			throw new QuereusError(`Error during xQuery on table '${tableSchema.name}': ${message}`, e instanceof QuereusError ? e.code : StatusCode.ERROR, e instanceof Error ? e : undefined);
+			throw new QuereusError(`Error during xQuery on table '${schema.name}': ${message}`, e instanceof QuereusError ? e.code : StatusCode.ERROR, e instanceof Error ? e : undefined);
 		} finally {
+			rowSlot.close();
 			// Properly disconnect the VirtualTable instance
 			await disconnectVTable(runtimeCtx, vtabInstance);
 		}
@@ -83,6 +85,6 @@ export function emitSeqScan(plan: TableScanNode, ctx: EmissionContext): Instruct
 		[],
 		run,
 		ctx,
-		`scan(${plan.source.tableSchema.name})`
+		`${plan.nodeType}(${schema.name})`
 	);
 }

package/src/runtime/emit/sort.ts

@@ -5,6 +5,7 @@ import { type SqlValue, type Row, type MaybePromise } from '../../common/types.j
 import type { EmissionContext } from '../emission-context.js';
 import { createOrderByComparatorFast, resolveCollation } from '../../util/comparison.js';
 import { buildRowDescriptor } from '../../util/row-descriptor.js';
+import { withAsyncRowContext } from '../context-helpers.js';
 
 export function emitSort(plan: SortNode, ctx: EmissionContext): Instruction {
 	const sourceInstruction = emitPlanNode(plan.source, ctx);
@@ -31,21 +32,17 @@ export function emitSort(plan: SortNode, ctx: EmissionContext): Instruction {
 		const rowsWithKeys: Array<{ row: Row; keys: SqlValue[] }> = [];
 
 		for await (const sourceRow of source) {
-			// Set up context for this row using row descriptor
-			ctx.context.set(sourceRowDescriptor, () => sourceRow);
-
-			try {
+			const keys = await withAsyncRowContext(ctx, sourceRowDescriptor, () => sourceRow, async () => {
 				// Evaluate sort key expressions
-				const keys: SqlValue[] = [];
+				const keyValues: SqlValue[] = [];
 				for (const keyFunc of sortKeyFunctions) {
-					keys.push(await keyFunc(ctx));
+					const result = keyFunc(ctx);
+					keyValues.push(await Promise.resolve(result));
 				}
+				return keyValues;
+			});
 
-				rowsWithKeys.push({ row: sourceRow, keys });
-			} finally {
-				// Clean up context for this row
-				ctx.context.delete(sourceRowDescriptor);
-			}
+			rowsWithKeys.push({ row: sourceRow, keys });
 		}
 
 		// Sort the collected rows using pre-created optimized comparators

package/src/runtime/emit/subquery.ts

@@ -1,16 +1,17 @@
-import type { Instruction, RuntimeContext } from '../types.js';
+import type { Instruction, InstructionRun, RuntimeContext } from '../types.js';
 import type { InNode, ScalarSubqueryNode, ExistsNode } from '../../planner/nodes/subquery.js';
 import { emitPlanNode } from '../emitters.js';
 import type { SqlValue, Row } from '../../common/types.js';
-import { compareSqlValuesFast, BINARY_COLLATION } from '../../util/comparison.js';
+import { compareSqlValuesFast, resolveCollation } from '../../util/comparison.js';
 import type { EmissionContext } from '../emission-context.js';
 import { QuereusError } from '../../common/errors.js';
 import { StatusCode } from '../../common/types.js';
 import { BTree } from 'inheritree';
+import { ConstantNode } from '../../planner/nodes/plan-node.js';
 
 export function emitScalarSubquery(plan: ScalarSubqueryNode, ctx: EmissionContext): Instruction {
 
-	async function run(ctx: RuntimeContext, input: AsyncIterable<Row>): Promise<SqlValue> {
+	async function run(rctx: RuntimeContext, input: AsyncIterable<Row>): Promise<SqlValue> {
 		let result: SqlValue = null;
 		let seen = false;
 
@@ -18,8 +19,11 @@ export function emitScalarSubquery(plan: ScalarSubqueryNode, ctx: EmissionContex
 			if (seen) {
 				throw new QuereusError('Scalar subquery returned more than one row', StatusCode.ERROR, undefined, plan.expression.loc?.start.line, plan.expression.loc?.start.column);
 			}
+			if (row.length > 1) {
+				throw new QuereusError('Subquery should return at most one column', StatusCode.ERROR);
+			}
+			result = row.length === 0 ? null : row[0];
 			seen = true;
-			result = row[0];
 		}
 
 		return result;
@@ -30,25 +34,25 @@ export function emitScalarSubquery(plan: ScalarSubqueryNode, ctx: EmissionContex
 	return {
 		params: [innerInstruction],
 		run,
-		note: 'scalar subquery'
+		note: 'SCALAR_SUBQUERY'
 	};
 }
 
 export function emitIn(plan: InNode, ctx: EmissionContext): Instruction {
+	// Extract collation from the condition expression
+	const conditionType = plan.condition.getType();
+	const collationName = conditionType.collationName || 'BINARY';
+	const collation = resolveCollation(collationName);
+
 	if (plan.source) {
 		// IN subquery: expr IN (SELECT ...)
-		async function runSubquery(ctx: RuntimeContext, input: AsyncIterable<Row>, condition: SqlValue): Promise<SqlValue> {
+		// Use streaming approach - check each row as we read it, return early on match
+		async function runSubqueryStreaming(_rctx: RuntimeContext, input: AsyncIterable<Row>, condition: SqlValue): Promise<SqlValue> {
 			// If condition is NULL, result is NULL
 			if (condition === null) {
 				return null;
 			}
 
-			// Build BTree of all values from subquery
-			const tree = new BTree<SqlValue, SqlValue>(
-				(val: SqlValue) => val,
-				(a: SqlValue, b: SqlValue) => compareSqlValuesFast(a, b, BINARY_COLLATION)
-			);
-
 			let hasNull = false;
 			for await (const row of input) {
 				if (row.length > 0) {
@@ -57,16 +61,13 @@ export function emitIn(plan: InNode, ctx: EmissionContext): Instruction {
 						hasNull = true;
 						continue;
 					}
-					tree.insert(rowValue);
+					// Check for match immediately - no need to materialize
+					if (compareSqlValuesFast(condition, rowValue, collation) === 0) {
+						return 1; // Found a match
+					}
 				}
 			}
 
-			// Check if condition exists in tree
-			const value = tree.get(condition);
-			if (value !== undefined) {
-				return 1; // Found a match
-			}
-
 			// No match found - if any value was NULL, result is NULL
 			return hasNull ? null : 0;
 		}
@@ -76,57 +77,116 @@ export function emitIn(plan: InNode, ctx: EmissionContext): Instruction {
 
 		return {
 			params: [sourceInstruction, conditionExpr],
-			run: runSubquery as any,
+			run: runSubqueryStreaming as any,
 			note: `IN (subquery)`
 		};
 	} else if (plan.values) {
 		// IN value list: expr IN (value1, value2, ...)
-		function runValues(ctx: RuntimeContext, condition: SqlValue, ...values: SqlValue[]): SqlValue {
-			// If condition is NULL, result is NULL
-			if (condition === null) {
-				return null;
-			}
 
-			// Build BTree of all values
+		// Check if all values are truly constant (can be evaluated at emit time)
+		const allConstant = plan.values.every(val => val.physical.constant);
+
+		if (allConstant) {
+			// Pre-build BTree at emit time for constant values
 			const tree = new BTree<SqlValue, SqlValue>(
 				(val: SqlValue) => val,
-				(a: SqlValue, b: SqlValue) => compareSqlValuesFast(a, b, BINARY_COLLATION)
+				(a: SqlValue, b: SqlValue) => compareSqlValuesFast(a, b, collation)
 			);
-
 			let hasNull = false;
-			for (const value of values) {
-				if (value === null) {
-					hasNull = true;
-					continue;
+
+			function innerConstantRun(_rctx: RuntimeContext, condition: SqlValue): SqlValue {
+				// If condition is NULL, result is NULL
+				if (condition === null) {
+					return null;
+				}
+
+				// Check if condition exists in pre-built tree
+				const path = tree.find(condition);
+				if (path.on) {
+					return 1; // Found a match
 				}
-				tree.insert(value);
+
+				// No match found - if any value was NULL, result is NULL
+				return hasNull ? null : 0;
 			}
 
-			// Check if condition exists in tree
-			const path = tree.find(condition);
-			if (path.on) {
-				return 1; // Found a match
+			const values = plan.values.map(val => (val as unknown as ConstantNode).getValue());
+
+			let runFunc: InstructionRun;
+
+			if (values.some(val => val instanceof Promise)) {
+				// Must resolve promises at runtime
+				runFunc = async (rctx: RuntimeContext, condition: SqlValue): Promise<SqlValue> => {
+					const resolved = await Promise.all(values);
+
+					for (const value of resolved) {
+						if (value === null) {
+							hasNull = true;
+							continue;
+						}
+						tree.insert(value as SqlValue);
+					}
+
+					return innerConstantRun(rctx, condition);
+				}
+			} else {
+				for (const value of values) {
+					if (value === null) {
+						hasNull = true;
+						continue;
+					}
+					tree.insert(value as SqlValue);
+				}
+			runFunc = innerConstantRun;
 			}
 
-			// No match found - if any value was NULL, result is NULL
-			return hasNull ? null : 0;
-		}
+			const conditionExpr = emitPlanNode(plan.condition, ctx);
+
+			return {
+				params: [conditionExpr],
+				run: runFunc as InstructionRun,
+				note: `IN (${plan.values.length} constant values)`
+			};
+		} else {
+			// Some values are expressions - build tree at runtime
+			function runDynamicValues(_rctx: RuntimeContext, condition: SqlValue, ...values: SqlValue[]): SqlValue {
+				// If condition is NULL, result is NULL
+				if (condition === null) {
+					return null;
+				}
 
-		const conditionExpr = emitPlanNode(plan.condition, ctx);
-		const valueExprs = plan.values.map(val => emitPlanNode(val, ctx));
+				// Linear scan is optimal since we're only doing one lookup per execution
+				let hasNull = false;
+				for (const value of values) {
+					if (value === null) {
+						hasNull = true;
+						continue;
+					}
+					if (compareSqlValuesFast(condition, value, collation) === 0) {
+						return 1; // Found a match
+					}
+				}
 
-		return {
-			params: [conditionExpr, ...valueExprs],
-			run: runValues as any,
-			note: `IN (${plan.values.length} values)`
-		};
+				// No match found - if any value was NULL, result is NULL
+				return hasNull ? null : 0;
+			}
+
+			const conditionExpr = emitPlanNode(plan.condition, ctx);
+			const valueExprs = plan.values.map(val => emitPlanNode(val, ctx));
+
+			return {
+				params: [conditionExpr, ...valueExprs],
+				run: runDynamicValues as any,
+				note: `IN (${plan.values.length} dynamic values)`
+			};
+		}
 	} else {
 		throw new QuereusError('IN node must have either source or values', StatusCode.INTERNAL);
 	}
 }
 
 export function emitExists(plan: ExistsNode, ctx: EmissionContext): Instruction {
-	async function run(ctx: RuntimeContext, input: AsyncIterable<Row>): Promise<SqlValue> {
+	async function run(_rctx: RuntimeContext, input: AsyncIterable<Row>): Promise<SqlValue> {
 		for await (const _row of input) {
 			return 1; // First row => TRUE
 		}

package/src/runtime/emit/table-valued-function.ts

@@ -7,6 +7,7 @@ import { isTableValuedFunctionSchema } from '../../schema/function.js';
 import type { EmissionContext } from '../emission-context.js';
 import type { TableFunctionCallNode } from '../../planner/nodes/table-function-call.js';
 import { buildRowDescriptor } from '../../util/row-descriptor.js';
+import { withRowContextGenerator } from '../context-helpers.js';
 
 export function emitTableValuedFunctionCall(plan: TableFunctionCallNode, ctx: EmissionContext): Instruction {
 	const functionName = plan.functionName.toLowerCase();
@@ -49,16 +50,9 @@ export function emitTableValuedFunctionCall(plan: TableFunctionCallNode, ctx: Em
 			// Handle both direct AsyncIterable and Promise<AsyncIterable>
 			const iterable = result instanceof Promise ? await result : result;
 
-			for await (const row of iterable) {
-				// Set up row context for column references
-				innerCtx.context.set(rowDescriptor, () => row);
-				try {
-					yield row;
-				} finally {
-					// Clean up row context
-					innerCtx.context.delete(rowDescriptor);
-				}
-			}
+			yield* withRowContextGenerator(innerCtx, rowDescriptor, iterable, async function* (row) {
+				yield row;
+			});
 		} catch (error: any) {
 			throw new QuereusError(`Table-valued function ${functionName} failed: ${error.message}`, StatusCode.ERROR, error);
 		}
@@ -92,16 +86,9 @@ export function emitTableValuedFunctionCall(plan: TableFunctionCallNode, ctx: Em
 			// Handle both direct AsyncIterable and Promise<AsyncIterable>
 			const iterable = result instanceof Promise ? await result : result;
 
-			for await (const row of iterable) {
-				// Set up row context for column references
-				innerCtx.context.set(rowDescriptor, () => row);
-				try {
-					yield row;
-				} finally {
-					// Clean up row context
-					innerCtx.context.delete(rowDescriptor);
-				}
-			}
+			yield* withRowContextGenerator(innerCtx, rowDescriptor, iterable, async function* (row) {
+				yield row;
+			});
 		} catch (error: any) {
 			throw new QuereusError(`Table-valued function ${functionName} failed: ${error.message}`, StatusCode.ERROR, error);
 		}

package/src/runtime/emit/update.ts

@@ -5,6 +5,7 @@ import { QuereusError } from '../../common/errors.js';
 import { StatusCode, type SqlValue, type Row } from '../../common/types.js';
 import type { EmissionContext } from '../emission-context.js';
 import { buildRowDescriptor, composeOldNewRow } from '../../util/row-descriptor.js';
+import { withRowContextGenerator } from '../context-helpers.js';
 
 export function emitUpdate(plan: UpdateNode, ctx: EmissionContext): Instruction {
 	const tableSchema = plan.table.tableSchema;
@@ -28,39 +29,31 @@ export function emitUpdate(plan: UpdateNode, ctx: EmissionContext): Instruction
 	);
 
 	async function* run(rctx: RuntimeContext, sourceRowsIterable: AsyncIterable<Row>, ...assignmentEvaluators: Array<(ctx: RuntimeContext) => SqlValue>): AsyncIterable<Row> {
-		for await (const sourceRow of sourceRowsIterable) {
-			// Set up row context for evaluating assignment expressions
-			rctx.context.set(sourceRowDescriptor, () => sourceRow);
-
-			try {
-				// Evaluate assignment expressions in the context of this row
-				const assignmentValues: SqlValue[] = [];
-				for (const evaluator of assignmentEvaluators) {
-					const value = evaluator(rctx) as SqlValue;
-					assignmentValues.push(value);
-				}
+		yield* withRowContextGenerator(rctx, sourceRowDescriptor, sourceRowsIterable, async function* (sourceRow) {
+			// Evaluate assignment expressions in the context of this row
+			const assignmentValues: SqlValue[] = [];
+			for (const evaluator of assignmentEvaluators) {
+				const value = evaluator(rctx) as SqlValue;
+				assignmentValues.push(value);
+			}
 
-				// Create a new row with updated values
-				const updatedRow = [...sourceRow]; // Copy the original row
+			// Create a new row with updated values
+			const updatedRow = [...sourceRow]; // Copy the original row
 
-				// Apply assignment values to the row
-				for (let i = 0; i < assignmentValues.length; i++) {
-					const targetColIdx = assignmentTargetIndices[i];
-					updatedRow[targetColIdx] = assignmentValues[i];
-				}
+			// Apply assignment values to the row
+			for (let i = 0; i < assignmentValues.length; i++) {
+				const targetColIdx = assignmentTargetIndices[i];
+				updatedRow[targetColIdx] = assignmentValues[i];
+			}
 
-				// Create flat row with OLD (source) and NEW (updated) values for constraint checking
-				const flatRow = composeOldNewRow(sourceRow, updatedRow, tableSchema.columns.length);
+			// Create flat row with OLD (source) and NEW (updated) values for constraint checking
+			const flatRow = composeOldNewRow(sourceRow, updatedRow, tableSchema.columns.length);
 
-				// Yield the flat row for constraint checking
-				// NOTE: UpdateNode only transforms rows - it does NOT execute the actual update
-				// The UpdateExecutorNode is responsible for calling vtab.xUpdate
-				yield flatRow;
-			} finally {
-				// Clean up row context
-				rctx.context.delete(sourceRowDescriptor);
-			}
-		}
+			// Yield the flat row for constraint checking
+			// NOTE: UpdateNode only transforms rows - it does NOT execute the actual update
+			// The UpdateExecutorNode is responsible for calling vtab.xUpdate
+			yield flatRow;
+		});
 	}
 
 	const sourceInstruction = emitPlanNode(plan.source, ctx);