89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
|
"""Optimize HNSW index parameters for better vector search recall.
|
||
|
|
|
||
|
|
Recreates existing HNSW indices with tuned parameters:
|
||
|
|
- ef_construction=128 (default 64, higher = better index quality, slower build)
|
||
|
|
- m=16 (default 16, higher = more memory, better recall)
|
||
|
|
|
||
|
|
IVFFlat Alternative (B-VEC-IVF):
|
||
|
|
-----------------------------
|
||
|
|
comm_messages uses IVFFlat with lists=100 (migration 0035).
|
||
|
|
Rule of thumb for IVFFlat: lists = sqrt(rows)
|
||
|
|
~10k rows → lists ≈ 100
|
||
|
|
~50k rows → lists ≈ 224
|
||
|
|
~100k rows → lists ≈ 316
|
||
|
|
IVFFlat builds faster but HNSW has better recall.
|
||
|
|
To switch: DROP INDEX + CREATE INDEX ... USING hnsw (embedding vector_cosine_ops)
|
||
|
|
WITH (ef_construction=128, m=16)
|
||
|
|
Config: vector_index_type setting in app/config.py (default 'hnsw', alternative 'ivfflat').
|
||
|
|
|
||
|
|
Revision ID: 0118
|
||
|
|
Revises: 0117
|
||
|
|
"""
|
||
|
|
|
||
|
|
from alembic import op
|
||
|
|
import sqlalchemy as sa
|
||
|
|
|
||
|
|
revision = "0118"
|
||
|
|
down_revision = "0117"
|
||
|
|
branch_labels = None
|
||
|
|
depends_on = None
|
||
|
|
|
||
|
|
# Optimized HNSW parameters
|
||
|
|
EF_CONSTRUCTION = 128
|
||
|
|
M = 16
|
||
|
|
|
||
|
|
# Tables with HNSW indices (from migration 0104)
|
||
|
|
# Format: (table_name, index_name)
|
||
|
|
HNSW_TABLES = [
|
||
|
|
("contacts", "ix_contacts_embedding"),
|
||
|
|
("mails", "ix_mails_embedding"),
|
||
|
|
("files", "ix_files_embedding"),
|
||
|
|
("calendar_entries", "ix_calendar_entries_embedding"),
|
||
|
|
("tags", "ix_tags_embedding"),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
conn = op.get_bind()
|
||
|
|
|
||
|
|
for table_name, index_name in HNSW_TABLES:
|
||
|
|
# Check if table exists
|
||
|
|
table_exists = conn.execute(sa.text(
|
||
|
|
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = :t)"
|
||
|
|
), {"t": table_name}).scalar()
|
||
|
|
|
||
|
|
if not table_exists:
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Drop existing HNSW index (regardless of parameters)
|
||
|
|
op.execute(f"DROP INDEX IF EXISTS {index_name}")
|
||
|
|
|
||
|
|
# Recreate with optimized parameters
|
||
|
|
op.execute(
|
||
|
|
f"CREATE INDEX IF NOT EXISTS {index_name} "
|
||
|
|
f"ON {table_name} USING hnsw (embedding vector_cosine_ops) "
|
||
|
|
f"WITH (ef_construction={EF_CONSTRUCTION}, m={M})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
"""Recreate HNSW indices with default parameters (no WITH clause)."""
|
||
|
|
conn = op.get_bind()
|
||
|
|
|
||
|
|
for table_name, index_name in HNSW_TABLES:
|
||
|
|
table_exists = conn.execute(sa.text(
|
||
|
|
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = :t)"
|
||
|
|
), {"t": table_name}).scalar()
|
||
|
|
|
||
|
|
if not table_exists:
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Drop optimized index
|
||
|
|
op.execute(f"DROP INDEX IF EXISTS {index_name}")
|
||
|
|
|
||
|
|
# Recreate with default parameters (no WITH clause = pgvector defaults)
|
||
|
|
op.execute(
|
||
|
|
f"CREATE INDEX IF NOT EXISTS {index_name} "
|
||
|
|
f"ON {table_name} USING hnsw (embedding vector_cosine_ops)"
|
||
|
|
)
|