iframe-proxy | Sunbelt Computer Software

History

64 lines (53 loc) · 1.79 KB

Raw

import pg from 'pg';

import pgvector from 'pgvector/pg';

import { from as copyFrom } from 'pg-copy-streams';

import { stdout } from 'process';

// generate random data

const rows = 100000;

const dimensions = 128;

const embeddings = Array.from({length: rows}, () => Array.from({length: dimensions}, () => Math.random()));

// connect

const client = new pg.Client({database: 'pgvector_example'});

await client.connect();

// enable extension

await client.query('CREATE EXTENSION IF NOT EXISTS vector');

await pgvector.registerTypes(client);

// create table

await client.query('DROP TABLE IF EXISTS items');

await client.query(`CREATE TABLE items (id bigserial, embedding vector(${dimensions}))`);

function copyRow(stream, line) {

return new Promise((resolve) => {

let ok = stream.write(line);

if (!ok) {

stream.once('drain', () => resolve());

} else {

resolve();

}

});

}

// load data

console.log(`Loading ${embeddings.length} rows`);

const stream = client.query(copyFrom('COPY items (embedding) FROM STDIN'));

for (const [i, embedding] of embeddings.entries()) {

const line = `${pgvector.toSql(embedding)}\n`;

await copyRow(stream, line);

// show progress

if (i % 10000 == 0) {

stdout.write('.');

}

stream.on('finish', async function () {

console.log('\nSuccess!');

// create any indexes *after* loading initial data (skipping for this example)

const createIndex = false;

if (createIndex) {

console.log('Creating index');

await client.query("SET maintenance_work_mem = '8GB'");

await client.query('SET max_parallel_maintenance_workers = 7');

await client.query('CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)');

}

// update planner statistics for good measure

await client.query('ANALYZE items');

client.end();

});

stream.end();

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

example.js

Latest commit

History

example.js

File metadata and controls