@lobomfz/db 0.3.9 → 0.4.1

package/src/vector-index/nearest.ts

@@ -0,0 +1,68 @@
+import { sql, type SelectQueryBuilder } from "kysely";
+
+import type { ResolvedVectorIndex } from "./resolve.js";
+
+type IndexMeta = Pick<ResolvedVectorIndex, "anchor" | "anchorPkColumn" | "column" | "dim">;
+
+const registry = new Map<string, IndexMeta>();
+
+export function registerVectorIndex(index: ResolvedVectorIndex) {
+	registry.set(index.name, {
+		anchor: index.anchor,
+		anchorPkColumn: index.anchorPkColumn,
+		column: index.column,
+		dim: index.dim,
+	});
+}
+
+export function unregisterVectorIndex(name: string) {
+	registry.delete(name);
+}
+
+export type NearestOptions<IndexName extends string = string> = {
+	indexName: IndexName;
+	vector: Float32Array | number[];
+	k: number;
+};
+
+type WithDistance<O> = O & { distance: number };
+
+function vectorToBytes(v: Float32Array | number[]) {
+	const f = v instanceof Float32Array ? v : Float32Array.from(v);
+
+	return new Uint8Array(f.buffer, f.byteOffset, f.byteLength);
+}
+
+export function nearest<IndexName extends string>(options: NearestOptions<IndexName>) {
+	return <DB extends { [K in IndexName]: unknown }, TB extends keyof DB & string, O>(
+		qb: SelectQueryBuilder<DB, TB, O>,
+	): SelectQueryBuilder<DB, TB | (IndexName & keyof DB & string), WithDistance<O>> => {
+		const meta = registry.get(options.indexName);
+
+		if (!meta) {
+			throw new Error(
+				`Vector index "${options.indexName}" is not registered. Declare it in the schema's vectorIndexes and construct the Database before calling nearest.`,
+			);
+		}
+
+		const bytes = vectorToBytes(options.vector);
+		const indexName = options.indexName;
+		const column = meta.column;
+		const anchor = meta.anchor;
+		const anchorPk = meta.anchorPkColumn;
+
+		const anyQb = qb as unknown as {
+			innerJoin: (a: any, b: any, c: any) => any;
+			where: (expr: any) => any;
+			select: (expr: any) => any;
+			orderBy: (expr: any) => any;
+		};
+
+		return anyQb
+			.innerJoin(indexName, `${indexName}.rowid`, `${anchor}.${anchorPk}`)
+			.where(sql<boolean>`${sql.ref(`${indexName}.${column}`)} MATCH ${bytes}`)
+			.where(sql<boolean>`${sql.ref(`${indexName}.k`)} = ${options.k}`)
+			.select(sql<number>`${sql.ref(`${indexName}.distance`)}`.as("distance"))
+			.orderBy(`${indexName}.distance`);
+	};
+}

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 });
+}