Skip to content

fix: remove retrieval reads of dropped namespace statistics #368

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
suguanYang merged 1 commit into from
Sep 1, 2026
Merged

fix: remove retrieval reads of dropped namespace statistics #368

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 9 additions & 46 deletions packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,6 @@ def load_persisted_score_corpus(
comes from ``document_map_unit_indexes`` (written at index time).
"""
from shared.services.retrieval.nav.persisted_score_load import (
average_idf_from_namespace_stats,
build_channel_bm25_stats,
combine_average_idf,
)
Expand Down Expand Up @@ -360,51 +359,15 @@ def load_persisted_score_corpus(
revision_key
]
else:
total_unit_count = sum(int(row[3] or 0) for row in index_rows)
try:
cur.execute(
"SELECT tokens.channel, tokens.token, "
"COUNT(DISTINCT tokens.map_unit_id) "
"FROM document_map_unit_tokens AS tokens "
"JOIN document_map_units AS units ON units.id = 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.channel = ANY(%s) "
"GROUP BY tokens.channel, tokens.token",
[*revision_params, ["path", "content"]],
)
namespace_token_dfs: dict[str, list[int]] = {
"path": [],
"content": [],
}
for channel, _token, document_frequency in cur.fetchall():
if str(channel) in namespace_token_dfs:
namespace_token_dfs[str(channel)].append(
int(document_frequency)
)
average_idf_path = average_idf_from_namespace_stats(
unit_count=total_unit_count,
token_document_frequencies=namespace_token_dfs["path"],
)
average_idf_content = average_idf_from_namespace_stats(
unit_count=total_unit_count,
token_document_frequencies=namespace_token_dfs["content"],
)
except Exception as exc:
_logger.warning(
"exact namespace IDF load failed; using revision averages: %s",
exc,
)
average_idf_path = combine_average_idf(
[(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows]
)
average_idf_content = combine_average_idf(
[
(float(row[5] or 0.0), int(row[3] or 0))
for row in index_rows
]
)
average_idf_path = combine_average_idf(
[(float(row[4] or 0.0), int(row[3] or 0)) for row in index_rows]
)
average_idf_content = combine_average_idf(
[
(float(row[5] or 0.0), int(row[3] or 0))
for row in index_rows
]
)
self._score_average_idf_cache[revision_key] = (
average_idf_path,
average_idf_content,
Expand Down
23 changes: 0 additions & 23 deletions packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,29 +38,6 @@ def combine_average_idf(parts: Sequence[tuple[float, int]]) -> float:
)


def average_idf_from_namespace_stats(
*,
unit_count: int,
token_document_frequencies: Sequence[int],
) -> float:
"""Compute the exact namespace-level average IDF used by rank_bm25.

Namespace token statistics already contain one document frequency per
token. Computing the mean from those rows avoids the incorrect
per-revision-average approximation when a namespace contains revisions
with different token distributions.
"""
if unit_count <= 0:
return 0.0
idfs = [
math.log(unit_count - int(frequency) + 0.5)
- math.log(int(frequency) + 0.5)
for frequency in token_document_frequencies
if 0 < int(frequency) <= unit_count
]
return sum(idfs) / len(idfs) if idfs else 0.0


def build_channel_bm25_stats(
*,
unit_rows: Sequence[Mapping[str, Any]],
Expand Down
110 changes: 12 additions & 98 deletions packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,9 @@
tokenize_query_for_ranker,
)
from shared.services.retrieval.nav.persisted_score_load import (
average_idf_from_namespace_stats,
build_channel_bm25_stats,
combine_average_idf,
)
from shared.services.retrieval.serving_manifest import decode_serving_manifest
from shared.services.retrieval.cache_service import record_retrieval_index_readiness
from shared.services.retrieval.search.scoring import normalize_row_scores
from shared.services.retrieval.search.section_filters import is_excluded_section
Expand Down Expand Up @@ -103,80 +101,6 @@ class DiscoveryResult:
error: str | None = None


async def _load_exact_namespace_average_idf(
db: AsyncSession,
*,
user_id: str,
namespace: str,
) -> tuple[float, float] | None:
"""Load exact namespace IDF floors from the published aggregate tables."""
generation_row = (
await db.execute(
text(
"SELECT generation FROM retrieval_namespace_generations "
"WHERE user_id = :user_id AND namespace = :namespace"
),
{"user_id": user_id, "namespace": namespace},
)
).first()
if generation_row is None:
return None
generation = int(generation_row[0])
stat_row = (
await db.execute(
text(
"SELECT payload_zlib, checksum, format_version "
"FROM retrieval_namespace_stats "
"WHERE user_id = :user_id AND namespace = :namespace "
"AND generation = :generation"
),
{
"user_id": user_id,
"namespace": namespace,
"generation": generation,
},
)
).first()
if stat_row is None:
return None
payload = decode_serving_manifest(
stat_row[0], checksum=str(stat_row[1]), format_version=int(stat_row[2])
)
unit_count = int(payload.get("unit_count") or 0)
if unit_count <= 0:
return None
token_rows = (
await db.execute(
text(
"SELECT channel, document_frequency "
"FROM retrieval_namespace_token_stats "
"WHERE user_id = :user_id AND namespace = :namespace "
"AND generation = :generation AND channel = ANY(:channels)"
),
{
"user_id": user_id,
"namespace": namespace,
"generation": generation,
"channels": ["path", "content"],
},
)
).all()
frequencies: dict[str, list[int]] = {"path": [], "content": []}
for channel, frequency in token_rows:
if str(channel) in frequencies:
frequencies[str(channel)].append(int(frequency))
return (
average_idf_from_namespace_stats(
unit_count=unit_count,
token_document_frequencies=frequencies["path"],
),
average_idf_from_namespace_stats(
unit_count=unit_count,
token_document_frequencies=frequencies["content"],
),
)


def _build_revision_scope(
revision_pins: Mapping[str, str] | None,
) -> tuple[str, str, dict[str, Any]]:
Expand Down Expand Up @@ -421,28 +345,18 @@ async def map_unit_discovery(
)
except Exception as exc:
logger.warning("retrieval index readiness publish failed: %s", exc)
exact_namespace_idf = None
if revision_pins is None:
exact_namespace_idf = await _load_exact_namespace_average_idf(
db,
user_id=user_id,
namespace=namespace,
)
if exact_namespace_idf is not None:
average_idf_path, average_idf_content = exact_namespace_idf
else:
average_idf_path = combine_average_idf(
[
(path_idf, unit_count)
for path_idf, _content_idf, unit_count in index_parts
]
)
average_idf_content = combine_average_idf(
[
(content_idf, unit_count)
for _path_idf, content_idf, unit_count in index_parts
]
)
average_idf_path = combine_average_idf(
[
(path_idf, unit_count)
for path_idf, _content_idf, unit_count in index_parts
]
)
average_idf_content = combine_average_idf(
[
(content_idf, unit_count)
for _path_idf, content_idf, unit_count in index_parts
]
)

path_stats = build_channel_bm25_stats(
unit_rows=unit_rows,
Expand Down
Loading