@lobomfz/db 0.4.0 → 0.5.0

package/src/vector-index/resolve.ts

@@ -0,0 +1,78 @@
+import type { VectorIndexEntry, VectorMetric } from "../types.js";
+
+export type VectorAnchorInfo = {
+	columns: Set<string>;
+	pkColumn: string | null;
+	pkIsInteger: boolean;
+	vectorColumns: Map<string, number>;
+};
+
+export type ResolvedVectorIndex = {
+	name: string;
+	anchor: string;
+	anchorPkColumn: string;
+	column: string;
+	dim: number;
+	metric: VectorMetric;
+};
+
+const DEFAULT_METRIC: VectorMetric = "cosine";
+
+const VALID_METRICS = new Set<VectorMetric>(["cosine", "l2", "l1"]);
+
+export function resolveVectorIndex(
+	name: string,
+	entry: VectorIndexEntry<any>,
+	anchors: Map<string, VectorAnchorInfo>,
+) {
+	const anchor = anchors.get(entry.anchor);
+
+	if (!anchor) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" is not declared in schema`,
+		);
+	}
+
+	if (!anchor.pkColumn) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" has no primary key; anchor PK must be an integer`,
+		);
+	}
+
+	if (!anchor.pkIsInteger) {
+		throw new Error(
+			`VectorIndex "${name}": anchor table "${entry.anchor}" primary key "${anchor.pkColumn}" must be an integer`,
+		);
+	}
+
+	if (!anchor.columns.has(entry.column)) {
+		throw new Error(
+			`VectorIndex "${name}": column "${entry.column}" does not exist on anchor "${entry.anchor}"`,
+		);
+	}
+
+	const dim = anchor.vectorColumns.get(entry.column);
+
+	if (dim === undefined) {
+		throw new Error(
+			`VectorIndex "${name}": column "${entry.anchor}.${entry.column}" is not a vector column`,
+		);
+	}
+
+	const metric = entry.metric ?? DEFAULT_METRIC;
+
+	if (!VALID_METRICS.has(metric)) {
+		throw new Error(
+			`VectorIndex "${name}": unknown metric "${metric}", expected one of cosine, l2, l1`,
+		);
+	}
+
+	return {
+		name,
+		anchor: entry.anchor,
+		anchorPkColumn: anchor.pkColumn,
+		column: entry.column,
+		dim,
+		metric,
+	};
+}

package/src/vector-index/sql.ts

@@ -0,0 +1,76 @@
+import type { ResolvedVectorIndex } from "./resolve.js";
+
+export const VEC_TRIGGER_PREFIX = "__vec_sync_";
+
+export type VecTriggerEvent = "AI" | "AU" | "AD";
+
+export function vecTriggerName(index: string, anchor: string, event: VecTriggerEvent) {
+	return `${VEC_TRIGGER_PREFIX}${index}_${anchor}_${event}`;
+}
+
+export type GeneratedTrigger = { name: string; sql: string };
+
+const VALID_IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]*$/;
+
+export function generateCreateVecSql(index: ResolvedVectorIndex) {
+	if (!VALID_IDENTIFIER.test(index.column)) {
+		throw new Error(
+			`VectorIndex "${index.name}": column "${index.column}" must be a simple identifier (vec0 does not accept quoted column names)`,
+		);
+	}
+
+	return (
+		`CREATE VIRTUAL TABLE "${index.name}" USING vec0(` +
+		`rowid INTEGER PRIMARY KEY, ` +
+		`${index.column} float[${index.dim}] distance_metric=${index.metric}` +
+		`)`
+	);
+}
+
+export function generateVectorTriggers(index: ResolvedVectorIndex): GeneratedTrigger[] {
+	const anchor = index.anchor;
+	const pk = index.anchorPkColumn;
+	const col = index.column;
+	const vec = index.name;
+
+	const aiName = vecTriggerName(index.name, anchor, "AI");
+	const auName = vecTriggerName(index.name, anchor, "AU");
+	const adName = vecTriggerName(index.name, anchor, "AD");
+
+	return [
+		{
+			name: aiName,
+			sql:
+				`CREATE TRIGGER "${aiName}" AFTER INSERT ON "${anchor}" ` +
+				`WHEN NEW."${col}" IS NOT NULL BEGIN ` +
+				`INSERT INTO "${vec}"(rowid, "${col}") VALUES (NEW."${pk}", NEW."${col}"); ` +
+				`END`,
+		},
+		{
+			name: auName,
+			sql:
+				`CREATE TRIGGER "${auName}" AFTER UPDATE OF "${col}" ON "${anchor}" BEGIN ` +
+				`DELETE FROM "${vec}" WHERE rowid = OLD."${pk}" AND NEW."${col}" IS NULL; ` +
+				`INSERT INTO "${vec}"(rowid, "${col}") ` +
+				`SELECT NEW."${pk}", NEW."${col}" WHERE OLD."${col}" IS NULL AND NEW."${col}" IS NOT NULL; ` +
+				`UPDATE "${vec}" SET "${col}" = NEW."${col}" ` +
+				`WHERE rowid = NEW."${pk}" AND OLD."${col}" IS NOT NULL AND NEW."${col}" IS NOT NULL; ` +
+				`END`,
+		},
+		{
+			name: adName,
+			sql:
+				`CREATE TRIGGER "${adName}" AFTER DELETE ON "${anchor}" BEGIN ` +
+				`DELETE FROM "${vec}" WHERE rowid = OLD."${pk}"; ` +
+				`END`,
+		},
+	];
+}
+
+export function generateRebuildVectorSql(index: ResolvedVectorIndex) {
+	return (
+		`INSERT INTO "${index.name}"(rowid, "${index.column}") ` +
+		`SELECT "${index.anchorPkColumn}", "${index.column}" FROM "${index.anchor}" ` +
+		`WHERE "${index.column}" IS NOT NULL`
+	);
+}

package/src/vector.ts

@@ -0,0 +1,13 @@
+import { type } from "arktype";
+
+export function vector(dim: number) {
+	if (!Number.isInteger(dim) || dim <= 0) {
+		throw new Error(`vector(dim) requires a positive integer dim, got ${dim}`);
+	}
+
+	return type
+		.instanceOf(Float32Array)
+		.or(type("number[]"))
+		.pipe((v) => (v instanceof Float32Array ? v : Float32Array.from(v)))
+		.configure({ _generated: "vector", _vectorDim: dim });
+}

package/src/write-validation-plugin.ts

@@ -30,6 +30,7 @@ type TableWriteSchema = {
 	columnSet: Set<string>;
 	optionalNonNullColumns: Set<string>;
 	insertNullColumns: Set<string>;
+	vectorColumns: Map<string, number>;
 };
 
 type InsertRow = {
@@ -57,6 +58,16 @@ export class WriteValidationPlugin implements KyselyPlugin {
 			const structureProps = (schema as any).structure?.props as StructureProp[] | undefined;
 			const columns = structureProps?.map((prop) => prop.key) ?? [];
 
+			const vectorColumns = new Map<string, number>();
+
+			for (const prop of structureProps ?? []) {
+				const meta = this.extractFieldMeta(prop);
+
+				if (meta?._generated === "vector" && typeof meta._vectorDim === "number") {
+					vectorColumns.set(prop.key, meta._vectorDim);
+				}
+			}
+
 			this.schemas.set(table, {
 				schema,
 				columns,
@@ -71,14 +82,91 @@ export class WriteValidationPlugin implements KyselyPlugin {
 						?.filter((prop) => this.acceptsNull(prop.value) && prop.inner.default === undefined)
 						.map((prop) => prop.key) ?? [],
 				),
+				vectorColumns,
 			});
 		}
 	}
 
+	private validateAndCoerceVector(
+		table: string,
+		column: string,
+		dim: number,
+		value: unknown,
+	): Uint8Array {
+		const vec =
+			value instanceof Float32Array
+				? value
+				: Array.isArray(value)
+					? Float32Array.from(value as number[])
+					: null;
+
+		if (!vec) {
+			throw new ValidationError(
+				table,
+				`expected Float32Array or number[] for vector column, got ${typeof value}`,
+				column,
+			);
+		}
+
+		if (vec.length !== dim) {
+			throw new ValidationError(
+				table,
+				`vector dim mismatch: expected ${dim}, received ${vec.length}`,
+				column,
+			);
+		}
+
+		for (let i = 0; i < vec.length; i++) {
+			if (!Number.isFinite(vec[i]!)) {
+				throw new ValidationError(
+					table,
+					`non-finite vector element at index ${i}: ${vec[i]}`,
+					column,
+				);
+			}
+		}
+
+		return new Uint8Array(vec.buffer, vec.byteOffset, vec.byteLength);
+	}
+
+	private coerceVectorsInRow(
+		table: string,
+		tableSchema: TableWriteSchema,
+		values: Record<string, unknown>,
+	) {
+		for (const [column, dim] of tableSchema.vectorColumns) {
+			if (!Object.prototype.hasOwnProperty.call(values, column)) {
+				continue;
+			}
+
+			const value = values[column];
+
+			if (value === null || value === undefined) {
+				continue;
+			}
+
+			values[column] = this.validateAndCoerceVector(table, column, dim, value);
+		}
+	}
+
 	private acceptsNull(field: StructureProp["value"]) {
 		return field.branches.some((branch) => branch.unit === null || branch.domain === "null");
 	}
 
+	private extractFieldMeta(prop: StructureProp) {
+		const rootMeta = (prop.value as any).meta as
+			| { _generated?: string; _vectorDim?: number }
+			| undefined;
+		const branchMeta = prop.value.branches.find(
+			(b) => (b as any).meta && Object.keys((b as any).meta).length > 0,
+		);
+		const fromBranch = (branchMeta as any)?.meta as
+			| { _generated?: string; _vectorDim?: number }
+			| undefined;
+
+		return { ...fromBranch, ...rootMeta };
+	}
+
 	private getTableFromNode(node: RootOperationNode) {
 		switch (node.kind) {
 			case "InsertQueryNode":
@@ -219,6 +307,8 @@ export class WriteValidationPlugin implements KyselyPlugin {
 			this.prepareInsertValue(tableSchema, schemaLiteralValues),
 		);
 
+		this.coerceVectorsInRow(table, tableSchema, morphedSchemaValues);
+
 		return {
 			values: { ...nonSchemaLiteralValues, ...morphedSchemaValues },
 			passthrough: row.passthrough,
@@ -364,6 +454,8 @@ export class WriteValidationPlugin implements KyselyPlugin {
 						this.stripNullOptionalFields(tableSchema, literalValues),
 					);
 
+		this.coerceVectorsInRow(table, tableSchema, morphed);
+
 		return updates.flatMap((update) => {
 			if (!ColumnNode.is(update.column) || !ValueNode.is(update.value)) {
 				return [update];