First commit

ankane · ankane · commit 4f95724be480 · 2021-06-18T14:58:52.000-07:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -0,0 +1,20 @@
+name: build
+on: [push, pull_request]
+jobs:
+  build:
+    if: "!contains(github.event.head_commit.message, '[skip ci]')"
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-node@v2
+    - run: npm install
+    - uses: ankane/setup-postgres@v1
+      with:
+        database: pgvector_node_test
+    - run: |
+        sudo apt-get update && sudo apt-get install postgresql-server-dev-13
+        git clone --branch v0.1.7 https://github.com/ankane/pgvector.git
+        cd pgvector
+        make
+        sudo make install
+    - run: npm test
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,2 @@
+/node_modules/
+/package-lock.json
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -0,0 +1,3 @@
+## 0.1.0 (unreleased)
+
+- First release
diff --git a/LICENSE.txt b/LICENSE.txt
@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2021 Andrew Kane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
diff --git a/README.md b/README.md
@@ -0,0 +1,101 @@
+# pgvector-node
+
+[pgvector](https://github.com/ankane/pgvector) support for Node.js
+
+Supports [Sequelize](https://github.com/sequelize/sequelize) and [node-postgres](https://github.com/brianc/node-postgres)
+
+[![Build Status](https://github.com/ankane/pgvector-node/workflows/build/badge.svg?branch=master)](https://github.com/ankane/pgvector-node/actions)
+
+## Installation
+
+Run:
+
+```sh
+npm install pgvector
+```
+
+And follow the instructions for your database library:
+
+- [Sequelize](#sequelize)
+- [node-postgres](#node-postgres)
+
+## Sequelize
+
+Register the type
+
+```js
+const { Sequelize } = require('sequelize');
+const pgvector = require('pgvector/sequelize');
+
+pgvector.registerType(Sequelize);
+```
+
+Add a vector field
+
+```js
+Item.init({
+  factors: {
+    type: DataTypes.VECTOR(3)
+  }
+}, ...);
+```
+
+Insert a vector
+
+```js
+await Item.create({factors: [1, 2, 3]});
+```
+
+Get the nearest neighbors to a vector
+
+```js
+const items = await Item.findAll({
+  order: [sequelize.literal(`factors <-> '[1, 2, 3]'`)],
+  limit: 5
+});
+```
+
+## node-postgres
+
+Register the type
+
+```js
+const pgvector = require('pgvector/pg');
+
+await pgvector.registerType(client);
+```
+
+Insert a vector
+
+```js
+const factors = [1, 2, 3];
+await client.query('INSERT INTO items (factors) VALUES ($1)', [pgvector.toSql(factors)]);
+```
+
+Get the nearest neighbors to a vector
+
+```js
+const result = await client.query('SELECT * FROM items ORDER BY factors <-> $1 LIMIT 5', [pgvector.toSql(factors)]);
+```
+
+## History
+
+View the [changelog](https://github.com/ankane/pgvector-node/blob/master/CHANGELOG.md)
+
+## Contributing
+
+Everyone is encouraged to help improve this project. Here are a few ways you can help:
+
+- [Report bugs](https://github.com/ankane/pgvector-node/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/pgvector-node/pulls)
+- Write, clarify, or fix documentation
+- Suggest or add new features
+
+To get started with development:
+
+```sh
+git clone https://github.com/ankane/pgvector-node.git
+cd pgvector-node
+npm install
+npm test
+```
diff --git a/package.json b/package.json
@@ -0,0 +1,35 @@
+{
+  "name": "pgvector",
+  "version": "0.1.0",
+  "description": "pgvector support for Node.js",
+  "homepage": "https://github.com/ankane/pgvector-node",
+  "license": "MIT",
+  "authors": [
+    "ankane"
+  ],
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/ankane/pgvector-node.js"
+  },
+  "exports": {
+    "./pg": "./pg/index.js",
+    "./sequelize": "./sequelize/index.js",
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "pg",
+    "sequelize",
+    "utils"
+  ],
+  "engines": {
+    "node": ">= 12"
+  },
+  "scripts": {
+    "test": "jest"
+  },
+  "devDependencies": {
+    "jest": "^27.0.4",
+    "pg": "^8.6.0",
+    "sequelize": "^6.6.2"
+  }
+}
diff --git a/pg/index.js b/pg/index.js
@@ -0,0 +1,21 @@
+const utils = require('../utils');
+
+async function registerType(client) {
+  const result = await client.query('SELECT typname, oid, typarray FROM pg_type WHERE typname = $1', ['vector']);
+  if (result.rowCount < 1) {
+    throw new Error('vector type not found in the database');
+  }
+  const oid = result.rows[0].oid;
+  client.setTypeParser(oid, 'text', function(value) {
+    return utils.fromSql(value);
+  });
+}
+
+function toSql(value) {
+  if (!Array.isArray(value)) {
+    throw new Error('expected array');
+  }
+  return utils.toSql(value);
+}
+
+module.exports = {registerType, toSql};
diff --git a/pg/index.test.js b/pg/index.test.js
@@ -0,0 +1,38 @@
+const { Client } = require('pg');
+const pgvector = require('./index');
+
+const client = new Client({database: 'pgvector_node_test'});
+
+beforeAll(async () => {
+  await client.connect();
+  const sql = `
+  CREATE EXTENSION IF NOT EXISTS vector;
+  DROP TABLE IF EXISTS items;
+  CREATE TABLE IF NOT EXISTS items (
+    id serial primary key,
+    factors vector(3)
+  );
+  `;
+  await client.query(sql);
+  await pgvector.registerType(client);
+});
+
+afterAll(async () => {
+  await client.end();
+});
+
+beforeEach(async () => {
+  await client.query('DELETE FROM items');
+});
+
+test('works', async () => {
+  await client.query('INSERT INTO items (factors) VALUES ($1)', [pgvector.toSql([1, 2, 3])]);
+  const { rows } = await client.query('SELECT * FROM items ORDER BY factors <-> $1 LIMIT 5', [pgvector.toSql([1, 2, 3])]);
+  expect(rows[0].factors).toStrictEqual([1, 2, 3]);
+});
+
+test('bad object', () => {
+  expect(() => {
+    pgvector.toSql({hello: 'world'});
+  }).toThrowError('expected array');
+});
diff --git a/sequelize/index.js b/sequelize/index.js
@@ -0,0 +1,51 @@
+const util = require('util');
+const utils = require('../utils');
+
+function registerType(Sequelize) {
+  const DataTypes = Sequelize.DataTypes;
+  const PgTypes = DataTypes.postgres;
+  const ABSTRACT = DataTypes.ABSTRACT.prototype.constructor;
+
+  class VECTOR extends ABSTRACT {
+    constructor(dimensions) {
+      super();
+      this._dimensions = dimensions;
+    }
+
+    toSql() {
+      if (this._dimensions === undefined) {
+        return 'VECTOR';
+      }
+      if (!Number.isInteger(this._dimensions)) {
+        throw new Error('expected integer');
+      }
+      return util.format('VECTOR(%d)', this._dimensions);
+    }
+
+    _stringify(value) {
+      return utils.toSql(value);
+    }
+
+    static parse(value) {
+      return utils.fromSql(value);
+    }
+  }
+
+  VECTOR.prototype.key = VECTOR.key = 'vector';
+
+  DataTypes.VECTOR = Sequelize.Utils.classToInvokable(VECTOR);
+  DataTypes.VECTOR.types.postgres = ['vector'];
+
+  PgTypes.VECTOR = function VECTOR() {
+    if (!(this instanceof PgTypes.VECTOR)) {
+      return new PgTypes.VECTOR();
+    }
+    DataTypes.VECTOR.apply(this, arguments);
+  };
+  util.inherits(PgTypes.VECTOR, DataTypes.VECTOR);
+  PgTypes.VECTOR.parse = DataTypes.VECTOR.parse;
+  PgTypes.VECTOR.types = {postgres: ['vector']};
+  DataTypes.postgres.VECTOR.key = 'vector';
+}
+
+module.exports = {registerType};
diff --git a/sequelize/index.test.js b/sequelize/index.test.js
@@ -0,0 +1,57 @@
+const { Sequelize, DataTypes, Model } = require('sequelize');
+const pgvector = require('./index');
+
+pgvector.registerType(Sequelize);
+
+class Item extends Model {}
+
+const sequelize = new Sequelize('postgres://localhost/pgvector_node_test', {
+  logging: false
+});
+
+Item.init({
+  factors: {
+    type: DataTypes.VECTOR(3)
+  }
+}, {
+  sequelize,
+  modelName: 'Item'
+});
+
+beforeAll(async () => {
+  await sequelize.authenticate();
+  await sequelize.query('CREATE EXTENSION IF NOT EXISTS vector');
+  await Item.sync({force: true});
+});
+
+afterAll(() => {
+  sequelize.close();
+});
+
+test('works', async () => {
+  await Item.create({factors: [1, 2, 3]});
+  const items = await Item.findAll({
+    order: [sequelize.literal(`factors <-> '[1, 1, 1]'`)],
+    limit: 5
+  });
+  expect(items[0].factors).toStrictEqual([1, 2, 3]);
+});
+
+test('bad value', () => {
+  expect.assertions(1);
+  return Item.create({factors: 'bad'}).catch(e => expect(e.message).toMatch('malformed vector literal'));
+});
+
+test('dimensions', () => {
+  expect(DataTypes.VECTOR(3).toSql()).toBe('VECTOR(3)');
+});
+
+test('no dimensions', () => {
+  expect(DataTypes.VECTOR().toSql()).toBe('VECTOR');
+});
+
+test('bad dimensions', () => {
+  expect(() => {
+    DataTypes.VECTOR('bad').toSql()
+  }).toThrowError('expected integer');
+});
diff --git a/utils/index.js b/utils/index.js
@@ -0,0 +1,9 @@
+function fromSql(value) {
+  return value.substring(1, value.length - 1).split(',').map((v) => parseFloat(v));
+}
+
+function toSql(value) {
+  return JSON.stringify(value);
+}
+
+module.exports = {fromSql, toSql};

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+/node_modules/`
	`2`	`+/package-lock.json`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+## 0.1.0 (unreleased)`
	`2`	`+`
	`3`	`+- First release`