okengine 0.18.5 → 0.19.1

package/src/elements/store/search-lsh.ts

@@ -0,0 +1,214 @@
+/**
+ * Random-hyperplane LSH — K fixed planes generated once per field lifetime.
+ */
+
+import { createHash } from "node:crypto";
+import { LSH_DEFAULT_K } from "./search-errors.ts";
+
+/**
+ * Stable seed string for a field's hyperplanes.
+ *
+ * @param table - SQL table
+ * @param column - SQL column
+ * @param dims - Embedding dimensionality
+ * @param k - Number of planes
+ */
+export function hyperplaneSeed(table: string, column: string, dims: number, k: number): string {
+  return `${table}|${column}|${dims}|${k}`;
+}
+
+/**
+ * Generate K unit-ish random hyperplanes in `dims` dimensions from a seed.
+ * Deterministic across hosts for the same seed.
+ *
+ * @param seed - Stable seed string
+ * @param dims - Vector dimensionality
+ * @param k - Plane count (default {@link LSH_DEFAULT_K})
+ */
+export function generateHyperplanes(
+  seed: string,
+  dims: number,
+  k: number = LSH_DEFAULT_K,
+): Float32Array[] {
+  const planes: Float32Array[] = [];
+  let counter = 0;
+  for (let i = 0; i < k; i++) {
+    const plane = new Float32Array(dims);
+    let filled = 0;
+    while (filled < dims) {
+      const h = createHash("sha256").update(`${seed}\0${counter++}`).digest();
+      for (let b = 0; b + 3 < h.length && filled < dims; b += 4) {
+        const u = h.readUInt32BE(b) / 0xffff_ffff;
+        plane[filled++] = u * 2 - 1;
+      }
+    }
+    // L2 normalize
+    let norm = 0;
+    for (let j = 0; j < dims; j++) norm += plane[j]! * plane[j]!;
+    norm = Math.sqrt(norm) || 1;
+    for (let j = 0; j < dims; j++) plane[j]! /= norm;
+    planes.push(plane);
+  }
+  return planes;
+}
+
+/**
+ * Pack LSH bucket bits into a bigint (K ≤ 64).
+ *
+ * @param vector - Embedding
+ * @param planes - Fixed hyperplanes
+ */
+export function lshBucket(vector: readonly number[], planes: readonly Float32Array[]): bigint {
+  if (planes.length > 64) {
+    throw new Error(`lshBucket: K=${planes.length} exceeds bigint packing (max 64)`);
+  }
+  let bits = 0n;
+  for (let i = 0; i < planes.length; i++) {
+    const plane = planes[i]!;
+    if (vector.length !== plane.length) {
+      throw new Error(`lshBucket: vector length ${vector.length} !== plane dims ${plane.length}`);
+    }
+    let dot = 0;
+    for (let j = 0; j < plane.length; j++) {
+      dot += vector[j]! * plane[j]!;
+    }
+    if (dot >= 0) bits |= 1n << BigInt(i);
+  }
+  return bits;
+}
+
+/**
+ * Encode an LSH bucket for PostgreSQL `bigint` (signed int64).
+ * K=64 bit patterns with the high bit set exceed unsigned range when passed
+ * as a decimal string — reinterpret as two's-complement signed.
+ *
+ * @param bucket - Unsigned bit pack from {@link lshBucket}
+ */
+export function lshBucketToSql(bucket: bigint): string {
+  const masked = bucket & 0xffff_ffff_ffff_ffffn;
+  if (masked >= 0x8000_0000_0000_0000n) {
+    return (masked - 0x1_0000_0000_0000_0000n).toString();
+  }
+  return masked.toString();
+}
+
+/**
+ * Decode a PostgreSQL `bigint` LSH value back to an unsigned bit pack.
+ *
+ * @param value - Driver value (string | number | bigint)
+ */
+export function lshBucketFromSql(value: unknown): bigint {
+  const n = typeof value === "bigint" ? value : BigInt(String(value));
+  return n < 0n ? n + 0x1_0000_0000_0000_0000n : n;
+}
+
+/**
+ * Cosine similarity in [−1, 1].
+ *
+ * @param a - Vector a
+ * @param b - Vector b
+ */
+export function cosineSimilarity(a: readonly number[], b: readonly number[]): number {
+  if (a.length !== b.length) {
+    throw new Error(`cosineSimilarity: length mismatch ${a.length} vs ${b.length}`);
+  }
+  let dot = 0;
+  let na = 0;
+  let nb = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i]! * b[i]!;
+    na += a[i]! * a[i]!;
+    nb += b[i]! * b[i]!;
+  }
+  const denom = Math.sqrt(na) * Math.sqrt(nb);
+  return denom === 0 ? 0 : dot / denom;
+}
+
+/**
+ * Hamming distance between two K-bit LSH packs (K ≤ 64).
+ *
+ * Random-hyperplane (Charikar) SimHash ranks neighbors by this distance.
+ * Hamming-1 equality is not a viable candidate set at {@link LSH_DEFAULT_K}:
+ * true near-neighbors typically sit at distance ~5–16, not 0–1.
+ *
+ * @param a - Bucket a
+ * @param b - Bucket b
+ */
+export function hammingDistance(a: bigint, b: bigint): number {
+  let x = (a ^ b) & 0xffff_ffff_ffff_ffffn;
+  let count = 0;
+  while (x) {
+    x &= x - 1n;
+    count++;
+  }
+  return count;
+}
+
+/**
+ * Candidate buckets: exact match plus Hamming distance 1 flips.
+ *
+ * Kept for tight-collision tests and callers that want a 65-wide equality
+ * probe. Query-time retrieval uses {@link lshHammingSql} ranking instead —
+ * Hamming-1 alone misses almost all true neighbors at K=64.
+ *
+ * @param bucket - Query bucket
+ * @param k - Bit width
+ */
+export function neighborBuckets(bucket: bigint, k: number): bigint[] {
+  const out: bigint[] = [bucket];
+  for (let i = 0; i < k; i++) {
+    out.push(bucket ^ (1n << BigInt(i)));
+  }
+  return out;
+}
+
+/**
+ * Postgres 14+ Hamming distance between a stored signed-int64 LSH column
+ * and a bound query bucket (`?::bigint`). Used as `ORDER BY … LIMIT k`.
+ *
+ * @param columnSql - Quoted LSH column identifier
+ * @param bucketPlaceholder - Bound-parameter placeholder (`?` before `$n` rewrite)
+ */
+export function lshHammingSql(columnSql: string, bucketPlaceholder: string = "?"): string {
+  return `bit_count((${columnSql} # ${bucketPlaceholder}::bigint)::bit(64))`;
+}
+
+/**
+ * Serialize planes for `oke_search_planes.planes` (Float32 little-endian).
+ *
+ * @param planes - Hyperplanes
+ */
+export function serializePlanes(planes: readonly Float32Array[]): Buffer {
+  const dims = planes[0]?.length ?? 0;
+  const buf = Buffer.allocUnsafe(4 + planes.length * dims * 4);
+  buf.writeUInt32LE(dims, 0);
+  let offset = 4;
+  for (const plane of planes) {
+    for (let i = 0; i < dims; i++) {
+      buf.writeFloatLE(plane[i]!, offset);
+      offset += 4;
+    }
+  }
+  return buf;
+}
+
+/**
+ * Deserialize planes from storage.
+ *
+ * @param buf - Stored bytes
+ * @param k - Expected plane count
+ */
+export function deserializePlanes(buf: Buffer, k: number): Float32Array[] {
+  const dims = buf.readUInt32LE(0);
+  const planes: Float32Array[] = [];
+  let offset = 4;
+  for (let i = 0; i < k; i++) {
+    const plane = new Float32Array(dims);
+    for (let j = 0; j < dims; j++) {
+      plane[j] = buf.readFloatLE(offset);
+      offset += 4;
+    }
+    planes.push(plane);
+  }
+  return planes;
+}

package/src/elements/store/search-runtime.pglite.test.ts

@@ -0,0 +1,186 @@
+/**
+ * Hybrid search runtime — PGlite end-to-end for LSH candidate SQL.
+ *
+ * Captures the literal SQL `runSqlSearch` issues and checks K=64 Hamming-rank
+ * recovery on the same synthetic hash-bag G17 uses (40-row hand case).
+ */
+
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+import { createHash } from "node:crypto";
+import type { DeclaredColumn } from "../../manifest/types.ts";
+import { pgliteDriver } from "../../drivers/pglite.ts";
+import type { SqlConnection } from "../../drivers/types.ts";
+import {
+  embColumn,
+  ensureHyperplaneInserts,
+  lshColumn,
+  OKE_SEARCH_PLANES,
+  searchDdlForTable,
+} from "./search-ddl.ts";
+import { LSH_DEFAULT_K } from "./search-errors.ts";
+import {
+  cosineSimilarity,
+  deserializePlanes,
+  lshBucket,
+  lshBucketFromSql,
+  lshBucketToSql,
+} from "./search-lsh.ts";
+import { runSqlSearch, type SearchColumnMeta } from "./search-runtime.ts";
+
+const TABLE = "lsh_recall_docs";
+const DIMS = 32;
+const TOPICS = [
+  "refund policy shipping delay",
+  "password reset two factor auth",
+  "invoice payment stripe webhook",
+  "search ranking bm25 hybrid",
+  "postgres gin btree index plan",
+  "tenant isolation row level security",
+  "durable flow journal resume",
+  "live query fanout subscriber",
+] as const;
+
+const COLS: Record<string, DeclaredColumn> = {
+  id: { type: "text", sqlName: "id", primaryKey: true },
+  title: { type: "text", sqlName: "title", searchable: { weight: 2 } },
+  body: {
+    type: "text",
+    sqlName: "body",
+    searchable: { weight: 1 },
+    embed: { dims: DIMS, model: "mock" },
+  },
+};
+
+const SEARCH_COLS: SearchColumnMeta[] = [
+  { key: "title", sqlName: "title", weight: 2 },
+  { key: "body", sqlName: "body", weight: 1, embed: { dims: DIMS, model: "mock" } },
+];
+
+let conn: SqlConnection;
+
+beforeAll(async () => {
+  conn = await pgliteDriver.connect({ url: "memory://lsh-recall-runtime", role: "primary" });
+}, 20_000);
+
+afterAll(async () => {
+  await conn.close();
+});
+
+/**
+ * G17 synthetic embedding — deterministic SHA-256 hash bag, L2-normalized.
+ *
+ * @param text - Source text
+ */
+function synthEmbed(text: string): number[] {
+  const v = Array.from({ length: DIMS }, () => 0);
+  const tokens = text.toLowerCase().split(/\W+/).filter(Boolean);
+  for (const t of tokens) {
+    const h = createHash("sha256").update(t).digest();
+    for (let i = 0; i < DIMS; i++) {
+      v[i]! += (h[i % h.length]! / 255) * 2 - 1;
+    }
+  }
+  let norm = 0;
+  for (const x of v) norm += x * x;
+  norm = Math.sqrt(norm) || 1;
+  return v.map((x) => x / norm);
+}
+
+describe("runSqlSearch LSH candidate path (PGlite)", () => {
+  test("issued SQL ranks by bit_count Hamming, not Hamming-1 ANY(); recovers exact top-10", async () => {
+    await conn.exec(`DROP TABLE IF EXISTS ${TABLE} CASCADE`);
+    await conn.exec(`
+      CREATE TABLE ${TABLE} (
+        id text PRIMARY KEY,
+        title text NOT NULL,
+        body text NOT NULL
+      )
+    `);
+    for (const stmt of searchDdlForTable(TABLE, COLS)) {
+      await conn.exec(stmt);
+    }
+    for (const ins of ensureHyperplaneInserts(TABLE, COLS)) {
+      await conn.exec(ins.sql, ins.params);
+    }
+
+    const planesRow = await conn.query(
+      `SELECT k, planes FROM ${OKE_SEARCH_PLANES} WHERE table_name = ? AND column_name = ?`,
+      [TABLE, "body"],
+    );
+    const prow = planesRow[0]!;
+    const k = Number(prow["k"] ?? LSH_DEFAULT_K);
+    const planes = deserializePlanes(Buffer.from(prow["planes"] as Buffer), k);
+
+    const vectors: number[][] = [];
+    const n = 40;
+    for (let i = 0; i < n; i++) {
+      const topic = TOPICS[i % TOPICS.length]!;
+      const id = `d${String(i).padStart(8, "0")}`;
+      const title = `${topic} item ${i}`;
+      const body = `${topic}. Document ${i} discusses ${topic} with enough tokens for BM25.`;
+      const vec = synthEmbed(body);
+      const bucket = lshBucket(vec, planes);
+      vectors.push(vec);
+      await conn.exec(
+        `INSERT INTO ${TABLE} (id, title, body, ${embColumn("body")}, ${lshColumn("body")})
+         VALUES (?, ?, ?, ?, ?)`,
+        [id, title, body, `{${vec.join(",")}}`, lshBucketToSql(bucket)],
+      );
+    }
+
+    // Write/query identity for one stored row vs re-projection of the same body.
+    const stored = await conn.query(`SELECT ${lshColumn("body")} AS b FROM ${TABLE} WHERE id = ?`, [
+      "d00000000",
+    ]);
+    const body0 = `${TOPICS[0]}. Document 0 discusses ${TOPICS[0]} with enough tokens for BM25.`;
+    const queryBucket = lshBucket(synthEmbed(body0), planes);
+    expect(lshBucketFromSql(stored[0]?.["b"])).toBe(queryBucket);
+
+    const captured: Array<{ sql: string; params: readonly unknown[] }> = [];
+    const tapped: SqlConnection = {
+      driverId: conn.driverId,
+      role: conn.role,
+      query: async (sql, params = []) => {
+        captured.push({ sql, params });
+        return conn.query(sql, params);
+      },
+      exec: (sql, params) => conn.exec(sql, params),
+      close: () => conn.close(),
+    };
+
+    const qVec = synthEmbed(TOPICS[0]!);
+    const result = await runSqlSearch({
+      conn: tapped,
+      table: TABLE,
+      columns: SEARCH_COLS,
+      pkSqlName: "id",
+      options: { query: `zxqwv ${TOPICS[0]}`, limit: 10 },
+      embedQuery: async () => qVec,
+    });
+
+    const candidate = captured.find((c) => /UNION/i.test(c.sql) && /bit_count/i.test(c.sql));
+    expect(candidate).toBeDefined();
+    expect(candidate!.sql).toContain("bit_count");
+    expect(candidate!.sql).toContain("ORDER BY");
+    expect(candidate!.sql).not.toMatch(/=\s*ANY\s*\(/i);
+    // Bound query bucket is the signed int64 pack of the query vector.
+    expect(candidate!.params).toContain(lshBucketToSql(lshBucket(qVec, planes)));
+
+    const scored: Array<{ i: number; s: number }> = [];
+    for (let i = 0; i < vectors.length; i++) {
+      scored.push({ i, s: cosineSimilarity(qVec, vectors[i]!) });
+    }
+    scored.sort((a, b) => b.s - a.s);
+    const exact = new Set(scored.slice(0, 10).map((x) => `d${String(x.i).padStart(8, "0")}`));
+    const approx = result.data.map((r) => String(r["id"]));
+    let hits = 0;
+    for (const id of approx.slice(0, 10)) {
+      if (exact.has(id)) hits += 1;
+    }
+    // Hamming-1 equality recovered 0/10 on this corpus. Hamming-rank UNION
+    // plus in-process cosine/RRF must recover a real majority of exact neighbors.
+    expect(result.data).toHaveLength(10);
+    expect(hits / 10).toBeGreaterThanOrEqual(0.5);
+    expect(result.meta.engine).toContain("lsh");
+  }, 30_000);
+});