From b4a1351628fab6088ac2e4b8a70d73c88a8f129d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 16:10:59 +0800 Subject: [PATCH 1/2] perf(retrieval): make evidence packing linear --- .../test_retrieval_lazy_tree_contract.py | 72 +++++ .../services/retrieval/nav/nav_compose.py | 280 +++++++++++++++--- 2 files changed, 304 insertions(+), 48 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py index 30994dad..a920bb39 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py @@ -5,6 +5,9 @@ from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow from shared.services.retrieval.nav.nav_map_scores import _walk_tree +from shared.services.retrieval.nav._compat import Chunk +from shared.services.retrieval.nav.nav_compose import pack_nav_evidence +from shared.services.retrieval.nav.nav_types import NavConfig, NavState class _MetadataForbiddenProvider(KnowhereProvider): @@ -31,3 +34,72 @@ def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> assert children == {"root": ["child"], "child": []} assert leaves == {"child"} assert titles == {"root": "Root", "child": "Child"} + + +def test_evidence_pack_reads_titles_without_materializing_subtree_metadata() -> None: + class _CountingMetadataProvider(_MetadataForbiddenProvider): + metadata_calls = 0 + + def node_meta(self, section_id: str) -> NodeMeta: + self.metadata_calls += 1 + return super().node_meta(section_id) + + provider = _CountingMetadataProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), + ], + units=(), + ) + toolspace = ProviderToolSpace(provider) + state = NavState(doc_id="doc", query="child") + chunk = Chunk( + node_id="child", + doc_id="doc", + text="evidence", + line_ids=(1,), + section_id="child", + ) + + result = pack_nav_evidence( + [(chunk, 1.0)], + toolspace, + state, + NavConfig(), + budget_chars=100, + ) + + assert result.evidence_text == "[E1]\n[§ Child]\nevidence" + assert provider.metadata_calls == 0 + + +def test_evidence_pack_identifies_header_owners_from_parent_chain() -> None: + provider = KnowhereProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("parent", "root", "Root / Parent", "Parent", 1, "", 1), + SectionRow( + "child", "parent", "Root / Parent / Child", "Child", 2, "", 2 + ), + ], + units=(), + ) + toolspace = ProviderToolSpace(provider) + state = NavState(doc_id="doc", query="child") + chunks = [ + Chunk("parent", "doc", "parent evidence", (1,), "parent"), + Chunk("child", "doc", "child evidence", (2,), "child"), + ] + + result = pack_nav_evidence( + [(chunks[0], 1.0), (chunks[1], 0.9)], + toolspace, + state, + NavConfig(), + budget_chars=200, + ) + + assert result.kept_chunks == [chunks[1]] + assert result.evidence_text == "[E1]\n[§ Child]\nchild evidence" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_compose.py b/packages/shared-python/shared/services/retrieval/nav/nav_compose.py index 4c75b2f9..2cdbfe3d 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_compose.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_compose.py @@ -111,6 +111,18 @@ def _clip(text: str) -> str: return t[:max_chars].rstrip() return t + # ``path_titles`` is a lightweight title lookup. Prefer it over + # ``get_structure`` because the latter may calculate the complete subtree + # chunk count, which is unnecessary while rendering evidence headers. + path_titles = getattr(ts, "path_titles", None) + if callable(path_titles): + try: + path = str(path_titles(sid, doc_id) or "").strip() + except TypeError: + path = str(path_titles(sid) or "").strip() + if path: + return _clip(path.rsplit(" / ", 1)[-1]) + # Prefer structure title (Knowhere / ProviderToolSpace); never parse ids. try: st = ts.get_structure(sid) @@ -172,25 +184,43 @@ def _child_final_score( return float(own_unit) + float(w_conf) * conf -def _is_header_only_owner(owner: str, owners: set[str], ts: ToolSpace, doc_id: str) -> bool: - """True if owner is a structural ancestor of another collected owner.""" - if not owner or owner not in owners: - return False - owner_doc = _section_doc_id(ts, owner, doc_id) - for other in owners: - if other == owner: - continue - if _section_doc_id(ts, other, doc_id) != owner_doc: +def _parent_id_for_compose(ts: ToolSpace, section_id: str, doc_id: str) -> Optional[str]: + """Resolve a direct parent without materializing a section subtree.""" + parent_fn = getattr(ts, "parent_id", None) + if callable(parent_fn): + try: + parent = parent_fn(section_id) + except Exception: + parent = None + if parent: + return str(parent) + return direct_parent_id(ts, section_id, doc_id) + + +def _header_only_owners( + owners: set[str], ts: ToolSpace, doc_id: str +) -> set[str]: + """Find collected owners that are ancestors of another collected owner. + + Walk each owner's parent chain once instead of comparing every owner with + every other owner. This keeps compose grouping proportional to the number + of owners and hierarchy depth, even when a broad collect returns thousands + of chunks. + """ + headers: set[str] = set() + for owner in owners: + if not owner: continue - cur = other + owner_doc = _section_doc_id(ts, owner, doc_id) + current = owner for _ in range(64): - p = direct_parent_id(ts, cur, owner_doc) - if p is None: + parent = _parent_id_for_compose(ts, current, owner_doc) + if parent is None: break - if p == owner: - return True - cur = p - return False + if parent in owners and _section_doc_id(ts, parent, doc_id) == owner_doc: + headers.add(parent) + current = parent + return headers @dataclass @@ -259,9 +289,7 @@ def _build_groups( items.append((chunk, owner, score)) owners = {owner for _c, owner, _s in items if owner} - header_owners = { - o for o in owners if _is_header_only_owner(o, owners, ts, state.doc_id) - } + header_owners = _header_only_owners(owners, ts, state.doc_id) groups: Dict[Optional[str], _ParentGroup] = {} for chunk, owner, score in items: @@ -407,27 +435,183 @@ def _render_kept( ) +def _count_power10_boundaries(lower: int, upper: int) -> int: + """Count decimal digit-boundary positions in ``(lower, upper]``.""" + if upper <= lower: + return 0 + boundary = 10 + count = 0 + while boundary <= upper: + if boundary > lower: + count += 1 + boundary *= 10 + return count + + +class _SelectionLength: + """Maintain exact rendered selection length under one-child mutations.""" + + def __init__(self, groups: Sequence[_ParentGroup], kept_ids: Set[str]) -> None: + self._groups = groups + self.kept_ids = kept_ids + self._group_by_node_id: Dict[str, int] = {} + self._plain_lengths: Dict[str, int] = {} + self._indented_lengths: Dict[str, int] = {} + self._counts: List[int] = [] + self._plain_totals: List[int] = [] + self._indented_totals: List[int] = [] + self._nonempty_count = 0 + self.total = 0 + + for group_index, group in enumerate(groups): + plain_total = 0 + indented_total = 0 + count = 0 + for child in group.children: + node_id = str(child.chunk.node_id or "") + if not node_id: + continue + body = _chunk_body(child.chunk) + plain_length = len(body) + indented_length = len( + "\n".join( + (" " + line if line.strip() else line) + for line in body.splitlines() + ) + ) + self._group_by_node_id[node_id] = group_index + self._plain_lengths[node_id] = plain_length + self._indented_lengths[node_id] = indented_length + if node_id in kept_ids: + count += 1 + plain_total += plain_length + indented_total += indented_length + self._counts.append(count) + self._plain_totals.append(plain_total) + self._indented_totals.append(indented_total) + + self._recalculate_total() + + @staticmethod + def _block_length( + group: _ParentGroup, + *, + count: int, + plain_total: int, + indented_total: int, + evidence_index: int, + ) -> int: + if count <= 0: + return 0 + prefix = len(f"[E{evidence_index}]") + if group.parent_title: + prefix += 1 + len(f"[§ {group.parent_title}]") + body_length = indented_total if count >= 2 else plain_total + return prefix + count + body_length + + def _recalculate_total(self) -> None: + total = 0 + evidence_index = 0 + for index, group in enumerate(self._groups): + count = self._counts[index] + if count <= 0: + continue + evidence_index += 1 + total += self._block_length( + group, + count=count, + plain_total=self._plain_totals[index], + indented_total=self._indented_totals[index], + evidence_index=evidence_index, + ) + self._nonempty_count = evidence_index + self.total = total + max(0, evidence_index - 1) * 2 + + def _evidence_index(self, group_index: int) -> int: + return sum(1 for count in self._counts[:group_index] if count > 0) + 1 + + def remove(self, node_id: str) -> None: + group_index = self._group_by_node_id.get(node_id) + if group_index is None or node_id not in self.kept_ids: + return + old_count = self._counts[group_index] + old_index = self._evidence_index(group_index) + old_block = self._block_length( + self._groups[group_index], + count=old_count, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=old_index, + ) + self.kept_ids.remove(node_id) + self._counts[group_index] -= 1 + self._plain_totals[group_index] -= self._plain_lengths[node_id] + self._indented_totals[group_index] -= self._indented_lengths[node_id] + if old_count > 1: + new_block = self._block_length( + self._groups[group_index], + count=old_count - 1, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=old_index, + ) + self.total += new_block - old_block + return + + self.total -= old_block + if self._nonempty_count > 1: + self.total -= 2 + self.total -= _count_power10_boundaries(old_index, self._nonempty_count) + self._nonempty_count -= 1 + + def add(self, node_id: str) -> bool: + group_index = self._group_by_node_id.get(node_id) + if group_index is None or node_id in self.kept_ids: + return False + old_count = self._counts[group_index] + if old_count > 0: + evidence_index = self._evidence_index(group_index) + old_block = self._block_length( + self._groups[group_index], + count=old_count, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=evidence_index, + ) + else: + evidence_index = self._evidence_index(group_index) + old_block = 0 + + self.kept_ids.add(node_id) + self._counts[group_index] += 1 + self._plain_totals[group_index] += self._plain_lengths[node_id] + self._indented_totals[group_index] += self._indented_lengths[node_id] + new_block = self._block_length( + self._groups[group_index], + count=self._counts[group_index], + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=evidence_index, + ) + if old_count > 0: + self.total += new_block - old_block + else: + old_nonempty_count = self._nonempty_count + self.total += new_block + if old_nonempty_count > 0: + self.total += 2 + self.total += _count_power10_boundaries( + evidence_index, old_nonempty_count + 1 + ) + self._nonempty_count += 1 + return True + + def _selection_chars( groups: Sequence[_ParentGroup], kept_ids: Set[str] ) -> int: - """Exact rendered length of the current full-text selection (no truncation).""" - parts: List[str] = [] - sep = "\n\n" - for g in groups: - entries = [ - c - for c in sorted(g.children, key=lambda x: x.line_key) - if c.chunk.node_id in kept_ids - ] - if not entries: - continue - block = _render_group( - g, entries, evidence_index=len(parts) + 1, indent=len(entries) >= 2 - ) - parts.append(block) - if not parts: - return 0 - return len(sep.join(parts)) + """Exact rendered length of the current full-text selection.""" + return _SelectionLength(groups, kept_ids).total def _refill_to_budget( @@ -460,18 +644,16 @@ def _refill_to_budget( c.line_key, ) ) - used = _selection_chars(groups, kept_ids) + selection = _SelectionLength(groups, kept_ids) for child in dropped: - if used >= budget_chars: + if selection.total >= budget_chars: break - if len(child.chunk.text or "") > budget_chars - used: + if len(child.chunk.text or "") > budget_chars - selection.total: continue - kept_ids.add(child.chunk.node_id) - size = _selection_chars(groups, kept_ids) - if size > budget_chars: - kept_ids.discard(child.chunk.node_id) + if not selection.add(child.chunk.node_id): continue - used = size + if selection.total > budget_chars: + selection.remove(child.chunk.node_id) def _pack_trim( @@ -497,8 +679,10 @@ def _pack_trim( if not kept_ids: return ComposeFillResult([], "", 0, 0, False, scored_flat, dropped_any=False) + selection = _SelectionLength(groups, kept_ids) + def fits() -> bool: - return _selection_chars(groups, kept_ids) <= budget_chars + return selection.total <= budget_chars if fits(): return _render_kept( @@ -516,7 +700,7 @@ def fits() -> bool: for child in candidates: if fits(): break - kept_ids.discard(child.chunk.node_id) + selection.remove(child.chunk.node_id) if fits(): break @@ -532,7 +716,7 @@ def fits() -> bool: break if len(kept_ids) <= 1: break - kept_ids.discard(child.chunk.node_id) + selection.remove(child.chunk.node_id) if fits() or len(kept_ids) <= 1: break From e44abb3ed0e334f7093daca9db5d813c05d74569 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 1 Sep 2026 17:18:39 +0800 Subject: [PATCH 2/2] perf(retrieval): drive map frequency lookup by token hash --- ...est_retrieval_classic_map_unit_contract.py | 7 +- .../test_retrieval_map_unit_index_contract.py | 72 +++++++++++++++++++ .../services/retrieval/nav/nav_knowhere.py | 22 ++++-- .../retrieval/search/map_unit_discovery.py | 42 ++++++----- 4 files changed, 116 insertions(+), 27 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index 796d2ba8..63941a1f 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -106,8 +106,9 @@ def capture_frequency_query( _executemany: bool, ) -> None: if ( - "FROM document_map_unit_tokens AS tokens" in statement - and "tokens.frequency" in statement + "FROM document_map_unit_tokens" in statement + and "frequency" in statement + and "token_hash" in statement ): statements.append(statement) @@ -168,6 +169,8 @@ def capture_frequency_query( assert statements assert "token_hash = ANY" in statements[-1] assert "token = ANY" not in statements[-1] + assert "matching_tokens AS MATERIALIZED" in statements[-1] + assert "FROM matching_tokens" in statements[-1] async def test_classic_route_image_filter_scores_only_units_with_images( diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index e95efd23..c44f04fa 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -15,6 +15,7 @@ ) from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace from shared.services.retrieval.nav._compat import Chunk, EpisodeResult +from shared.services.retrieval.nav import nav_knowhere from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, @@ -41,6 +42,77 @@ _USER_ID = "local-dev-user" +def test_read_only_score_loader_drives_frequency_lookup_from_token_hash( + monkeypatch, +) -> None: + document_id = "doc-frequency" + job_result_id = "revision-frequency" + executions: list[tuple[str, object]] = [] + + class FakeCursor: + def __init__(self) -> None: + self.rows: list[tuple[object, ...]] = [] + + def execute(self, statement: str, parameters: object = None) -> None: + executions.append((statement, parameters)) + if "document_map_unit_indexes" in statement: + self.rows = [(document_id, job_result_id, 1, 1, 0.0, 0.0)] + elif "FROM document_map_units AS units" in statement: + self.rows = [("unit-frequency", document_id, "chunk-frequency", "section-frequency", 1, 1)] + elif "matching_tokens AS MATERIALIZED" in statement: + self.rows = [("unit-frequency", "path", "retrieval", 1)] + else: + self.rows = [] + + def fetchall(self) -> list[tuple[object, ...]]: + return self.rows + + def close(self) -> None: + return None + + class FakeConnection: + def __init__(self) -> None: + self.cursor_instance = FakeCursor() + + def set_session(self, *, readonly: bool, autocommit: bool) -> None: + assert readonly is True + assert autocommit is True + + def cursor(self) -> FakeCursor: + return self.cursor_instance + + def close(self) -> None: + return None + + connection = FakeConnection() + monkeypatch.setattr(nav_knowhere, "_connect", lambda _dsn: connection) + store = nav_knowhere.ReadOnlyChunkStore( + dsn="postgresql://test", + revisions={document_id: job_result_id}, + ) + + corpus = store.load_persisted_score_corpus( + [document_id], + {document_id: ["section-frequency"]}, + ["retrieval"], + ) + + assert corpus is not None + frequency_executions = [ + (statement, parameters) + for statement, parameters in executions + if "matching_tokens AS MATERIALIZED" in statement + ] + assert len(frequency_executions) == 1 + statement, parameters = frequency_executions[0] + assert "FROM matching_tokens" in statement + assert "token_hash = ANY" in statement + assert isinstance(parameters, list) + assert parameters[0] == [ + "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7" + ] + + class _IncompleteIndexStore: """Minimal lazy store whose missing index returns no persisted scores.""" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 38cef224..68d76d3e 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -24,6 +24,7 @@ import logging import time from dataclasses import dataclass, field +from hashlib import sha256 from typing import ( Any, Callable, @@ -323,6 +324,9 @@ def load_persisted_score_corpus( for token in tokenize_query_for_ranker(query) ) ) + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] cur = self._connection().cursor() try: revision_key = tuple(revisions) @@ -432,14 +436,20 @@ def load_persisted_score_corpus( if unit_rows and query_tokens: stage_started = time.perf_counter() cur.execute( - "SELECT units.id, tokens.channel, tokens.token, tokens.frequency " - "FROM document_map_unit_tokens AS tokens " - "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " + "WITH matching_tokens AS MATERIALIZED (" + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE token_hash = ANY(%s) AND channel = ANY(%s)" + ") " + "SELECT matching_tokens.map_unit_id, matching_tokens.channel, " + "matching_tokens.token, matching_tokens.frequency " + "FROM matching_tokens " + "JOIN document_map_units AS units " + "ON units.id = matching_tokens.map_unit_id " f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id " - "WHERE tokens.token = ANY(%s) AND tokens.channel = ANY(%s)", - [*revision_params, list(query_tokens), list(_MAP_SCORE_CHANNELS)], + "AND units.job_result_id = revisions.job_result_id", + [list(query_token_hashes), list(_MAP_SCORE_CHANNELS), *revision_params], ) allowed_map_unit_ids = {str(row["map_unit_id"]) for row in unit_rows} for map_unit_id, channel, token, frequency in cur.fetchall(): diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index dc68d896..95436458 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -232,26 +232,30 @@ async def map_unit_discovery( if not unit_rows: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) - frequency_result = await db.execute( - text( - ( - cte - if signal_paths - else _SCOPED_UNIT_IDS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, - exclude_clause=exclude_clause, - type_clause=type_clause, - ) - ) - + """ - SELECT tokens.map_unit_id, tokens.channel, tokens.token, tokens.frequency - FROM document_map_unit_tokens AS tokens - JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id - WHERE tokens.channel = ANY(:channels) - AND tokens.token_hash = ANY(:token_hashes) + frequency_scope_cte = cte if signal_paths else _SCOPED_UNIT_IDS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + ) + frequency_query = text( + "WITH matching_tokens AS MATERIALIZED (" + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE channel = ANY(:channels) " + "AND token_hash = ANY(:token_hashes)" + "), " + + frequency_scope_cte.lstrip().removeprefix("WITH ") + + """ + SELECT matching_tokens.map_unit_id, matching_tokens.channel, + matching_tokens.token, matching_tokens.frequency + FROM matching_tokens + JOIN scoped_units + ON scoped_units.map_unit_id = matching_tokens.map_unit_id """ - ), + ) + frequency_result = await db.execute( + frequency_query, { **params, "channels": list(_MAP_SCORE_CHANNELS),