Add vector support

Fixes typeorm#10056
Lunix-AI · Mar 24, 2024 · 08ff7cd · 08ff7cd
1 parent 83567f5
commit 08ff7cd
Show file tree

Hide file tree

Showing 6 changed files with 448 additions and 4 deletions.
diff --git a/docs/entities.md b/docs/entities.md
@@ -393,6 +393,8 @@ or
 `tsrange`, `tstzrange`, `daterange`, `int4multirange`, `int8multirange`, `nummultirange`,
 `tsmultirange`, `tstzmultirange`, `multidaterange`, `geometry`, `geography`, `cube`, `ltree`
 
+> Note: `vector` is also available when the [pgvector](https://github.com/pgvector/pgvector) extension is installed.
+
 ### Column types for `cockroachdb`
 
 `array`, `bool`, `boolean`, `bytes`, `bytea`, `blob`, `date`, `numeric`, `decimal`, `dec`, `float`,

diff --git a/src/driver/postgres/PostgresDriver.ts b/src/driver/postgres/PostgresDriver.ts
@@ -190,6 +190,7 @@ export class PostgresDriver implements Driver {
         "geography",
         "cube",
         "ltree",
+        "vector",
     ]
 
     /**
@@ -213,6 +214,7 @@ export class PostgresDriver implements Driver {
         "bit",
         "varbit",
         "bit varying",
+        "vector",
     ]
 
     /**
@@ -421,6 +423,7 @@ export class PostgresDriver implements Driver {
             hasCubeColumns,
             hasGeometryColumns,
             hasLtreeColumns,
+            hasVectorColumns,
             hasExclusionConstraints,
         } = extensionsMetadata
 
@@ -500,6 +503,18 @@ export class PostgresDriver implements Driver {
                     "At least one of the entities has a ltree column, but the 'ltree' extension cannot be installed automatically. Please install it manually using superuser rights",
                 )
             }
+        if (hasVectorColumns)
+            try {
+                await this.executeQuery(
+                    connection,
+                    `CREATE EXTENSION IF NOT EXISTS "vector"`,
+                )
+            } catch (_) {
+                logger.log(
+                    "warn",
+                    "At least one of the entities has a vector column, but the 'vector' extension cannot be installed automatically. Please install it manually using superuser rights",
+                )
+            }
         if (hasExclusionConstraints)
             try {
                 // The btree_gist extension provides operator support in PostgreSQL exclusion constraints
@@ -568,6 +583,14 @@ export class PostgresDriver implements Driver {
                 )
             },
         )
+        const hasVectorColumns = this.connection.entityMetadatas.some(
+            (metadata) => {
+                return (
+                    metadata.columns.filter((column) => column.type === "vector")
+                        .length > 0
+                )
+            },
+        )
         const hasExclusionConstraints = this.connection.entityMetadatas.some(
             (metadata) => {
                 return metadata.exclusions.length > 0
@@ -581,6 +604,7 @@ export class PostgresDriver implements Driver {
             hasCubeColumns,
             hasGeometryColumns,
             hasLtreeColumns,
+            hasVectorColumns,
             hasExclusionConstraints,
             hasExtensions:
                 hasUuidColumns ||
@@ -589,6 +613,7 @@ export class PostgresDriver implements Driver {
                 hasGeometryColumns ||
                 hasCubeColumns ||
                 hasLtreeColumns ||
+                hasVectorColumns ||
                 hasExclusionConstraints,
         }
     }

diff --git a/src/driver/postgres/PostgresQueryRunner.ts b/src/driver/postgres/PostgresQueryRunner.ts
@@ -601,7 +601,7 @@ export class PostgresQueryRunner
                 downQueries.push(this.dropIndexSql(table, index))
             })
         }
-        
+
         if (table.comment) {
             upQueries.push(new Query("COMMENT ON TABLE " + this.escapePath(table) + " IS '" + table.comment + "'"));
             downQueries.push(new Query("COMMENT ON TABLE " + this.escapePath(table) + " IS NULL"));
@@ -3542,8 +3542,8 @@ export class PostgresQueryRunner
                             // check if column has user-defined data type.
                             // NOTE: if ENUM type defined with "array:true" it comes with ARRAY type instead of USER-DEFINED
                             if (
-                                dbColumn["data_type"] === "USER-DEFINED" ||
-                                dbColumn["data_type"] === "ARRAY"
+                                (dbColumn["data_type"] === "USER-DEFINED" ||
+                                    dbColumn["data_type"] === "ARRAY") && dbColumn["udt_name"] !== "vector"
                             ) {
                                 const { name } =
                                     await this.getUserDefinedTypeName(
@@ -3595,6 +3595,20 @@ export class PostgresQueryRunner
                                 }
                             }
 
+                            if (
+                                dbColumn["data_type"] === "USER-DEFINED" &&
+                                tableColumn.type === "vector"
+                            ) {
+                                // vector length is from the format_type e.g. `format_type = vector(n)`
+                                const match = /vector\((\d+)\)/.exec(
+                                    dbColumn["format_type"],
+                                )
+
+                                if (match) {
+                                    tableColumn.length = match[1]
+                                }
+                            }
+
                             if (
                                 tableColumn.type === "geometry" ||
                                 tableColumn.type === "geography"
@@ -4744,7 +4758,7 @@ export class PostgresQueryRunner
 
         newComment = this.escapeComment(newComment)
         const comment = this.escapeComment(table.comment)
-        
+
         if (newComment === comment) {
             return
         }

diff --git a/src/driver/types/ColumnTypes.ts b/src/driver/types/ColumnTypes.ts
@@ -75,6 +75,7 @@ export type WithLengthColumnType =
     | "binary" // mssql
     | "varbinary" // mssql, sap
     | "string" // cockroachdb, spanner
+    | "vector" // postgres pgvector
 
 export type WithWidthColumnType =
     | "tinyint" // mysql

diff --git a/test/github-issues/10056/entity/message.ts b/test/github-issues/10056/entity/message.ts
@@ -0,0 +1,13 @@
+import { Column, Entity, PrimaryGeneratedColumn } from "../../../../src"
+
+@Entity({
+    name: "message",
+})
+export class Message {
+    @PrimaryGeneratedColumn()
+    id: number
+
+    // create a vector embedding with 5 dimensions
+    @Column("vector", { length: 5 })
+    embedding: string;
+}