diff --git a/api/db/services/doc_metadata_service.py b/api/db/services/doc_metadata_service.py index ab2bcafbfb..9e75b6d43a 100644 --- a/api/db/services/doc_metadata_service.py +++ b/api/db/services/doc_metadata_service.py @@ -32,6 +32,9 @@ from common.doc_store.doc_store_base import OrderByExpr from common.metadata_utils import dedupe_list +METADATA_ID_BATCH_SIZE = 10000 + + def _es_response_total(response: Any) -> int | None: """Extract the exact total hit count from an ES search response. @@ -199,6 +202,18 @@ class DocMetadataService: Returns: Search results from ES/Infinity, or empty list if index doesn't exist """ + if condition is None: + condition = {"kb_id": kb_id} + + doc_ids = condition.get("id") + if isinstance(doc_ids, list) and len(doc_ids) > METADATA_ID_BATCH_SIZE: + all_results = [] + for offset in range(0, len(doc_ids), METADATA_ID_BATCH_SIZE): + batch_condition = dict(condition) + batch_condition["id"] = doc_ids[offset : offset + METADATA_ID_BATCH_SIZE] + all_results.extend(cls._search_metadata(kb_id, condition=batch_condition)) + return all_results + kb = Knowledgebase.get_by_id(kb_id) if not kb: return [] @@ -215,9 +230,6 @@ class DocMetadataService: return [] logging.debug(f"Successfully created metadata index {index_name}") - if condition is None: - condition = {"kb_id": kb_id} - # Add sort by id for ES to enable search_after on large data order_by = OrderByExpr() if not settings.DOC_ENGINE_INFINITY: diff --git a/test/unit_test/api/db/services/test_doc_metadata_service_pagination.py b/test/unit_test/api/db/services/test_doc_metadata_service_pagination.py index 014a2d03e5..2b9dc4bff2 100644 --- a/test/unit_test/api/db/services/test_doc_metadata_service_pagination.py +++ b/test/unit_test/api/db/services/test_doc_metadata_service_pagination.py @@ -1,93 +1,119 @@ -# -# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Regression test for #16524: a manual metadata filter over a knowledge base -with more documents than the ES push-down cap (``filter_doc_ids_by_meta_pushdown``'s -default ``limit=10000``) must still see every document once the request falls -back to the in-memory path, not just the first page. - -Exercises ``DocMetadataService.get_flatted_meta_by_kbs`` end-to-end against a -fake, paginated ``docStoreConn`` standing in for Elasticsearch, then feeds the -result into ``meta_filter`` with the same ``not in`` condition from the -original report. -""" - -from types import SimpleNamespace - -import pytest - -from common import settings -from common.metadata_utils import meta_filter -from api.db.services.doc_metadata_service import DocMetadataService -from api.db.db_models import DB - -pytestmark = pytest.mark.p2 - -TOTAL_DOCS = 12000 -CANON_ZERO_COUNT = 30 # a small minority tagged "0"; the rest are "1" - - -class _FakeDocStoreConn: - """Stands in for the ES connection's paginated ``search``. - - Mirrors the shape ``DocMetadataService._iter_search_results`` expects - (``{"hits": {"hits": [{"_id": ..., "_source": {...}}]}}``) and actually - honors ``offset``/``limit`` so a caller that stops paginating too early - provably sees a truncated result, the way the reported bug did. - """ - - def __init__(self, total: int, canon_zero_count: int): - self._docs = [] - for i in range(total): - canon = "0" if i < canon_zero_count else "1" - self._docs.append({"_id": f"doc-{i}", "_source": {"meta_fields": {"canon": canon}}}) - - def index_exist(self, index_name, kb_id): - return True - - def search(self, select_fields, highlight_fields, condition, match_expressions, order_by, offset, limit, index_names, knowledgebase_ids, agg_fields=None, rank_feature=None): - page = self._docs[offset : offset + limit] - return {"hits": {"hits": page}} - - -def test_get_flatted_meta_by_kbs_returns_every_document_beyond_pushdown_cap(monkeypatch): - monkeypatch.setattr(DB, "connect", lambda *args, **kwargs: None) - monkeypatch.setattr(DB, "close", lambda *args, **kwargs: None) - monkeypatch.setattr(settings, "docStoreConn", _FakeDocStoreConn(TOTAL_DOCS, CANON_ZERO_COUNT)) - monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False) - fake_kb = SimpleNamespace(tenant_id="tenant-1") - monkeypatch.setattr("api.db.services.doc_metadata_service.Knowledgebase.get_by_id", lambda kb_id: fake_kb) - - metas = DocMetadataService.get_flatted_meta_by_kbs(["kb-1"]) - - assert len(metas["canon"]["1"]) == TOTAL_DOCS - CANON_ZERO_COUNT - assert len(metas["canon"]["0"]) == CANON_ZERO_COUNT - - -def test_manual_not_in_filter_matches_every_document_beyond_pushdown_cap(monkeypatch): - # Same scenario as the #16524 report: a "canon Not in ['0']" manual filter - # over a KB whose match set (TOTAL_DOCS - CANON_ZERO_COUNT) exceeds the - # push-down cap, so this exercises the in-memory fallback exclusively. - monkeypatch.setattr(DB, "connect", lambda *args, **kwargs: None) - monkeypatch.setattr(DB, "close", lambda *args, **kwargs: None) - monkeypatch.setattr(settings, "docStoreConn", _FakeDocStoreConn(TOTAL_DOCS, CANON_ZERO_COUNT)) - monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False) - fake_kb = SimpleNamespace(tenant_id="tenant-1") - monkeypatch.setattr("api.db.services.doc_metadata_service.Knowledgebase.get_by_id", lambda kb_id: fake_kb) - - metas = DocMetadataService.get_flatted_meta_by_kbs(["kb-1"]) - doc_ids = meta_filter(metas, [{"key": "canon", "op": "not in", "value": ["0"]}]) - - assert len(doc_ids) == TOTAL_DOCS - CANON_ZERO_COUNT +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Regression test for #16524: a manual metadata filter over a knowledge base +with more documents than the ES push-down cap (``filter_doc_ids_by_meta_pushdown``'s +default ``limit=10000``) must still see every document once the request falls +back to the in-memory path, not just the first page. + +Exercises ``DocMetadataService.get_flatted_meta_by_kbs`` end-to-end against a +fake, paginated ``docStoreConn`` standing in for Elasticsearch, then feeds the +result into ``meta_filter`` with the same ``not in`` condition from the +original report. +""" + +from types import SimpleNamespace + +import pytest + +from common import settings +from common.metadata_utils import meta_filter +from api.db.services.doc_metadata_service import DocMetadataService, METADATA_ID_BATCH_SIZE +from api.db.db_models import DB + +pytestmark = pytest.mark.p2 + +TOTAL_DOCS = 12000 +CANON_ZERO_COUNT = 30 # a small minority tagged "0"; the rest are "1" + + +class _FakeDocStoreConn: + """Stands in for the ES connection's paginated ``search``. + + Mirrors the shape ``DocMetadataService._iter_search_results`` expects + (``{"hits": {"hits": [{"_id": ..., "_source": {...}}]}}``) and actually + honors ``offset``/``limit`` so a caller that stops paginating too early + provably sees a truncated result, the way the reported bug did. + """ + + def __init__(self, total: int, canon_zero_count: int): + self.conditions = [] + self._docs = [] + for i in range(total): + canon = "0" if i < canon_zero_count else "1" + self._docs.append({"_id": f"doc-{i}", "_source": {"meta_fields": {"canon": canon}}}) + + def index_exist(self, index_name, kb_id): + return True + + def search(self, select_fields, highlight_fields, condition, match_expressions, order_by, offset, limit, index_names, knowledgebase_ids, agg_fields=None, rank_feature=None): + self.conditions.append(condition.copy()) + docs = self._docs + if condition.get("id"): + doc_ids = set(condition["id"]) + docs = [doc for doc in docs if doc["_id"] in doc_ids] + page = docs[offset : offset + limit] + return {"hits": {"hits": page, "total": {"value": len(docs)}}} + + +def test_get_flatted_meta_by_kbs_returns_every_document_beyond_pushdown_cap(monkeypatch): + monkeypatch.setattr(DB, "connect", lambda *args, **kwargs: None) + monkeypatch.setattr(DB, "close", lambda *args, **kwargs: None) + monkeypatch.setattr(settings, "docStoreConn", _FakeDocStoreConn(TOTAL_DOCS, CANON_ZERO_COUNT)) + monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False) + fake_kb = SimpleNamespace(tenant_id="tenant-1") + monkeypatch.setattr("api.db.services.doc_metadata_service.Knowledgebase.get_by_id", lambda kb_id: fake_kb) + + metas = DocMetadataService.get_flatted_meta_by_kbs(["kb-1"]) + + assert len(metas["canon"]["1"]) == TOTAL_DOCS - CANON_ZERO_COUNT + assert len(metas["canon"]["0"]) == CANON_ZERO_COUNT + + +def test_manual_not_in_filter_matches_every_document_beyond_pushdown_cap(monkeypatch): + # Same scenario as the #16524 report: a "canon Not in ['0']" manual filter + # over a KB whose match set (TOTAL_DOCS - CANON_ZERO_COUNT) exceeds the + # push-down cap, so this exercises the in-memory fallback exclusively. + monkeypatch.setattr(DB, "connect", lambda *args, **kwargs: None) + monkeypatch.setattr(DB, "close", lambda *args, **kwargs: None) + monkeypatch.setattr(settings, "docStoreConn", _FakeDocStoreConn(TOTAL_DOCS, CANON_ZERO_COUNT)) + monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False) + fake_kb = SimpleNamespace(tenant_id="tenant-1") + monkeypatch.setattr("api.db.services.doc_metadata_service.Knowledgebase.get_by_id", lambda kb_id: fake_kb) + + metas = DocMetadataService.get_flatted_meta_by_kbs(["kb-1"]) + doc_ids = meta_filter(metas, [{"key": "canon", "op": "not in", "value": ["0"]}]) + + assert len(doc_ids) == TOTAL_DOCS - CANON_ZERO_COUNT + + +def test_get_metadata_for_documents_batches_large_id_filters(monkeypatch): + total = METADATA_ID_BATCH_SIZE * 2 + 1 + store = _FakeDocStoreConn(total, 0) + monkeypatch.setattr(DB, "connect", lambda *args, **kwargs: None) + monkeypatch.setattr(DB, "close", lambda *args, **kwargs: None) + monkeypatch.setattr(settings, "docStoreConn", store) + monkeypatch.setattr(settings, "DOC_ENGINE_INFINITY", False) + fake_kb = SimpleNamespace(tenant_id="tenant-1") + monkeypatch.setattr("api.db.services.doc_metadata_service.Knowledgebase.get_by_id", lambda kb_id: fake_kb) + + doc_ids = [f"doc-{index}" for index in range(total)] + metadata = DocMetadataService.get_metadata_for_documents(doc_ids, "kb-1") + + filtered_conditions = [condition for condition in store.conditions if "id" in condition] + requested_batches = {tuple(condition["id"]) for condition in filtered_conditions} + assert len(metadata) == total + assert len(requested_batches) == 3 + assert all(len(condition["id"]) <= METADATA_ID_BATCH_SIZE for condition in filtered_conditions)