Skip to content

Commit

Permalink
add vector data type, operator experessions
Browse files Browse the repository at this point in the history
  • Loading branch information
mitchwadair committed Jul 23, 2024
1 parent 21529bb commit 2783479
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 0 deletions.
67 changes: 67 additions & 0 deletions drizzle-orm/src/singlestore-core/columns/vector.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
import { ColumnBaseConfig } from "~/column";
import { ColumnBuilderBaseConfig, ColumnBuilderRuntimeConfig, MakeColumnConfig } from "~/column-builder";
import { entityKind } from "~/entity";
import { SingleStoreColumn, SingleStoreColumnBuilder } from "./common";
import { AnySingleStoreTable } from "~/singlestore-core/table";

export type SingleStoreVectorBuilderInitial<TName extends string> = SingleStoreVectorBuilder<{
name: TName;
dataType: "array";
columnType: "SingleStoreVector";
data: Array<number>;
driverParam: Array<number>;
enumValues: undefined;
generated: undefined;
}>

export class SingleStoreVectorBuilder<T extends ColumnBuilderBaseConfig<"array", "SingleStoreVector">> extends SingleStoreColumnBuilder<T, SingleStoreVectorConfig> {
static readonly [entityKind]: string = "SingleStoreVectorBuilder"

constructor(name: T["name"], config: SingleStoreVectorConfig) {
super(name, "array", "SingleStoreVector")
this.config.dimensions = config.dimensions
this.config.elementType = config.elementType
}

/** @internal */
override build<TTableName extends string>(table: AnySingleStoreTable<{name: TTableName}>): SingleStoreVector<MakeColumnConfig<T, TTableName>> {
return new SingleStoreVector(table, this.config as ColumnBuilderRuntimeConfig<any, any>);
}
}

export class SingleStoreVector<T extends ColumnBaseConfig<"array", "SingleStoreVector">> extends SingleStoreColumn<T> {
static readonly [entityKind]: string = "SingleStoreVector"

readonly dimensions: number;
readonly elementType: ElementType | undefined;

constructor(table: AnySingleStoreTable<{name: T["tableName"]}>, config: SingleStoreVectorBuilder<T>["config"]) {
super(table, config)
this.dimensions = config.dimensions
this.elementType = config.elementType
}

getSQLType(): string {
const et = this.elementType === undefined ? "" : `, ${this.elementType}`;
return `vector(${this.dimensions}${et})`
}

override mapToDriverValue(value: Array<number>) {
return JSON.stringify(value)
}

override mapFromDriverValue(value: string): Array<number> {
return JSON.parse(value)
}
}

type ElementType = "I8" | "I16" | "I32" | "I64" | "F32" | "F64";

export interface SingleStoreVectorConfig {
dimensions: number;
elementType: ElementType;
}

export function vector<TName extends string>(name: TName, config: SingleStoreVectorConfig): SingleStoreVectorBuilderInitial<TName> {
return new SingleStoreVectorBuilder(name, config)
}
10 changes: 10 additions & 0 deletions drizzle-orm/src/singlestore-core/expressions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,3 +23,13 @@ export function substring(
chunks.push(sql`)`);
return sql.join(chunks);
}

// Vectors

export function dotProduct(column: SingleStoreColumn | SQL.Aliased, value: Array<number>) {
return sql`${column} <*> ${JSON.stringify(value)}`
}

export function euclideanDistance(column: SingleStoreColumn | SQL.Aliased, value: Array<number>) {
return sql`${column} <-> ${JSON.stringify(value)}`
}

0 comments on commit 2783479

Please sign in to comment.