From 8728357f5d15e70c7e7a9b81c71690b20a3d1944 Mon Sep 17 00:00:00 2001 From: Dane Urban Date: Sat, 17 Oct 2026 23:24:06 -0700 Subject: [PATCH] fix(indexing): stop the re-chunk when the live index is swapped The runner binds its document index once at startup. An INSTANT swap landing mid-run sent writes to the index being retired while lowering the shared chunk_count, so the new live index kept the oversized chunks and no later scan could select the document again. Re-check the live index name per page and stop, keeping the cursor so the next tick resumes once the tick's own guard allows it. One batch remains exposed, since the check cannot reach inside the pipeline write. Also trim comments that restated the code beside them. Reported by greptile on #16153. Co-Authored-By: Claude Opus 5 (1M context) --- .../celery/tasks/tabular_rechunk/tasks.py | 18 +++-- .../indexing/run_tabular_rechunk.py | 17 +++-- backend/onyx/db/document.py | 10 ++- .../indexing/test_run_tabular_rechunk.py | 69 ++++++++++++++++++- 4 files changed, 91 insertions(+), 23 deletions(-) diff --git a/backend/onyx/background/celery/tasks/tabular_rechunk/tasks.py b/backend/onyx/background/celery/tasks/tabular_rechunk/tasks.py index 4a3e9ed7cb..0c55db7604 100644 --- a/backend/onyx/background/celery/tasks/tabular_rechunk/tasks.py +++ b/backend/onyx/background/celery/tasks/tabular_rechunk/tasks.py @@ -35,8 +35,8 @@ from onyx.db.search_settings import get_active_search_settings from onyx.key_value_store.factory import get_kv_store from onyx.key_value_store.interface import KvKeyNotFoundError -# Outlives one run so a crashed worker's claim expires; extended after each -# batch, since one batch can outlast the budget under provider throttling. +# Outlives one run so a crashed worker's claim expires. Extended per batch, +# since one batch can outlast the budget under provider throttling. _LOCK_TIMEOUT_S = TABULAR_RECHUNK_TIME_BUDGET_S + 5 * 60 # Bounded so a permanently broken document cannot rescan for ever. _MAX_PASSES = 3 @@ -78,13 +78,12 @@ def _blocked_reason(db_session: Session) -> str | None: search_settings = get_active_search_settings(db_session) if search_settings.secondary is not None: - # The pipeline writes to the secondary whenever a FUTURE exists, - # leaving the live index on the old chunks. + # The pipeline writes to the secondary whenever a FUTURE exists. return "a reindex is in flight" - # Only while an INSTANT swap is still backfilling: the port copies chunks - # boundary for boundary over finished work. Without the source-id check - # this blocks for ever, because "pending work" means a cc_pair has no - # settled port attempt - the normal state. + # Only while an INSTANT swap is backfilling, since the port copies chunks + # boundary for boundary. The source-id check is load-bearing: "pending + # work" alone means a cc_pair has no settled port attempt - the normal + # state - so without it this blocks for ever. if search_settings.primary.port_backfill_source_id is not None and ( port_backfill_has_pending_work(db_session, search_settings.primary.id) ): @@ -153,8 +152,7 @@ def run_tabular_rechunk_tick(tenant_id: str) -> int | None: on_batch=extend_lock, ) - # Re-read: the run took minutes, so write onto current state rather - # than a stale copy. + # The run took minutes; write onto current state, not a stale copy. progress = load_progress() progress.cursor = result.last_document_id progress.documents_indexed += result.documents_indexed diff --git a/backend/onyx/background/indexing/run_tabular_rechunk.py b/backend/onyx/background/indexing/run_tabular_rechunk.py index dbbf72126b..a13d3d9411 100644 --- a/backend/onyx/background/indexing/run_tabular_rechunk.py +++ b/backend/onyx/background/indexing/run_tabular_rechunk.py @@ -65,7 +65,7 @@ logger = setup_logger() TABULAR_RECHUNK_TIME_BUDGET_S = 4 * 60 # Per-workbook cap plus the large chunks multipass derives from it. Counting -# those even when multipass is off only loosens the bound. +# those when multipass is off only loosens the bound. TABULAR_CHUNK_CEILING = TABULAR_MAX_CHUNKS_PER_WORKBOOK + ( TABULAR_MAX_CHUNKS_PER_WORKBOOK // LARGE_CHUNK_RATIO ) @@ -121,7 +121,7 @@ def _rebuild_document( id=db_doc.id, source=DocumentSource(head.source_type), semantic_identifier=db_doc.semantic_id, - # None would fall back to semantic_identifier and add a title vector + # None would fall back to semantic_identifier, adding a title vector # the document never had. title=head.title if head.title is not None else "", metadata=( @@ -196,8 +196,7 @@ def _rechunk_batch( # never moves, because content_hash() ignores TabularSections. ignore_time_skip=True, force_update=True, - # Only the chunks and chunk_count change here; nothing re-fetched - # the document, so its row has nothing newer to receive. + # Nothing re-fetched the document, so its row has nothing newer. skip_document_row_writes=True, tenant_id=tenant_id, document_batch=documents, @@ -250,9 +249,16 @@ def run_tabular_rechunk( search_settings=search_settings, callback=None ) document_index = build_opensearch_document_index(search_settings) + index_name = search_settings.index_name # read while attached while time.monotonic() - start < time_budget_s: with get_session_with_current_tenant() as db_session: + # An INSTANT swap retires the index this run holds: writing on + # would lower chunk_count while the new live index keeps the old + # chunks, stranding it. The cursor is kept for the next tick. + if get_current_search_settings(db_session).index_name != index_name: + logger.info("tabular re-chunk: the live index changed, stopping") + break targets = get_oversized_tabular_documents( db_session, chunk_count_above=TABULAR_CHUNK_CEILING, @@ -262,8 +268,7 @@ def run_tabular_rechunk( # A short page means no further matches before the end of the table. finished = len(targets) < INDEX_BATCH_SIZE - # The pipeline adapter is per connector + credential, so a batch cannot - # mix cc-pairs. + # The pipeline adapter is per connector + credential. grouped: dict[tuple[int, int], list[TabularRechunkTarget]] = defaultdict(list) for target in targets: grouped[(target.connector_id, target.credential_id)].append(target) diff --git a/backend/onyx/db/document.py b/backend/onyx/db/document.py index 24ce24db59..c45af138ad 100644 --- a/backend/onyx/db/document.py +++ b/backend/onyx/db/document.py @@ -338,10 +338,9 @@ def get_oversized_tabular_documents( Left out: documents whose raw file was never stored, because only a crawl can rebuild their sections, and documents whose cc-pairs are all DELETING. """ - # One row per document, the lowest cc-pair winning, as any of them would - # write the same chunks. A LATERAL rather than DISTINCT ON, which would - # order by (id, connector_id, credential_id) over the join and so could - # not use the ordered scan the cursor depends on. + # One row per document, lowest cc-pair winning, as any of them would write + # the same chunks. A LATERAL, not DISTINCT ON: that would order the join by + # (id, connector_id, credential_id) and lose the ordered scan. cc_pair = ( select( DocumentByConnectorCredentialPair.connector_id, @@ -383,8 +382,7 @@ def get_oversized_tabular_documents( DbDocument.chunk_count > chunk_count_above, DbDocument.file_id.isnot(None), or_( - # A Drive-exported spreadsheet keeps its original name, which - # often has no extension, so the stored file's type decides. + # A Drive sheet's name often has no extension; its type does. FileRecord.file_type.in_(sorted(OnyxMimeTypes.TABULAR_MIME_TYPES)), *[ func.lower(DbDocument.semantic_id).like(f"%{extension}") diff --git a/backend/tests/unit/onyx/background/indexing/test_run_tabular_rechunk.py b/backend/tests/unit/onyx/background/indexing/test_run_tabular_rechunk.py index ec003130ba..a4fa25430b 100644 --- a/backend/tests/unit/onyx/background/indexing/test_run_tabular_rechunk.py +++ b/backend/tests/unit/onyx/background/indexing/test_run_tabular_rechunk.py @@ -7,7 +7,9 @@ survive the round trip. """ from io import BytesIO -from typing import IO +from itertools import chain, repeat +from types import SimpleNamespace +from typing import IO, Any from unittest.mock import MagicMock, patch import pytest @@ -15,6 +17,7 @@ import pytest from onyx.background.indexing.run_tabular_rechunk import ( _rebuild_document, _tabular_file_name, + run_tabular_rechunk, ) from onyx.configs.constants import DocumentSource from onyx.connectors.models import ( @@ -96,3 +99,67 @@ def test_tabular_file_name_supplies_a_missing_extension() -> None: _tabular_file_name("notes", "application/pdf") # And the rebuild routes an extensionless name through it. assert _rebuild(semantic_id="Q3 Budget").semantic_identifier == "Q3 Budget" + + +class _FakeTarget: + def __init__(self, index: int) -> None: + self.document = SimpleNamespace(id=f"doc-{index:02d}") + self.connector_id = 1 + self.credential_id = 1 + self.file_type = "text/csv" + + +def _run_against(index_names: list[str], pages_before_the_end: int) -> Any: + """Runs the scan with `get_current_search_settings` answering in turn.""" + names = chain(index_names, repeat(index_names[-1])) + pages = 0 + + def settings(_db_session: object) -> SimpleNamespace: + return SimpleNamespace(index_name=next(names), id=1) + + def targets(*args: Any, **kwargs: Any) -> list[_FakeTarget]: # noqa: ARG001 + nonlocal pages + pages += 1 + # A short page ends the scan, as a real one does at the table's end. + return [ + _FakeTarget(i) for i in range(16 if pages < pages_before_the_end else 3) + ] + + def batch( + targets: list[_FakeTarget], + **kwargs: Any, # noqa: ARG001 + ) -> tuple[int, list[str]]: + return len(targets), [] + + module = "onyx.background.indexing.run_tabular_rechunk" + with ( + patch(f"{module}.get_current_search_settings", settings), + patch(f"{module}.get_oversized_tabular_documents", targets), + patch(f"{module}._rechunk_batch", batch), + patch(f"{module}.get_session_with_current_tenant"), + patch(f"{module}.DefaultIndexingEmbedder"), + patch(f"{module}.build_opensearch_document_index"), + ): + result = run_tabular_rechunk( + tenant_id="t", after_doc_id=None, on_batch=lambda: None + ) + return result, pages + + +def test_the_scan_stops_when_the_live_index_is_swapped() -> None: + """An INSTANT swap retires the index this run holds. Writing on would + replace chunks in the index being dropped while lowering the shared + chunk_count, stranding the new live index oversized and unselectable. + """ + result, pages = _run_against(["idx-a", "idx-a", "idx-a", "idx-b"], 99) + + assert pages == 2 + # Not "complete": the pass has to be resumable, not marked done. + assert not result.finished + assert result.last_document_id == "doc-15" + + +def test_an_undisturbed_scan_still_runs_to_the_end() -> None: + result, pages = _run_against(["idx-a"], pages_before_the_end=3) + assert pages == 3 + assert result.finished