pg-read-write-gz.ts

import fs from "node:fs/promises";
import { createReadStream } from "node:fs";
import zlib from "node:zlib";
import { pipeline } from "node:stream/promises";
import { PGlite } from "@electric-sql/pglite";

async function writePg(filePath: string, pglite: PGlite) {
  await fs.writeFile(filePath, (await pglite.dumpDataDir("gzip")).stream());
}
async function readPg(filePath: string) {
  const chunks: Array<ArrayBufferLike> = [];
  await pipeline(
    createReadStream(filePath),
    zlib.createGunzip(),
    async function* (source: AsyncIterable<Buffer>) {
      for await (const chunk of source) {
        chunks.push(chunk);
      }
    }
  );
  return new Blob(chunks);
}

// with vector
import { vector } from "@electric-sql/pglite/vector";
{
  // init db
  const pglite = new PGlite({
    extensions: { vector },
  });
  await pglite.sql`CREATE EXTENSION IF NOT EXISTS vector`;
  await pglite.sql`CREATE TABLE IF NOT EXISTS items (id bigserial PRIMARY KEY, embedding vector(3))`;
  await pglite.sql`CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)`;
  //await pg.sql`INSERT INTO items (embedding) VALUES ('[1, 2, 3]')`;
  for (let i = 0; i < 1000; i++) {
    await pglite.sql`INSERT INTO items (embedding) VALUES (${`[${Math.random()}, ${Math.random()}, ${Math.random()}]`})`;
  }
  await writePg("foo.gz", pglite);
  await pglite.close();
}

{
  // query
  const pg = new PGlite({
    extensions: { vector },
    loadDataDir: await readPg("foo.gz"),
  });

  const q = `[${Math.random()}, ${Math.random()}, ${Math.random()}]`;
  console.log(
    await pg.sql`SELECT * FROM items ORDER BY embedding <-> ${q} LIMIT 10`
  );
}