From 4c3076f2a474fd66233a1a0f9fe0331423aea546 Mon Sep 17 00:00:00 2001 From: huangzhuo1949 <167434202+huangzhuo1949@users.noreply.github.com> Date: Wed, 22 Jan 2025 17:07:18 +0800 Subject: [PATCH] feat: add pg vector index (#12338) Co-authored-by: huangzhuo --- api/core/rag/datasource/vdb/pgvector/pgvector.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/api/core/rag/datasource/vdb/pgvector/pgvector.py b/api/core/rag/datasource/vdb/pgvector/pgvector.py index de443ba580..c8a1e4f90c 100644 --- a/api/core/rag/datasource/vdb/pgvector/pgvector.py +++ b/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)