add indexes for query documents (#111)

* add indexes for query documents

* rm concurrently from migration

* remove compound index

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")

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
+        ),
     )