Skip to content

Commit

Permalink
feat: support pgvector (#305)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sixmen
sixmen authored Nov 13, 2024
1 parent 6d4b772 commit 74ccc0f
Show file tree
Hide file tree
Showing 19 changed files with 421 additions and 40 deletions.
2 changes: 0 additions & 2 deletions docker-compose.yml
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
version: '3'

services:
mysql:
image: mysql:8.0
Expand Down
2 changes: 2 additions & 0 deletions packages/cormo/lib/cjs/adapters/base.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,11 @@ export interface Schemas {
[table_name: string]: any;
};
}
export type VectorOrderOption = Record<string, Record<string, number[] | undefined> | undefined>;
export interface AdapterFindOptions {
lean: boolean;
orders: any[];
vector_order?: VectorOrderOption;
near?: any;
select?: string[];
select_raw?: string[];
Expand Down
73 changes: 64 additions & 9 deletions packages/cormo/lib/cjs/adapters/postgresql.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,10 @@ function _typeToSQL(property) {
return 'BIGINT';
case types.GeoPoint:
return 'GEOMETRY(POINT)';
case types.Vector:
return property.type.dimension
? `VECTOR(${property.type.dimension})`
: 'VECTOR';
case types.Date:
return 'TIMESTAMP WITHOUT TIME ZONE';
case types.Object:
Expand Down Expand Up @@ -454,6 +458,7 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
callback();
};
this._pool.connect().then((client) => {
this._connection._logger.logQuery(sql, params);
client
.query(new QueryStream(sql, params))
.on('end', () => {
Expand Down Expand Up @@ -648,10 +653,12 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
await this._connection._promise_connection;
}
if (transaction && transaction._adapter_connection) {
this._connection._logger.logQuery(text, values);
transaction.checkFinished();
return await transaction._adapter_connection.query(text, values);
}
else {
this._connection._logger.logQuery(text, values);
return await this._pool.query(text, values);
}
}
Expand All @@ -670,6 +677,9 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
}
return String(value);
}
else if (property.type_class === types.Vector) {
return JSON.parse(value);
}
return value;
}
/** @internal */
Expand Down Expand Up @@ -744,15 +754,17 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
column.udt_schema === 'public' &&
column.udt_name === 'geometry'
? new types.GeoPoint()
: column.data_type === 'timestamp without time zone'
? new types.Date()
: column.data_type === 'json'
? new types.Object()
: column.data_type === 'text'
? new types.Text()
: column.data_type === 'bytea'
? new types.Blob()
: undefined;
: column.data_type === 'vector'
? new types.Vector()
: column.data_type === 'timestamp without time zone'
? new types.Date()
: column.data_type === 'json'
? new types.Object()
: column.data_type === 'text'
? new types.Text()
: column.data_type === 'bytea'
? new types.Blob()
: undefined;
let adapter_type_string = column.data_type.toUpperCase();
if (column.data_type === 'character varying') {
adapter_type_string += `(${column.character_maximum_length || 255})`;
Expand Down Expand Up @@ -814,6 +826,16 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
fields.push(`"${dbname}"=ST_Point($${values.length - 1}, $${values.length})`);
}
}
else if (property.type_class === types.Vector) {
values.push(value != null ? JSON.stringify(value) : null);
if (insert) {
fields.push(`"${dbname}"`);
places.push(`$${values.length}`);
}
else {
fields.push(`"${dbname}"=$${values.length}`);
}
}
else if (value && value.$inc != null) {
values.push(value.$inc);
fields.push(`"${dbname}"="${dbname}"+$${values.length}`);
Expand Down Expand Up @@ -936,6 +958,39 @@ class PostgreSQLAdapter extends sql_base_js_1.SQLAdapterBase {
}
sql += ' ORDER BY ' + orders.join(',');
}
else if (options.vector_order &&
typeof options.vector_order === 'object' &&
Object.keys(options.vector_order).length === 1) {
const column = Object.keys(options.vector_order)[0];
const cond = options.vector_order[column];
if (cond && typeof cond === 'object' && Object.keys(cond).length === 1) {
const key = Object.keys(cond)[0];
const value = Object.values(cond)[0];
let op = '';
if (key === '$l2_distance') {
op = '<->';
}
else if (key === '$l1_distance') {
op = '<+>';
}
else if (key === '$cosine_distance') {
op = '<=>';
}
else if (key === '$negative_inner_product') {
op = '<#>';
}
else if (key === '$hamming_distance') {
op = '<~>';
}
else if (key === '$jaccard_distance') {
op = '<%>';
}
if (op) {
params.push(value != null ? JSON.stringify(value) : null);
sql += ` ORDER BY _Base."${column}" ${op} $${params.length}`;
}
}
}
if (options.limit) {
sql += ' LIMIT ' + options.limit;
if (options.skip) {
Expand Down
6 changes: 6 additions & 0 deletions packages/cormo/lib/cjs/query.d.ts
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import stream from 'stream';
import { VectorOrderOption } from './adapters/base.js';
import { BaseModel, ModelColumnNamesWithId } from './model/index.js';
import { Transaction } from './transaction.js';
import { RecordID } from './types.js';
interface QueryOptions {
lean: boolean;
orders?: string;
vector_order?: VectorOrderOption;
near?: any;
select_columns?: string[];
select_single: boolean;
Expand Down Expand Up @@ -43,6 +45,7 @@ export interface QuerySingle<M extends BaseModel, T = M> extends PromiseLike<T>
select<K extends ModelColumnNamesWithId<M>>(columns?: string): QuerySingle<M, Pick<M, K>>;
selectSingle<K extends ModelColumnNamesWithId<M>>(column: K): QuerySingle<M, M[K]>;
order(orders?: string): QuerySingle<M, T>;
vector_order(order: VectorOrderOption): QuerySingle<M, T>;
group<G extends ModelColumnNamesWithId<M>, F>(group_by: G | G[], fields?: F): QuerySingle<M, {
[field in keyof F]: number;
} & Pick<M, G>>;
Expand Down Expand Up @@ -95,6 +98,7 @@ interface QuerySingleNull<M extends BaseModel, T = M> extends PromiseLike<T | nu
select<K extends ModelColumnNamesWithId<M>>(columns?: string): QuerySingleNull<M, Pick<M, K>>;
selectSingle<K extends ModelColumnNamesWithId<M>>(column: K): QuerySingleNull<M, M[K]>;
order(orders?: string): QuerySingleNull<M, T>;
vector_order(order: VectorOrderOption): QuerySingleNull<M, T>;
group<G extends ModelColumnNamesWithId<M>, F>(group_by: G | G[], fields?: F): QuerySingleNull<M, {
[field in keyof F]: number;
} & Pick<M, G>>;
Expand Down Expand Up @@ -147,6 +151,7 @@ export interface QueryArray<M extends BaseModel, T = M> extends PromiseLike<T[]>
select<K extends ModelColumnNamesWithId<M>>(columns?: string): QueryArray<M, Pick<M, K>>;
selectSingle<K extends ModelColumnNamesWithId<M>>(column: K): QueryArray<M, M[K]>;
order(orders?: string): QueryArray<M, T>;
vector_order(order: VectorOrderOption): QueryArray<M, T>;
group<G extends ModelColumnNamesWithId<M>, F>(group_by: G | G[], fields?: F): QueryArray<M, {
[field in keyof F]: number;
} & Pick<M, G>>;
Expand Down Expand Up @@ -238,6 +243,7 @@ declare class Query<M extends BaseModel, T = M> implements QuerySingle<M, T>, Qu
* Specifies orders of result
*/
order(orders?: string): this;
vector_order(order: VectorOrderOption): this;
/**
* Groups result records
*/
Expand Down
8 changes: 8 additions & 0 deletions packages/cormo/lib/cjs/query.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,13 @@ class Query {
this._options.orders = orders;
return this;
}
vector_order(order) {
if (!this._current_if) {
return this;
}
this._options.vector_order = order;
return this;
}
group(group_by, fields) {
if (!this._current_if) {
return this;
Expand Down Expand Up @@ -595,6 +602,7 @@ class Query {
node: this._options.node,
index_hint: this._options.index_hint,
orders,
vector_order: this._options.vector_order,
skip: this._options.skip,
transaction: this._options.transaction,
distinct: this._options.distinct,
Expand Down
24 changes: 20 additions & 4 deletions packages/cormo/lib/cjs/types.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,22 @@ export interface CormoTypesGeoPointConstructor {
(): CormoTypesGeoPoint;
}
declare const CormoTypesGeoPoint: CormoTypesGeoPointConstructor;
/**
* Represents a vector, used in model schemas.
*
* This type is supported only in PostgreSQL with pgvector
* @namespace types
* @class Vector
*/
export interface CormoTypesVector {
_type: 'vector';
dimension?: number;
}
export interface CormoTypesVectorConstructor {
new (dimension?: number): CormoTypesVector;
(dimension?: number): CormoTypesVector;
}
declare const CormoTypesVector: CormoTypesVectorConstructor;
/**
* Represents a date, used in model schemas.
* @namespace types
Expand Down Expand Up @@ -152,17 +168,17 @@ export interface CormoTypesBlobConstructor {
(): CormoTypesBlob;
}
declare const CormoTypesBlob: CormoTypesBlobConstructor;
export type ColumnTypeInternal = CormoTypesString | CormoTypesNumber | CormoTypesBoolean | CormoTypesDate | CormoTypesObject | CormoTypesInteger | CormoTypesBigInteger | CormoTypesGeoPoint | CormoTypesRecordID | CormoTypesText | CormoTypesBlob;
export type ColumnTypeInternalConstructor = CormoTypesStringConstructor | CormoTypesNumberConstructor | CormoTypesBooleanConstructor | CormoTypesDateConstructor | CormoTypesObjectConstructor | CormoTypesIntegerConstructor | CormoTypesBigIntegerConstructor | CormoTypesGeoPointConstructor | CormoTypesRecordIDConstructor | CormoTypesTextConstructor | CormoTypesBlobConstructor;
export type ColumnTypeInternal = CormoTypesString | CormoTypesNumber | CormoTypesBoolean | CormoTypesDate | CormoTypesObject | CormoTypesInteger | CormoTypesBigInteger | CormoTypesGeoPoint | CormoTypesVector | CormoTypesRecordID | CormoTypesText | CormoTypesBlob;
export type ColumnTypeInternalConstructor = CormoTypesStringConstructor | CormoTypesNumberConstructor | CormoTypesBooleanConstructor | CormoTypesDateConstructor | CormoTypesObjectConstructor | CormoTypesIntegerConstructor | CormoTypesBigIntegerConstructor | CormoTypesGeoPointConstructor | CormoTypesVectorConstructor | CormoTypesRecordIDConstructor | CormoTypesTextConstructor | CormoTypesBlobConstructor;
type ColumnTypeNativeConstructor = StringConstructor | NumberConstructor | BooleanConstructor | DateConstructor | ObjectConstructor;
type ColumnTypeString = 'string' | 'number' | 'boolean' | 'date' | 'object' | 'integer' | 'biginteger' | 'geopoint' | 'recordid' | 'text' | 'blob';
type ColumnTypeString = 'string' | 'number' | 'boolean' | 'date' | 'object' | 'integer' | 'biginteger' | 'geopoint' | 'vector' | 'recordid' | 'text' | 'blob';
export type ColumnType = ColumnTypeInternal | ColumnTypeInternalConstructor | ColumnTypeNativeConstructor | ColumnTypeString;
/**
* Converts JavaScript built-in class to CORMO type
* @private
*/
declare function _toCORMOType(type: ColumnType): ColumnTypeInternal;
export { CormoTypesString as String, CormoTypesNumber as Number, CormoTypesBoolean as Boolean, CormoTypesInteger as Integer, CormoTypesBigInteger as BigInteger, CormoTypesGeoPoint as GeoPoint, CormoTypesDate as Date, CormoTypesObject as Object, CormoTypesRecordID as RecordID, CormoTypesText as Text, CormoTypesBlob as Blob, _toCORMOType, };
export { CormoTypesString as String, CormoTypesNumber as Number, CormoTypesBoolean as Boolean, CormoTypesInteger as Integer, CormoTypesBigInteger as BigInteger, CormoTypesGeoPoint as GeoPoint, CormoTypesVector as Vector, CormoTypesDate as Date, CormoTypesObject as Object, CormoTypesRecordID as RecordID, CormoTypesText as Text, CormoTypesBlob as Blob, _toCORMOType, };
/**
* A pseudo type represents a record's unique identifier.
*
Expand Down
12 changes: 11 additions & 1 deletion packages/cormo/lib/cjs/types.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
* @namespace cormo
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.Blob = exports.Text = exports.RecordID = exports.Object = exports.Date = exports.GeoPoint = exports.BigInteger = exports.Integer = exports.Boolean = exports.Number = exports.String = void 0;
exports.Blob = exports.Text = exports.RecordID = exports.Object = exports.Date = exports.Vector = exports.GeoPoint = exports.BigInteger = exports.Integer = exports.Boolean = exports.Number = exports.String = void 0;
exports._toCORMOType = _toCORMOType;
const CormoTypesString = function (length) {
if (!(this instanceof CormoTypesString)) {
Expand Down Expand Up @@ -50,6 +50,14 @@ const CormoTypesGeoPoint = function () {
this.toString = () => 'geopoint';
};
exports.GeoPoint = CormoTypesGeoPoint;
const CormoTypesVector = function (dimension) {
if (!(this instanceof CormoTypesVector)) {
return new CormoTypesVector(dimension);
}
this.dimension = dimension;
this.toString = () => (this.dimension ? `vector(${this.dimension})` : 'vector');
};
exports.Vector = CormoTypesVector;
const CormoTypesDate = function () {
if (!(this instanceof CormoTypesDate)) {
return new CormoTypesDate();
Expand Down Expand Up @@ -108,6 +116,8 @@ function _toCORMOType(type) {
return new CormoTypesBigInteger();
case 'geopoint':
return new CormoTypesGeoPoint();
case 'vector':
return new CormoTypesVector();
case 'date':
return new CormoTypesDate();
case 'object':
Expand Down
2 changes: 2 additions & 0 deletions packages/cormo/lib/esm/adapters/base.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,11 @@ export interface Schemas {
[table_name: string]: any;
};
}
export type VectorOrderOption = Record<string, Record<string, number[] | undefined> | undefined>;
export interface AdapterFindOptions {
lean: boolean;
orders: any[];
vector_order?: VectorOrderOption;
near?: any;
select?: string[];
select_raw?: string[];
Expand Down
Loading

0 comments on commit 74ccc0f

Please sign in to comment.