Преглед на файлове

feat: add pg vector index (#12338)

Co-authored-by: huangzhuo <huangzhuo1@xiaomi.com>
huangzhuo1949 преди 3 месеца
родител
ревизия
4c3076f2a4
променени са 1 файла, в които са добавени 9 реда и са изтрити 1 реда
  1. 9 1
      api/core/rag/datasource/vdb/pgvector/pgvector.py

+ 9 - 1
api/core/rag/datasource/vdb/pgvector/pgvector.py

@@ -57,6 +57,11 @@ CREATE TABLE IF NOT EXISTS {table_name} (
 ) using heap;
 """
 
+SQL_CREATE_INDEX = """
+CREATE INDEX IF NOT EXISTS embedding_cosine_v1_idx ON {table_name} 
+USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);
+"""
+
 
 class PGVector(BaseVector):
     def __init__(self, collection_name: str, config: PGVectorConfig):
@@ -205,7 +210,10 @@ class PGVector(BaseVector):
             with self._get_cursor() as cur:
                 cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
                 cur.execute(SQL_CREATE_TABLE.format(table_name=self.table_name, dimension=dimension))
-                # TODO: create index https://github.com/pgvector/pgvector?tab=readme-ov-file#indexing
+                # PG hnsw index only support 2000 dimension or less
+                # ref: https://github.com/pgvector/pgvector?tab=readme-ov-file#indexing
+                if dimension <= 2000:
+                    cur.execute(SQL_CREATE_INDEX.format(table_name=self.table_name))
             redis_client.set(collection_exist_cache_key, 1, ex=3600)