perf: persist map-nav scoring index for large corpora

apps/api/alembic/versions/1d2e3f4a5b6c_add_document_map_units.py

@@ -0,0 +1,122 @@
+"""Add revision-pinned map-nav score units."""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = "1d2e3f4a5b6c"
+down_revision = "0c1d2e3f4a5b"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    if not inspector.has_table("document_map_unit_indexes"):
+        op.create_table(
+            "document_map_unit_indexes",
+            sa.Column("id", sa.String(length=100), nullable=False),
+            sa.Column("document_id", sa.String(length=36), nullable=False),
+            sa.Column("job_result_id", sa.String(length=36), nullable=False),
+            sa.Column("format_version", sa.Integer(), nullable=False),
+            sa.Column("unit_count", sa.Integer(), nullable=False),
+            sa.Column("token_count", sa.Integer(), nullable=False),
+            sa.Column("created_at", sa.DateTime(), nullable=False),
+            sa.ForeignKeyConstraint(
+                ["document_id"], ["documents.document_id"], ondelete="CASCADE"
+            ),
+            sa.ForeignKeyConstraint(
+                ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
+            ),
+            sa.PrimaryKeyConstraint("id"),
+            sa.UniqueConstraint(
+                "document_id",
+                "job_result_id",
+                name="uq_document_map_unit_indexes_revision",
+            ),
+        )
+    if not inspector.has_table("document_map_units"):
+        op.create_table(
+            "document_map_units",
+            sa.Column("id", sa.String(length=160), nullable=False),
+            sa.Column("document_id", sa.String(length=36), nullable=False),
+            sa.Column("job_result_id", sa.String(length=36), nullable=False),
+            sa.Column("unit_id", sa.String(length=128), nullable=False),
+            sa.Column("section_id", sa.String(length=36), nullable=False),
+            sa.Column("unit_kind", sa.String(length=32), nullable=False),
+            sa.Column("path_token_count", sa.Integer(), nullable=False),
+            sa.Column("content_token_count", sa.Integer(), nullable=False),
+            sa.Column("term_search_text_lower", sa.Text(), nullable=False),
+            sa.Column("sort_order", sa.Integer(), nullable=False),
+            sa.Column("created_at", sa.DateTime(), nullable=False),
+            sa.ForeignKeyConstraint(
+                ["document_id"], ["documents.document_id"], ondelete="CASCADE"
+            ),
+            sa.ForeignKeyConstraint(
+                ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
+            ),
+            sa.PrimaryKeyConstraint("id"),
+        )
+    if not inspector.has_table("document_map_unit_tokens"):
+        op.create_table(
+            "document_map_unit_tokens",
+            sa.Column("id", sa.String(length=36), nullable=False),
+            sa.Column("map_unit_id", sa.String(length=160), nullable=False),
+            sa.Column("channel", sa.String(length=16), nullable=False),
+            sa.Column("token", sa.Text(), nullable=False),
+            sa.Column("token_hash", sa.String(length=64), nullable=False),
+            sa.Column("frequency", sa.Integer(), nullable=False),
+            sa.ForeignKeyConstraint(
+                ["map_unit_id"], ["document_map_units.id"], ondelete="CASCADE"
+            ),
+            sa.PrimaryKeyConstraint("id"),
+        )
+    inspector = sa.inspect(bind)
+    indexes = {
+        item["name"] for item in inspector.get_indexes("document_map_unit_indexes")
+    }
+    if "idx_document_map_unit_indexes_revision" not in indexes:
+        op.create_index(
+            "idx_document_map_unit_indexes_revision",
+            "document_map_unit_indexes",
+            ["document_id", "job_result_id"],
+        )
+    indexes = {item["name"] for item in inspector.get_indexes("document_map_units")}
+    if "idx_document_map_units_revision_order" not in indexes:
+        op.create_index(
+            "idx_document_map_units_revision_order",
+            "document_map_units",
+            ["document_id", "job_result_id", "sort_order", "unit_id"],
+        )
+    if "idx_document_map_units_section" not in indexes:
+        op.create_index(
+            "idx_document_map_units_section", "document_map_units", ["section_id"]
+        )
+    indexes = {
+        item["name"] for item in inspector.get_indexes("document_map_unit_tokens")
+    }
+    if "idx_document_map_unit_tokens_lookup" not in indexes:
+        op.create_index(
+            "idx_document_map_unit_tokens_lookup",
+            "document_map_unit_tokens",
+            ["channel", "token_hash", "map_unit_id"],
+        )
+    if "idx_document_map_unit_tokens_unit" not in indexes:
+        op.create_index(
+            "idx_document_map_unit_tokens_unit",
+            "document_map_unit_tokens",
+            ["map_unit_id", "channel"],
+        )
+
+
+def downgrade() -> None:
+    existing_tables = set(sa.inspect(op.get_bind()).get_table_names())
+    for table_name in (
+        "document_map_unit_tokens",
+        "document_map_units",
+        "document_map_unit_indexes",
+    ):
+        if table_name in existing_tables:
+            op.drop_table(table_name)

apps/api/scripts/backfill_map_unit_indexes.py

@@ -0,0 +1,97 @@
+"""Backfill persisted MAP-NAV lexical indexes for existing revisions.
+
+The migration creates empty derived tables intentionally. Run this command
+after deployment with ``--apply`` so each revision is rebuilt and committed
+independently; without ``--apply`` it is a read-only inventory.
+"""
+
+# ruff: noqa: E402
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+
+def _bootstrap_python_path() -> None:
+    api_root = Path(__file__).resolve().parents[1]
+    repo_root = api_root.parents[1]
+    shared_root = repo_root / "packages" / "shared-python"
+    for path in (api_root, shared_root):
+        value = os.fspath(path)
+        if value not in sys.path:
+            sys.path.insert(0, value)
+
+
+_bootstrap_python_path()
+
+from sqlalchemy import select
+
+from shared.core.database_sync import get_sync_session_factory
+from shared.models.database.document import Document
+from shared.services.retrieval.map_unit_index import replace_document_map_units
+from shared.services.retrieval.publication_models import DocumentPublicationScope
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Backfill MAP-NAV indexes for current document revisions."
+    )
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="Build and commit each current revision index.",
+    )
+    parser.add_argument("--document-id", default="", help="Limit the backfill to one document.")
+    return parser
+
+
+def _load_documents(document_id: str) -> list[Document]:
+    session_factory = get_sync_session_factory()
+    with session_factory() as db:
+        statement = select(Document).where(Document.current_job_result_id.is_not(None))
+        normalized_document_id = document_id.strip()
+        if normalized_document_id:
+            statement = statement.where(Document.document_id == normalized_document_id)
+        return list(db.scalars(statement).all())
+
+
+def backfill_map_unit_indexes(*, apply: bool, document_id: str = "") -> int:
+    documents = _load_documents(document_id)
+    if not apply:
+        for document in documents:
+            print(f"would backfill document={document.document_id} revision={document.current_job_result_id}")
+        return len(documents)
+
+    session_factory = get_sync_session_factory()
+    for document in documents:
+        job_result_id = document.current_job_result_id
+        if not job_result_id:
+            continue
+        scope = DocumentPublicationScope(
+            user_id=document.user_id,
+            namespace=document.namespace,
+            document_id=document.document_id,
+            job_result_id=job_result_id,
+            source_file_name=str(document.source_file_name or ""),
+        )
+        with session_factory() as db:
+            replace_document_map_units(db, scope=scope)
+            db.commit()
+        print(f"backfilled document={document.document_id} revision={job_result_id}")
+    return len(documents)
+
+
+def main() -> None:
+    arguments = _build_parser().parse_args()
+    count = backfill_map_unit_indexes(
+        apply=bool(arguments.apply), document_id=str(arguments.document_id)
+    )
+    action = "backfilled" if arguments.apply else "found"
+    print(f"{action} revisions={count}")
+
+
+if __name__ == "__main__":
+    main()