chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,383 @@
|
||||
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',
|
||||
}]);
|
||||
});
|
||||
Reference in New Issue
Block a user