wehub-resource-sync
384 lines
9.6 KiB
TypeScript
384 lines
9.6 KiB
TypeScript
import Docker from 'dockerode';
|
|
import { eq, hammingDistance, jaccardDistance, l2Distance, not, sql } from 'drizzle-orm';
|
|
import { bigserial, bit, halfvec, pgTable, sparsevec, vector } from 'drizzle-orm/pg-core';
|
|
import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js';
|
|
import getPort from 'get-port';
|
|
import postgres, { type Sql } from 'postgres';
|
|
import { v4 as uuid } from 'uuid';
|
|
import { afterAll, beforeAll, beforeEach, expect, test } from 'vitest';
|
|
|
|
const ENABLE_LOGGING = false;
|
|
|
|
let pgContainer: Docker.Container;
|
|
let docker: Docker;
|
|
let client: Sql;
|
|
let db: PostgresJsDatabase;
|
|
|
|
async function createDockerDB(): Promise<string> {
|
|
const inDocker = (docker = new Docker());
|
|
const port = await getPort({ port: 5432 });
|
|
const image = 'pgvector/pgvector:pg16';
|
|
|
|
const pullStream = await docker.pull(image);
|
|
await new Promise((resolve, reject) =>
|
|
inDocker.modem.followProgress(pullStream, (err) => (err ? reject(err) : resolve(err)))
|
|
);
|
|
|
|
pgContainer = await docker.createContainer({
|
|
Image: image,
|
|
Env: ['POSTGRES_PASSWORD=postgres', 'POSTGRES_USER=postgres', 'POSTGRES_DB=postgres'],
|
|
name: `drizzle-integration-tests-${uuid()}`,
|
|
HostConfig: {
|
|
AutoRemove: true,
|
|
PortBindings: {
|
|
'5432/tcp': [{ HostPort: `${port}` }],
|
|
},
|
|
},
|
|
});
|
|
|
|
await pgContainer.start();
|
|
|
|
return `postgres://postgres:postgres@localhost:${port}/postgres`;
|
|
}
|
|
|
|
beforeAll(async () => {
|
|
const connectionString = process.env['PG_VECTOR_CONNECTION_STRING'] ?? (await createDockerDB());
|
|
|
|
const sleep = 250;
|
|
let timeLeft = 5000;
|
|
let connected = false;
|
|
let lastError: unknown | undefined;
|
|
do {
|
|
try {
|
|
client = postgres(connectionString, {
|
|
max: 1,
|
|
onnotice: () => {
|
|
// disable notices
|
|
},
|
|
});
|
|
await client`select 1`;
|
|
connected = true;
|
|
break;
|
|
} catch (e) {
|
|
lastError = e;
|
|
await new Promise((resolve) => setTimeout(resolve, sleep));
|
|
timeLeft -= sleep;
|
|
}
|
|
} while (timeLeft > 0);
|
|
if (!connected) {
|
|
console.error('Cannot connect to Postgres');
|
|
await client?.end().catch(console.error);
|
|
await pgContainer?.stop().catch(console.error);
|
|
throw lastError;
|
|
}
|
|
db = drizzle(client, { logger: ENABLE_LOGGING });
|
|
|
|
await db.execute(sql`CREATE EXTENSION IF NOT EXISTS vector;`);
|
|
});
|
|
|
|
afterAll(async () => {
|
|
await client?.end().catch(console.error);
|
|
await pgContainer?.stop().catch(console.error);
|
|
});
|
|
|
|
const items = pgTable('items', {
|
|
id: bigserial('id', { mode: 'number' }).primaryKey(),
|
|
vector: vector('vector', { dimensions: 3 }),
|
|
bit: bit('bit', { dimensions: 3 }),
|
|
halfvec: halfvec('halfvec', { dimensions: 3 }),
|
|
sparsevec: sparsevec('sparsevec', { dimensions: 5 }),
|
|
});
|
|
|
|
beforeEach(async () => {
|
|
await db.execute(sql`drop table if exists items cascade`);
|
|
await db.execute(sql`
|
|
CREATE TABLE items (
|
|
id bigserial PRIMARY KEY,
|
|
"vector" vector(3),
|
|
"bit" bit(3),
|
|
"halfvec" halfvec(3),
|
|
"sparsevec" sparsevec(5)
|
|
);
|
|
`);
|
|
});
|
|
|
|
test('insert + partial select', async () => {
|
|
const insertedValues = await db.insert(items).values([{
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]).returning();
|
|
|
|
const query = db.select({ distance: l2Distance(items.vector, [3, 1, 2]) }).from(items);
|
|
|
|
expect(query.toSQL()).toStrictEqual({ sql: 'select "vector" <-> $1 from "items"', params: ['[3,1,2]'] });
|
|
|
|
const response = await query;
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
expect(response).toStrictEqual([{ distance: 0 }]);
|
|
});
|
|
|
|
// SELECT * FROM items WHERE embedding <-> '[3,1,2]' < 5;
|
|
test('insert + complex where', async () => {
|
|
const insertedValues = await db.insert(items).values([{
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]).returning();
|
|
|
|
const query = db.select().from(items)
|
|
.where(sql`${l2Distance(items.vector, [3, 1, 2])} < ${5}`)
|
|
.limit(5);
|
|
|
|
expect(query.toSQL()).toStrictEqual({
|
|
sql:
|
|
'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" where "items"."vector" <-> $1 < $2 limit $3',
|
|
params: ['[3,1,2]', 5, 5],
|
|
});
|
|
const res = await query;
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
expect(res).toStrictEqual([
|
|
{
|
|
id: 1,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
},
|
|
]);
|
|
});
|
|
|
|
// SELECT * FROM items WHERE id != 1 ORDER BY embedding <-> (SELECT embedding FROM items WHERE id = 1) LIMIT 5;
|
|
test('insert + order by subquery', async () => {
|
|
const insertedValues = await db.insert(items).values([{
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}, {
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]).returning();
|
|
|
|
const subquery = db.select({ vector: items.vector }).from(items).where(eq(items.id, 1));
|
|
|
|
const query = db.select().from(items)
|
|
.where(not(eq(items.id, 1)))
|
|
.orderBy(l2Distance(items.vector, subquery))
|
|
.limit(5);
|
|
|
|
expect(query.toSQL()).toStrictEqual({
|
|
sql:
|
|
'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" where not "items"."id" = $1 order by "items"."vector" <-> (select "vector" from "items" where "items"."id" = $2) limit $3',
|
|
params: [1, 1, 5],
|
|
});
|
|
const res = await query;
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}, {
|
|
id: 2,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
expect(res).toStrictEqual([
|
|
{
|
|
id: 2,
|
|
vector: [3, 1, 2],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
},
|
|
]);
|
|
});
|
|
|
|
test('insert + select order by jaccard distance', async () => {
|
|
const insertedValues = await db.insert(items).values({
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}).returning();
|
|
|
|
const bitQuery = db.select().from(items).orderBy(jaccardDistance(items.bit, '101')).limit(5);
|
|
|
|
expect(bitQuery.toSQL()).toStrictEqual({
|
|
params: [
|
|
'101',
|
|
5,
|
|
],
|
|
sql: 'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" order by "items"."bit" <%> $1 limit $2',
|
|
});
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
const response = await bitQuery;
|
|
|
|
expect(response).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
});
|
|
|
|
test('insert + select order by hamming distance', async () => {
|
|
const insertedValues = await db.insert(items).values({
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}).returning();
|
|
|
|
const bitQuery = db.select().from(items).orderBy(hammingDistance(items.bit, '101')).limit(5);
|
|
|
|
expect(bitQuery.toSQL()).toStrictEqual({
|
|
params: [
|
|
'101',
|
|
5,
|
|
],
|
|
sql: 'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" order by "items"."bit" <~> $1 limit $2',
|
|
});
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
const response = await bitQuery;
|
|
|
|
expect(response).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
});
|
|
|
|
test('insert + select order by l2 distance', async () => {
|
|
const insertedValues = await db.insert(items).values({
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}).returning();
|
|
|
|
const queryVector = db.select().from(items).orderBy(l2Distance(items.vector, [3, 1, 2])).limit(5);
|
|
const queryHalfvec = db.select().from(items).orderBy(l2Distance(items.halfvec, [3, 1, 2])).limit(5);
|
|
const querySparsevec = db.select().from(items).orderBy(l2Distance(items.sparsevec, '{1:3,3:1,5:2}/5')).limit(5);
|
|
|
|
expect(queryVector.toSQL()).toStrictEqual({
|
|
params: [
|
|
'[3,1,2]',
|
|
5,
|
|
],
|
|
sql: 'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" order by "items"."vector" <-> $1 limit $2',
|
|
});
|
|
|
|
expect(queryHalfvec.toSQL()).toStrictEqual({
|
|
params: [
|
|
'[3,1,2]',
|
|
5,
|
|
],
|
|
sql: 'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" order by "items"."halfvec" <-> $1 limit $2',
|
|
});
|
|
|
|
expect(querySparsevec.toSQL()).toStrictEqual({
|
|
params: [
|
|
'{1:3,3:1,5:2}/5',
|
|
5,
|
|
],
|
|
sql:
|
|
'select "id", "vector", "bit", "halfvec", "sparsevec" from "items" order by "items"."sparsevec" <-> $1 limit $2',
|
|
});
|
|
|
|
const vectorRes = await queryVector;
|
|
const halfvecRes = await queryHalfvec;
|
|
const sparsevecRes = await querySparsevec;
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
const expectedResponse = [{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}];
|
|
|
|
expect(vectorRes).toStrictEqual(expectedResponse);
|
|
expect(halfvecRes).toStrictEqual(expectedResponse);
|
|
expect(sparsevecRes).toStrictEqual(expectedResponse);
|
|
});
|
|
|
|
test('select + insert all vectors', async () => {
|
|
const insertedValues = await db.insert(items).values({
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}).returning();
|
|
|
|
const response = await db.select().from(items);
|
|
|
|
expect(insertedValues).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
|
|
expect(response).toStrictEqual([{
|
|
id: 1,
|
|
vector: [1, 2, 3],
|
|
bit: '000',
|
|
halfvec: [1, 2, 3],
|
|
sparsevec: '{1:1,3:2,5:3}/5',
|
|
}]);
|
|
});
|