iframe-proxy | Sunbelt Computer Software

History

106 lines (89 loc) · 3.49 KB

Raw

100

101

102

103

104

105

106

import assert from 'node:assert';

import test from 'node:test';

import Knex from 'knex';

import { Model } from 'objection';

import pgvector from 'pgvector/objection';

import { l2Distance, maxInnerProduct, cosineDistance, l1Distance, hammingDistance, jaccardDistance } from 'pgvector/objection';

import { SparseVector } from 'pgvector';

test('objection example', async () => {

const knex = Knex({

client: 'pg',

connection: {database: 'pgvector_node_test'}

});

Model.knex(knex);

class Item extends Model {

/** @type {number} */

id;

/** @type {?string} */

embedding;

static get tableName() {

return 'objection_items';

}

// @ts-ignore

// https://github.com/knex/knex/issues/3853#issuecomment-628955467

await knex.schema.createExtensionIfNotExists('vector');

await knex.schema.dropTableIfExists('objection_items');

await knex.schema.createTable('objection_items', (table) => {

table.increments('id');

table.vector('embedding', 3);

table.halfvec('half_embedding', 3);

// @ts-ignore

table.bit('binary_embedding', {length: 3});

table.sparsevec('sparse_embedding', 3);

});

const newItems = [

{embedding: pgvector.toSql([1, 1, 1]), half_embedding: pgvector.toSql([1, 1, 1]), binary_embedding: '000', sparse_embedding: new SparseVector([1, 1, 1])},

{embedding: pgvector.toSql([2, 2, 2]), half_embedding: pgvector.toSql([2, 2, 2]), binary_embedding: '101', sparse_embedding: new SparseVector([2, 2, 2])},

{embedding: pgvector.toSql([1, 1, 2]), half_embedding: pgvector.toSql([1, 1, 2]), binary_embedding: '111', sparse_embedding: new SparseVector([1, 1, 2])},

{embedding: null}

];

await Item.query().insert(newItems);

// L2 distance

let items = await Item.query()

.orderBy(l2Distance('embedding', [1, 1, 1]))

.limit(5);

assert.deepEqual(items.map(v => v.id), [1, 3, 2, 4]);

assert.deepEqual(pgvector.fromSql(items[0].embedding), [1, 1, 1]);

assert.deepEqual(pgvector.fromSql(items[1].embedding), [1, 1, 2]);

assert.deepEqual(pgvector.fromSql(items[2].embedding), [2, 2, 2]);

// L2 distance - halfvec

items = await Item.query()

.orderBy(l2Distance('half_embedding', [1, 1, 1]))

.limit(5);

assert.deepEqual(items.map(v => v.id), [1, 3, 2, 4]);

// L2 distance - sparsevec

items = await Item.query()

.orderBy(l2Distance('sparse_embedding', new SparseVector([1, 1, 1])))

.limit(5);

assert.deepEqual(items.map(v => v.id), [1, 3, 2, 4]);

// max inner product

items = await Item.query()

.orderBy(maxInnerProduct('embedding', [1, 1, 1]))

.limit(5);

assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);

// cosine distance

items = await Item.query()

.orderBy(cosineDistance('embedding', [1, 1, 1]))

.limit(5);

assert.deepEqual(items.map(v => v.id).slice(2), [3, 4]);

// L1 distance

items = await Item.query()

.orderBy(l1Distance('embedding', [1, 1, 1]))

.limit(5);

assert.deepEqual(items.map(v => v.id), [1, 3, 2, 4]);

// Hamming distance

items = await Item.query()

.orderBy(hammingDistance('binary_embedding', '101'))

.limit(5);

assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);

// Jaccard distance

items = await Item.query()

.orderBy(jaccardDistance('binary_embedding', '101'))

.limit(5);

assert.deepEqual(items.map(v => v.id), [2, 3, 1, 4]);

await knex.schema.alterTable('objection_items', function (table) {

table.index([knex.raw('embedding vector_l2_ops')], 'objection_items_embedding_idx', 'hnsw');

});

await knex.destroy();

});

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

objection.test.js

Latest commit

History

objection.test.js

File metadata and controls