diff --git a/migrations/versions/66e63cf2cf77_add_indexes_to_documents_table.py b/migrations/versions/66e63cf2cf77_add_indexes_to_documents_table.py new file mode 100644 index 00000000..7221496f --- /dev/null +++ b/migrations/versions/66e63cf2cf77_add_indexes_to_documents_table.py @@ -0,0 +1,48 @@ +"""add hnsw index to documents table + +Revision ID: 66e63cf2cf77 +Revises: 20f89a421aff +Create Date: 2025-05-19 17:00:18.151735 + +""" +from typing import Sequence, Union +from os import getenv + +from alembic import op +from sqlalchemy import text + + +# revision identifiers, used by Alembic. +revision: str = '66e63cf2cf77' +down_revision: Union[str, None] = '20f89a421aff' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None +schema = getenv("DATABASE_SCHEMA", "public") + + +def upgrade() -> None: + + # Create HNSW index on the embedding column for the documents table for cosine distance + # Parameters: + # - m: max number of connections (edges) per node (default=16) + # - ef_construction: size of the candidate list during index construction (default=64) + print(f"Creating HNSW index idx_documents_embedding_hnsw on {schema}.documents table") + try: + op.execute( + text( + f""" + CREATE INDEX idx_documents_embedding_hnsw ON {schema}.documents + USING hnsw (embedding vector_cosine_ops) + WITH (m=16, ef_construction=64); + """ + ) + ) + print(f"HNSW index idx_documents_embedding_hnsw created on {schema}.documents table") + except Exception as e: + print(f"Error creating HNSW index idx_documents_embedding_hnsw on {schema}.documents table: {e}") + + +def downgrade() -> None: + print(f"Dropping HNSW index idx_documents_embedding_hnsw from {schema}.documents table") + op.execute(text(f"DROP INDEX IF EXISTS {schema}.idx_documents_embedding_hnsw;")) + print(f"HNSW index idx_documents_embedding_hnsw dropped from {schema}.documents table") diff --git a/src/models.py b/src/models.py index 3811e93a..71587653 100644 --- a/src/models.py +++ b/src/models.py @@ -274,6 +274,14 @@ class Document(Base): CheckConstraint("length(public_id) = 21", name="public_id_length"), CheckConstraint("length(content) <= 65535", name="content_length"), CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), + # HNSW index on embedding column + Index( + "idx_documents_embedding_hnsw", + "embedding", + postgresql_using="hnsw", # HNSW index type + postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters + postgresql_ops={"embedding": "vector_cosine_ops"}, # Cosine distance operator + ), )