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) <noreply@anthropic.com>
This commit is contained in:
Dane UrbanandClaude Opus 5 committed 2026-10-17 23:24:30 -07:00
1 parent b328a04cc6
commit 8728357f5d
4 files changed
+91 -23

No files matched your search

@@ -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
@@ -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)
+4 -6
View File
@@ -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}")
@@ -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