Repository navigation
Expand file tree
/
Copy pathexample.js
More file actions
64 lines (53 loc) · 1.79 KB
/
Copy pathexample.js
File metadata and controls
64 lines (53 loc) · 1.79 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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();