59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
"""add hnsw index to documents table
|
|
|
|
Revision ID: 66e63cf2cf77
|
|
Revises: 20f89a421aff
|
|
Create Date: 2025-05-19 17:00:18.151735
|
|
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
from alembic import op
|
|
from sqlalchemy import text
|
|
|
|
from src.config import settings
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "66e63cf2cf77"
|
|
down_revision: str | None = "20f89a421aff"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
schema = settings.DB.SCHEMA
|
|
|
|
|
|
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"
|
|
)
|