@kernl-sdk/pg 0.1.11 → 0.1.12

package/src/pgvector/index.ts

@@ -0,0 +1,7 @@
+/**
+ * pgvector search index.
+ */
+
+export { PGSearchIndex, type PGSearchIndexConfig } from "./search";
+export { PGIndexHandle } from "./handle";
+export type { PGIndexConfig, PGFieldBinding } from "./types";

package/src/pgvector/search.ts

@@ -0,0 +1,330 @@
+import type { Pool } from "pg";
+import { CursorPage, type CursorPageResponse } from "@kernl-sdk/shared";
+import { KERNL_SCHEMA_NAME } from "@kernl-sdk/storage";
+import type {
+  SearchIndex,
+  IndexHandle,
+  NewIndexParams,
+  ListIndexesParams,
+  IndexSummary,
+  IndexStats,
+  UnknownDocument,
+  SearchCapabilities,
+} from "@kernl-sdk/retrieval";
+
+import { PGIndexHandle } from "./handle";
+import { FIELD_TYPE, SIMILARITY } from "./sql";
+import type { PGIndexConfig, PGFieldBinding } from "./types";
+
+const META_TABLE = "search_indexes";
+
+export interface PGSearchIndexConfig {
+  pool: Pool;
+  ensureInit?: () => Promise<void>;
+}
+
+/**
+ * pgvector-backed SearchIndex implementation.
+ */
+export class PGSearchIndex implements SearchIndex<PGIndexConfig> {
+  readonly id = "pgvector";
+
+  private pool: Pool;
+  private userInit: () => Promise<void>;
+  private configs = new Map<string, PGIndexConfig>();
+  private ready = false;
+
+  constructor(config: PGSearchIndexConfig) {
+    this.pool = config.pool;
+    this.userInit = config.ensureInit ?? (() => Promise.resolve());
+  }
+
+  /**
+   * Create a new index table.
+   *
+   * @param params.id - Table name
+   * @param params.schema - Field definitions (one field must have pk: true)
+   * @param params.providerOptions.schema - Postgres schema (default: "public")
+   */
+  async createIndex(params: NewIndexParams): Promise<void> {
+    await this.ensureInit();
+
+    const schemaName = (params.providerOptions?.schema as string) ?? "public";
+
+    // find primary key field
+    const pkEntry = Object.entries(params.schema).find(([, f]) => f.pk);
+    if (!pkEntry) {
+      throw new Error("schema must have a field with pk: true");
+    }
+    const pkey = pkEntry[0];
+
+    const columns: string[] = [];
+    const vectorFields: Array<{
+      name: string;
+      dimensions: number;
+      similarity?: string;
+    }> = [];
+
+    for (const [name, field] of Object.entries(params.schema)) {
+      const colDef = `"${name}" ${FIELD_TYPE.encode(field)}${field.pk ? " PRIMARY KEY" : ""}`;
+      columns.push(colDef);
+
+      if (field.type === "vector") {
+        vectorFields.push({
+          name,
+          dimensions: field.dimensions,
+          similarity: field.similarity,
+        });
+      }
+    }
+
+    // create table
+    await this.pool.query(`
+      CREATE TABLE "${schemaName}"."${params.id}" (
+        ${columns.join(",\n        ")}
+      )
+    `);
+
+    // create HNSW indexes for vector fields
+    for (const vf of vectorFields) {
+      await this.pool.query(`
+        CREATE INDEX "${params.id}_${vf.name}_idx"
+        ON "${schemaName}"."${params.id}"
+        USING hnsw ("${vf.name}" ${SIMILARITY.encode(vf.similarity as "cosine" | "euclidean" | "dot_product" | undefined)})
+      `);
+    }
+
+    // auto-bind the created table
+    const fields: Record<string, PGFieldBinding> = {};
+    for (const [name, field] of Object.entries(params.schema)) {
+      fields[name] = {
+        column: name,
+        type: field.type,
+        ...(field.type === "vector" && {
+          dimensions: field.dimensions,
+          similarity: field.similarity,
+        }),
+      };
+    }
+
+    const config: PGIndexConfig = {
+      schema: schemaName,
+      table: params.id,
+      pkey,
+      fields,
+    };
+
+    // persist to metadata table
+    await this.pool.query(
+      `INSERT INTO "${KERNL_SCHEMA_NAME}"."${META_TABLE}" (id, backend, config, created_at)
+       VALUES ($1, $2, $3, $4)`,
+      [params.id, this.id, JSON.stringify(config), Date.now()],
+    );
+
+    this.configs.set(params.id, config);
+  }
+
+  /**
+   * List all indexes.
+   */
+  async listIndexes(
+    params?: ListIndexesParams,
+  ): Promise<CursorPage<IndexSummary>> {
+    await this.ensureInit();
+
+    const loader = async (
+      p: ListIndexesParams,
+    ): Promise<CursorPageResponse<IndexSummary>> => {
+      const limit = p.limit ?? 100;
+
+      let sql = `
+        SELECT id FROM "${KERNL_SCHEMA_NAME}"."${META_TABLE}"
+        WHERE backend = $1
+      `;
+      const sqlParams: unknown[] = [this.id];
+      let idx = 2;
+
+      if (p.prefix) {
+        sql += ` AND id LIKE $${idx}`;
+        sqlParams.push(`${p.prefix}%`);
+        idx++;
+      }
+
+      if (p.cursor) {
+        sql += ` AND id > $${idx}`;
+        sqlParams.push(p.cursor);
+        idx++;
+      }
+
+      sql += ` ORDER BY id ASC LIMIT $${idx}`;
+      sqlParams.push(limit + 1);
+
+      const result = await this.pool.query<{ id: string }>(sql, sqlParams);
+
+      const hasMore = result.rows.length > limit;
+      const rows = hasMore ? result.rows.slice(0, -1) : result.rows;
+
+      const data: IndexSummary[] = rows.map((row) => ({
+        id: row.id,
+        status: "ready" as const,
+      }));
+
+      return {
+        data,
+        next: hasMore ? (rows[rows.length - 1]?.id ?? null) : null,
+        last: !hasMore,
+      };
+    };
+
+    const response = await loader(params ?? {});
+
+    return new CursorPage({
+      params: params ?? {},
+      response,
+      loader,
+    });
+  }
+
+  /**
+   * Get index statistics.
+   */
+  async describeIndex(id: string): Promise<IndexStats> {
+    await this.ensureInit();
+
+    const cfg = this.configs.get(id);
+    if (!cfg) {
+      throw new Error(`Index "${id}" not bound`);
+    }
+
+    // get row count
+    const countRes = await this.pool.query<{ count: string }>(
+      `SELECT COUNT(*) as count FROM "${cfg.schema}"."${cfg.table}"`,
+    );
+    const count = parseInt(countRes.rows[0]?.count ?? "0", 10);
+
+    // get table size in bytes
+    const sizeRes = await this.pool.query<{ size: string }>(
+      `SELECT pg_total_relation_size('"${cfg.schema}"."${cfg.table}"') as size`,
+    );
+    const sizeb = parseInt(sizeRes.rows[0]?.size ?? "0", 10);
+
+    // find vector field for dimensions/similarity
+    const vectorField = Object.values(cfg.fields).find(
+      (f) => f.type === "vector",
+    );
+
+    return {
+      id,
+      count,
+      sizeb,
+      dimensions: vectorField?.dimensions,
+      similarity: vectorField?.similarity,
+      status: "ready",
+    };
+  }
+
+  /**
+   * Delete an index and all its documents.
+   */
+  async deleteIndex(id: string): Promise<void> {
+    await this.ensureInit();
+
+    const cfg = this.configs.get(id);
+    if (!cfg) {
+      throw new Error(`Index "${id}" not bound`);
+    }
+
+    await this.pool.query(
+      `DROP TABLE IF EXISTS "${cfg.schema}"."${cfg.table}"`,
+    );
+    await this.pool.query(
+      `DELETE FROM "${KERNL_SCHEMA_NAME}"."${META_TABLE}" WHERE id = $1`,
+      [id],
+    );
+    this.configs.delete(id);
+  }
+
+  /**
+   * No-op for pgvector.
+   */
+  async warm(_id: string): Promise<void> {}
+
+  /**
+   * pgvector capabilities.
+   */
+  capabilities(): SearchCapabilities {
+    return {
+      modes: new Set(["vector"]),
+      multiSignal: false,
+      multiVector: false,
+      filters: true,
+      orderBy: true,
+    };
+  }
+
+  /**
+   * Get a handle for operating on a specific index.
+   */
+  index<TDocument = UnknownDocument>(id: string): IndexHandle<TDocument> {
+    const cfg = this.configs.get(id);
+    return new PGIndexHandle<TDocument>(
+      this.pool,
+      () => this.ensureInit(),
+      id,
+      cfg,
+    );
+  }
+
+  /**
+   * Bind an existing Postgres table as an index.
+   */
+  async bindIndex(id: string, config: PGIndexConfig): Promise<void> {
+    await this.ensureInit();
+
+    // upsert to metadata table
+    await this.pool.query(
+      `INSERT INTO "${KERNL_SCHEMA_NAME}"."${META_TABLE}" (id, backend, config, created_at)
+       VALUES ($1, $2, $3, $4)
+       ON CONFLICT (id) DO UPDATE SET config = $3`,
+      [id, this.id, JSON.stringify(config), Date.now()],
+    );
+
+    this.configs.set(id, config);
+  }
+
+  /* --- internal utils --- */
+
+  /**
+   * Ensure metadata table exists and load configs.
+   */
+  private async ensureInit(): Promise<void> {
+    if (this.ready) return;
+
+    await this.userInit();
+
+    // create metadata table
+    await this.pool.query(`
+      CREATE TABLE IF NOT EXISTS "${KERNL_SCHEMA_NAME}"."${META_TABLE}" (
+        id TEXT PRIMARY KEY,
+        backend TEXT NOT NULL,
+        config JSONB NOT NULL,
+        created_at BIGINT NOT NULL
+      )
+    `);
+
+    // load existing configs for this backend
+    const result = await this.pool.query<{
+      id: string;
+      config: PGIndexConfig;
+    }>(
+      `SELECT id, config FROM "${KERNL_SCHEMA_NAME}"."${META_TABLE}" WHERE backend = $1`,
+      [this.id],
+    );
+
+    for (const row of result.rows) {
+      this.configs.set(row.id, row.config);
+    }
+
+    this.ready = true;
+  }
+}

package/src/pgvector/sql/__tests__/limit.test.ts

@@ -0,0 +1,180 @@
+import { describe, it, expect } from "vitest";
+import { SQL_LIMIT } from "../limit";
+
+describe("SQL_LIMIT", () => {
+  describe("encode", () => {
+    it("builds LIMIT clause", () => {
+      const result = SQL_LIMIT.encode({
+        topK: 10,
+        offset: 0,
+        startIdx: 1,
+      });
+      expect(result.sql).toBe("LIMIT $1");
+      expect(result.params).toEqual([10]);
+    });
+
+    it("respects startIdx for parameter numbering", () => {
+      const result = SQL_LIMIT.encode({
+        topK: 10,
+        offset: 0,
+        startIdx: 5,
+      });
+      expect(result.sql).toBe("LIMIT $5");
+      expect(result.params).toEqual([10]);
+    });
+
+    it("includes OFFSET when offset > 0", () => {
+      const result = SQL_LIMIT.encode({
+        topK: 10,
+        offset: 20,
+        startIdx: 1,
+      });
+      expect(result.sql).toBe("LIMIT $1 OFFSET $2");
+      expect(result.params).toEqual([10, 20]);
+    });
+
+    it("skips OFFSET when offset is 0", () => {
+      const result = SQL_LIMIT.encode({
+        topK: 25,
+        offset: 0,
+        startIdx: 3,
+      });
+      expect(result.sql).toBe("LIMIT $3");
+      expect(result.params).toEqual([25]);
+    });
+
+    it("handles pagination correctly", () => {
+      // Page 1: offset 0
+      const page1 = SQL_LIMIT.encode({
+        topK: 20,
+        offset: 0,
+        startIdx: 1,
+      });
+      expect(page1.sql).toBe("LIMIT $1");
+      expect(page1.params).toEqual([20]);
+
+      // Page 2: offset 20
+      const page2 = SQL_LIMIT.encode({
+        topK: 20,
+        offset: 20,
+        startIdx: 1,
+      });
+      expect(page2.sql).toBe("LIMIT $1 OFFSET $2");
+      expect(page2.params).toEqual([20, 20]);
+
+      // Page 3: offset 40
+      const page3 = SQL_LIMIT.encode({
+        topK: 20,
+        offset: 40,
+        startIdx: 1,
+      });
+      expect(page3.sql).toBe("LIMIT $1 OFFSET $2");
+      expect(page3.params).toEqual([20, 40]);
+    });
+
+    it("correctly increments param index after SELECT and WHERE", () => {
+      // Simulating: SELECT uses $1, WHERE uses $2-$4
+      // LIMIT should start at $5
+      const result = SQL_LIMIT.encode({
+        topK: 10,
+        offset: 50,
+        startIdx: 5,
+      });
+      expect(result.sql).toBe("LIMIT $5 OFFSET $6");
+      expect(result.params).toEqual([10, 50]);
+    });
+
+    describe("edge values", () => {
+      it("handles topK: 0", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 0,
+          offset: 0,
+          startIdx: 1,
+        });
+        // LIMIT 0 is valid SQL - returns no rows
+        expect(result.sql).toBe("LIMIT $1");
+        expect(result.params).toEqual([0]);
+      });
+
+      it("handles topK: 1", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 1,
+          offset: 0,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1");
+        expect(result.params).toEqual([1]);
+      });
+
+      it("handles very large topK", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 1000000,
+          offset: 0,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1");
+        expect(result.params).toEqual([1000000]);
+      });
+
+      it("handles very large offset", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 10,
+          offset: 999999,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1 OFFSET $2");
+        expect(result.params).toEqual([10, 999999]);
+      });
+
+      it("handles very large startIdx", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 10,
+          offset: 20,
+          startIdx: 50,
+        });
+        expect(result.sql).toBe("LIMIT $50 OFFSET $51");
+        expect(result.params).toEqual([10, 20]);
+      });
+
+      it("handles startIdx: 1 with both topK and offset", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 25,
+          offset: 100,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1 OFFSET $2");
+        expect(result.params).toEqual([25, 100]);
+      });
+    });
+
+    describe("offset boundary", () => {
+      it("includes OFFSET when offset is exactly 1", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 10,
+          offset: 1,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1 OFFSET $2");
+        expect(result.params).toEqual([10, 1]);
+      });
+
+      it("does not include OFFSET when offset is exactly 0", () => {
+        const result = SQL_LIMIT.encode({
+          topK: 10,
+          offset: 0,
+          startIdx: 1,
+        });
+        expect(result.sql).toBe("LIMIT $1");
+        expect(result.params).toEqual([10]);
+      });
+    });
+  });
+
+  describe("decode", () => {
+    it("throws not implemented", () => {
+      expect(() => SQL_LIMIT.decode({} as any)).toThrow(
+        "SQL_LIMIT.decode not implemented",
+      );
+    });
+  });
+});