From 7d83678804353fa680fa34b3743184b1ace65f7d Mon Sep 17 00:00:00 2001 From: Mini256 Date: Fri, 23 Feb 2024 17:21:17 +0800 Subject: [PATCH 1/3] feat: add vector function support --- src/index.ts | 26 +++++++++++++++++++++++--- src/vector.ts | 38 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) create mode 100644 src/vector.ts diff --git a/src/index.ts b/src/index.ts index 84dbd6b..9c19452 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,3 @@ -import { connect, Connection, Tx, Config, FullResult } from '@tidbcloud/serverless'; - import { CompiledQuery, DatabaseConnection, @@ -8,11 +6,22 @@ import { Driver, Kysely, MysqlAdapter, - MysqlQueryCompiler, MysqlIntrospector, + MysqlQueryCompiler, QueryCompiler, QueryResult, } from 'kysely'; +import { Config, Connection, FullResult, Tx, connect } from '@tidbcloud/serverless'; +import { + cosineDistance, + cosineSimilarity, + innerProduct, + l1Distance, + l2Distance, + negativeInnerProduct, + vectorFromSql, + vectorToSql, +} from './vector'; /** * Config for the TiDB Serverless dialect. @@ -171,3 +180,14 @@ class TiDBServerlessTransaction { await this.#tx.rollback(); } } + +export { + cosineDistance, + cosineSimilarity, + innerProduct, + l1Distance, + l2Distance, + negativeInnerProduct, + vectorFromSql, + vectorToSql, +}; diff --git a/src/vector.ts b/src/vector.ts new file mode 100644 index 0000000..cf020ba --- /dev/null +++ b/src/vector.ts @@ -0,0 +1,38 @@ +import { RawBuilder, sql } from 'kysely'; + +type VectorLike = Float32Array | number[]; + +export function l1Distance(column: string, value: VectorLike): RawBuilder { + return sql`VEC_L1_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function l2Distance(column: string, value: VectorLike): RawBuilder { + return sql`VEC_L2_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function negativeInnerProduct(column: string, value: VectorLike): RawBuilder { + return sql`VEC_NEGATIVE_INNER_PRODUCT(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function innerProduct(column: string, value: VectorLike): RawBuilder { + return sql`1 - VEC_NEGATIVE_INNER_PRODUCT(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function cosineDistance(column: string, value: VectorLike): RawBuilder { + return sql`VEC_COSINE_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function cosineSimilarity(column: any, value: VectorLike): RawBuilder { + return sql`1 - VEC_COSINE_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +} + +export function vectorFromSql(value: string) { + return value + .substring(1, value.length - 1) + .split(',') + .map((v) => parseFloat(v)); +} + +export function vectorToSql(vector: VectorLike) { + return `[${vector.join(',')}]`; +} From 1a8e8dce117daa47ab88bef10fdd30ead07d1000 Mon Sep 17 00:00:00 2001 From: Mini256 Date: Thu, 29 Feb 2024 23:27:13 +0800 Subject: [PATCH 2/3] more type-safe impl --- src/vector.ts | 38 +++++++++++++++++++++++++------------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/src/vector.ts b/src/vector.ts index cf020ba..9990047 100644 --- a/src/vector.ts +++ b/src/vector.ts @@ -1,29 +1,41 @@ -import { RawBuilder, sql } from 'kysely'; +import { ExpressionBuilder, ReferenceExpression } from 'kysely'; type VectorLike = Float32Array | number[]; -export function l1Distance(column: string, value: VectorLike): RawBuilder { - return sql`VEC_L1_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +export function l1Distance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.fn('VEC_L1_DISTANCE', [ + column, + eb => eb.val(vectorToSql(value))], + ); } -export function l2Distance(column: string, value: VectorLike): RawBuilder { - return sql`VEC_L2_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +export function l2Distance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.fn('VEC_L2_DISTANCE', [ + column, + eb => eb.val(vectorToSql(value))], + ); } -export function negativeInnerProduct(column: string, value: VectorLike): RawBuilder { - return sql`VEC_NEGATIVE_INNER_PRODUCT(${sql.ref(column)}, ${vectorToSql(value)})`; +export function negativeInnerProduct = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.fn('VEC_NEGATIVE_INNER_PRODUCT', [ + column, + eb => eb.val(vectorToSql(value))], + ); } -export function innerProduct(column: string, value: VectorLike): RawBuilder { - return sql`1 - VEC_NEGATIVE_INNER_PRODUCT(${sql.ref(column)}, ${vectorToSql(value)})`; +export function innerProduct = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.neg(negativeInnerProduct(eb, column, value)); } -export function cosineDistance(column: string, value: VectorLike): RawBuilder { - return sql`VEC_COSINE_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +export function cosineDistance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.fn('VEC_COSINE_DISTANCE', [ + column, + eb => eb.val(vectorToSql(value))], + ); } -export function cosineSimilarity(column: any, value: VectorLike): RawBuilder { - return sql`1 - VEC_COSINE_DISTANCE(${sql.ref(column)}, ${vectorToSql(value)})`; +export function cosineSimilarity = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb(eb.lit(1), '-', cosineDistance(eb, column, value)); } export function vectorFromSql(value: string) { From 912a4521a09c3bd34bb74d43a3e1df877d4a9a0a Mon Sep 17 00:00:00 2001 From: Mini256 Date: Thu, 29 Feb 2024 23:29:20 +0800 Subject: [PATCH 3/3] fix ci --- src/vector.ts | 68 ++++++++++++++++++++++++++++++--------------------- 1 file changed, 40 insertions(+), 28 deletions(-) diff --git a/src/vector.ts b/src/vector.ts index 9990047..1137988 100644 --- a/src/vector.ts +++ b/src/vector.ts @@ -2,39 +2,51 @@ import { ExpressionBuilder, ReferenceExpression } from 'kysely'; type VectorLike = Float32Array | number[]; -export function l1Distance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { - return eb.fn('VEC_L1_DISTANCE', [ - column, - eb => eb.val(vectorToSql(value))], - ); -} - -export function l2Distance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { - return eb.fn('VEC_L2_DISTANCE', [ - column, - eb => eb.val(vectorToSql(value))], - ); -} - -export function negativeInnerProduct = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { - return eb.fn('VEC_NEGATIVE_INNER_PRODUCT', [ - column, - eb => eb.val(vectorToSql(value))], - ); -} - -export function innerProduct = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { +export function l1Distance = ReferenceExpression>( + eb: ExpressionBuilder, + column: RE, + value: VectorLike +) { + return eb.fn('VEC_L1_DISTANCE', [column, (eb) => eb.val(vectorToSql(value))]); +} + +export function l2Distance = ReferenceExpression>( + eb: ExpressionBuilder, + column: RE, + value: VectorLike +) { + return eb.fn('VEC_L2_DISTANCE', [column, (eb) => eb.val(vectorToSql(value))]); +} + +export function negativeInnerProduct< + DB, + TB extends keyof DB, + RE extends ReferenceExpression = ReferenceExpression +>(eb: ExpressionBuilder, column: RE, value: VectorLike) { + return eb.fn('VEC_NEGATIVE_INNER_PRODUCT', [column, (eb) => eb.val(vectorToSql(value))]); +} + +export function innerProduct = ReferenceExpression>( + eb: ExpressionBuilder, + column: RE, + value: VectorLike +) { return eb.neg(negativeInnerProduct(eb, column, value)); } -export function cosineDistance = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { - return eb.fn('VEC_COSINE_DISTANCE', [ - column, - eb => eb.val(vectorToSql(value))], - ); +export function cosineDistance = ReferenceExpression>( + eb: ExpressionBuilder, + column: RE, + value: VectorLike +) { + return eb.fn('VEC_COSINE_DISTANCE', [column, (eb) => eb.val(vectorToSql(value))]); } -export function cosineSimilarity = ReferenceExpression>(eb: ExpressionBuilder, column: RE, value: VectorLike) { +export function cosineSimilarity = ReferenceExpression>( + eb: ExpressionBuilder, + column: RE, + value: VectorLike +) { return eb(eb.lit(1), '-', cosineDistance(eb, column, value)); }