zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-27 18:56:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(metadata): batch large document ID filters (#17394)

Browse Source 
## Summary

- batch `DocMetadataService._search_metadata()` requests whose
`condition["id"]` list exceeds 10,000 IDs
- combine the per-batch results so metadata reads, summaries, and batch
updates avoid a single oversized Elasticsearch `terms` query
- add regression coverage for 20,001 IDs, verifying three bounded
backend searches and complete metadata results

Fixes #17393
This commit is contained in:
Jiang, Guomin
2026-08-21 15:21:45 +08:00
committed by GitHub
parent 928743ff40
commit 09af859357
2 changed files with 134 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
api/db/services/doc_metadata_service.py
View File

@@ -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:

212
test/unit_test/api/db/services/test_doc_metadata_service_pagination.py
View File

@@ -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)