From 278347987673cc3cff0b7f5d72b56cac1d9224b3 Mon Sep 17 00:00:00 2001
From: Mitchell Adair <madair@singlestore.com>
Date: Tue, 23 Jul 2024 14:30:37 -0400
Subject: [PATCH] add vector data type, operator experessions

---
 .../src/singlestore-core/columns/vector.ts    | 67 +++++++++++++++++++
 .../src/singlestore-core/expressions.ts       | 10 +++
 2 files changed, 77 insertions(+)
 create mode 100644 drizzle-orm/src/singlestore-core/columns/vector.ts

diff --git a/drizzle-orm/src/singlestore-core/columns/vector.ts b/drizzle-orm/src/singlestore-core/columns/vector.ts
new file mode 100644
index 000000000..ad1a09780
--- /dev/null
+++ b/drizzle-orm/src/singlestore-core/columns/vector.ts
@@ -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)
+}
\ No newline at end of file
diff --git a/drizzle-orm/src/singlestore-core/expressions.ts b/drizzle-orm/src/singlestore-core/expressions.ts
index 6d4284d18..aa77fbd48 100644
--- a/drizzle-orm/src/singlestore-core/expressions.ts
+++ b/drizzle-orm/src/singlestore-core/expressions.ts
@@ -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)}`
+}