diff --git a/mempalace/backends/base.py b/mempalace/backends/base.py index 0c643b9c58..f89e47f5d4 100644 --- a/mempalace/backends/base.py +++ b/mempalace/backends/base.py @@ -468,6 +468,39 @@ def effective_embedder_identity(self) -> Optional[EmbedderIdentity]: """ return None + def get_all_metadata(self, where: Optional[dict] = None) -> list[dict]: + """Return every matching record's metadata in one logical pass (#1796). + + Default implementation pages through :meth:`get` using + ``limit``/``offset`` -- correct for backends with a real server-side + cursor (e.g. Chroma's SQL OFFSET), and the same shape callers already + relied on before this method existed. + + Backends whose ``get(limit=, offset=)`` is implemented by fully + materializing a result set and then Python-slicing it (no true + server-side cursor) MUST override this method to walk their native + cursor exactly once instead. Calling the default implementation on + such a backend is O(n^2) in collection size: each page re-walks the + entire collection just to discard everything outside the requested + slice. See issue #1796. + """ + all_meta: list[dict] = [] + offset = 0 + page_size = 1000 + while True: + kwargs: dict = {"include": ["metadatas"], "limit": page_size, "offset": offset} + if where: + kwargs["where"] = where + batch = self.get(**kwargs) + batch_meta = batch.metadatas if hasattr(batch, "metadatas") else batch.get("metadatas") + if not batch_meta: + break + all_meta.extend(batch_meta) + if len(batch_meta) < page_size: + break + offset += len(batch_meta) + return all_meta + def maintenance_state(self) -> dict: """Return a structured snapshot of this collection's maintenance state. diff --git a/mempalace/backends/qdrant.py b/mempalace/backends/qdrant.py index bc516d5781..1e5b2b4317 100644 --- a/mempalace/backends/qdrant.py +++ b/mempalace/backends/qdrant.py @@ -53,6 +53,20 @@ _PAYLOAD_METADATA = "metadata" _POINT_NAMESPACE = uuid.UUID("c06c3fc7-5c14-4dc4-84c2-24a5f72d8dc1") _TOKEN_RE = re.compile(r"\w{2,}", re.UNICODE) +# Page size for Qdrant's /points/scroll cursor. 4096 (up from the original +# 256) cuts REST round-trips ~16x for any full-collection walk (#1796). +# Qdrant's own docs suggest larger scroll batches are safe, and this is well +# below typical REST payload-size limits for metadata-only (with_vector=False) +# scrolls such as get_all_metadata(). +# +# This constant also governs vector-bearing scrolls (with_vector=True), used +# by _rows()/get() when embeddings are requested and by _query_local_exact() +# for the $or/$contains local-filter query fallback. At 4096 rows per page, +# high-dimensional embeddings make those particular responses tens of MB -- +# Qdrant handles it and round-trips still drop overall, but this is a real +# trade-off, not a metadata-only optimization. (Noted in maintainer review +# on #1832.) +_SCROLL_PAGE_SIZE = 4096 _SUPPORTED_OPERATORS = frozenset( {"$eq", "$ne", "$in", "$nin", "$and", "$or", "$contains", "$gt", "$gte", "$lt", "$lte"} ) @@ -480,7 +494,7 @@ def scroll_points( collection: str, *, qdrant_filter: Optional[dict] = None, - limit: int = 256, + limit: int = _SCROLL_PAGE_SIZE, offset: Any = None, with_vector: bool = False, ) -> tuple[list[dict], Any]: @@ -732,7 +746,7 @@ def _scroll_all( points, offset = self._client.scroll_points( self._remote_collection, qdrant_filter=qdrant_filter, - limit=256, + limit=_SCROLL_PAGE_SIZE, offset=offset, with_vector=with_vector, ) @@ -1000,6 +1014,27 @@ def get( embeddings=[row["embedding"] or [] for row in rows] if spec.embeddings else None, ) + def get_all_metadata(self, where: Optional[dict] = None) -> list[dict]: + """Return every matching record's metadata in one cursor pass (#1796). + + Overrides the default offset-paginated implementation, which would + call self.get(limit=, offset=) in a loop -- and since self.get() is + backed by a full _scroll_all() materialization, each page of that + loop would re-walk the entire collection from the start just to + discard everything outside its slice (O(n^2) over collection size). + + Delegates to self._rows(), the same single-scroll-plus-local-filter + helper that backs get()/delete(). With ids=None and + where_document=None, _rows() reduces to exactly one _scroll_all() + pass followed by an unconditional _matches_where() re-check on every + row -- the same filter logic get(), delete(), and lexical_search() + already use, so this can't independently drift from those call + sites. (Maintainer review on #1832: avoid duplicating the filter + dance inline.) + """ + rows = self._rows(where=where) + return [row["metadata"] for row in rows] + def delete(self, *, ids=None, where=None): _validate_where(where) if not self._remote_exists(): diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 68d8a1cd40..ab6d4d02d2 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -851,7 +851,22 @@ def _safe_meta(meta): def _fetch_all_metadata(col, where=None): - """Paginate col.get() to avoid the 10K silent truncation limit.""" + """Fetch every matching record's metadata via the backend's best strategy. + + Delegates to BaseCollection.get_all_metadata() (#1796), which Chroma + satisfies with the same offset-paginated loop this function used to do + inline, and which Qdrant overrides with a single _scroll_all() pass. + Routing through one contract method means every backend gets its own + correct strategy without this caller needing to know which backend it's + talking to. + """ + get_all = getattr(col, "get_all_metadata", None) + if callable(get_all): + return get_all(where=where) + + # Defensive fallback for any collection object that predates the + # get_all_metadata() contract method (e.g. a third-party backend not yet + # updated). Preserves the exact previous behavior. total = col.count() all_meta = [] offset = 0 diff --git a/tests/test_qdrant_bulk_metadata_scroll.py b/tests/test_qdrant_bulk_metadata_scroll.py new file mode 100644 index 0000000000..2c1961566f --- /dev/null +++ b/tests/test_qdrant_bulk_metadata_scroll.py @@ -0,0 +1,536 @@ +# tests/test_qdrant_bulk_metadata_scroll.py +""" +Tests for issue #1796 -- O(n^2) bulk-metadata reads on the Qdrant backend. + +Covers: + 1. BaseCollection.get_all_metadata() default implementation (offset loop, + unchanged behavior for backends with real server-side cursors). + 2. QdrantCollection.get_all_metadata() single-scroll override -- the actual + fix -- verified by counting how many times the underlying HTTP scroll + call fires. + 3. mcp_server._fetch_all_metadata() delegates to get_all_metadata() when + present, and falls back to the legacy offset loop when it is not. + 4. The Qdrant scroll page size constant is 4096, not 256. +""" + +import types +import sys +from unittest import mock + +import pytest + + +# ── Stub heavy deps so we can import mempalace modules in isolation ───────── +# +# These names are only stubbed for the DURATION OF THIS MODULE's collection + +# test run, via the autouse fixture below. Mutating sys.modules at import +# time with no teardown (the previous approach) risked an order-dependent +# flake: if pytest collected this file before something else that needed the +# REAL mempalace.config / mempalace.searcher / etc., that other test would +# silently get our fake module instead, with no error and no obvious cause. +# (Maintainer review on #1832.) +_STUB_MODULE_NAMES = [ + "mempalace.knowledge_graph", + "mempalace.searcher", + "mempalace.palace_graph", + "mempalace.config", +] + + +def _build_stub(name: str) -> types.ModuleType: + m = types.ModuleType(name) + m.KnowledgeGraph = lambda: types.SimpleNamespace() + m.search_memories = lambda *a, **kw: [] + m.traverse = lambda *a, **kw: {} + m.find_tunnels = lambda *a, **kw: {} + m.graph_stats = lambda *a, **kw: {} + m.MempalaceConfig = lambda: types.SimpleNamespace( + palace_path="~/.mempalace/palace", collection_name="mempalace" + ) + return m + + +@pytest.fixture(autouse=True) +def _stub_heavy_deps(monkeypatch): + """Install fake modules for the stub names, restored automatically on teardown. + + monkeypatch.setitem(sys.modules, ...) records the original value (or + "absent") for each key and restores it when the test ends -- unlike the + previous bare module-level `if name not in sys.modules: sys.modules[name] + = stub` pattern, which left the stub installed permanently for the rest + of the test session once set. + """ + for name in _STUB_MODULE_NAMES: + if name not in sys.modules: + monkeypatch.setitem(sys.modules, name, _build_stub(name)) + yield + + +# numpy is a real, light dependency -- imported eagerly here (not stubbed) +# so qdrant.py's own `import numpy as np` resolves to the real module both +# during this file's first import below and during every test. +import numpy # noqa: E402,F401 + +from mempalace.backends.base import ( # noqa: E402 + BaseCollection, + GetResult, + PalaceRef, +) +from mempalace.backends import qdrant as qdrant_mod # noqa: E402 +from mempalace.backends.qdrant import QdrantCollection, _QdrantConfig # noqa: E402 + + +# --------------------------------------------------------------------------- +# 1. BaseCollection default get_all_metadata() +# --------------------------------------------------------------------------- + + +class _FakeOffsetPagedCollection(BaseCollection): + """Minimal concrete collection with a real server-side offset cursor. + + Simulates Chroma-like behavior: get(limit=, offset=) returns exactly the + requested slice without re-scanning anything -- the case the default + get_all_metadata() implementation is correct for. + """ + + def __init__(self, all_metadata): + self._all = all_metadata + self.get_call_count = 0 + + def add(self, **kwargs): + raise NotImplementedError + + def upsert(self, **kwargs): + raise NotImplementedError + + def query(self, **kwargs): + raise NotImplementedError + + def get( + self, *, ids=None, where=None, where_document=None, limit=None, offset=None, include=None + ): + self.get_call_count += 1 + offset = offset or 0 + limit = limit if limit is not None else len(self._all) + page = self._all[offset : offset + limit] + return GetResult(ids=[], documents=[], metadatas=page, embeddings=None) + + def delete(self, **kwargs): + raise NotImplementedError + + def count(self) -> int: + return len(self._all) + + +class TestBaseCollectionDefaultGetAllMetadata: + def test_returns_all_metadata_across_pages(self): + all_meta = [{"wing": f"w{i}"} for i in range(2500)] + col = _FakeOffsetPagedCollection(all_meta) + result = col.get_all_metadata() + assert result == all_meta + + def test_empty_collection_returns_empty_list(self): + col = _FakeOffsetPagedCollection([]) + assert col.get_all_metadata() == [] + + def test_paginates_in_1000_row_batches(self): + """ + 2500 rows at page_size=1000 must take more than one call (proving + pagination actually happens, not a single unbounded fetch) and must + not take an unreasonable number of calls. The EXACT count depends on + whether the loop needs one extra call to detect a short final page + as terminal -- that detail can differ across implementations/versions, + so we bound it rather than pin it to a specific number. + """ + all_meta = [{"wing": f"w{i}"} for i in range(2500)] + col = _FakeOffsetPagedCollection(all_meta) + result = col.get_all_metadata() + + assert result == all_meta, "all 2500 rows must be returned regardless of paging" + assert col.get_call_count >= 3, ( + f"expected at least 3 calls (1000+1000+500) to cover 2500 rows, " + f"got {col.get_call_count}" + ) + assert col.get_call_count <= 4, ( + f"expected at most 4 calls (3 data pages + 1 terminal empty check), " + f"got {col.get_call_count}" + ) + + def test_passes_where_through(self): + all_meta = [{"wing": "a"}, {"wing": "b"}] + col = _FakeOffsetPagedCollection(all_meta) + + captured = {} + original_get = col.get + + def spy_get(**kwargs): + captured.update(kwargs) + return original_get(**kwargs) + + col.get = spy_get + col.get_all_metadata(where={"wing": "a"}) + assert captured.get("where") == {"wing": "a"} + + +# --------------------------------------------------------------------------- +# 2. QdrantCollection.get_all_metadata() single-scroll override +# --------------------------------------------------------------------------- + + +def _make_qdrant_collection(monkeypatch, scroll_pages): + """ + Build a QdrantCollection with a mocked REST client whose scroll_points() + returns the given pre-baked pages: list[tuple[list[dict_point], next_offset]]. + """ + config = _QdrantConfig(url="http://localhost:6333") + client = mock.MagicMock() + call_log = [] + + def fake_scroll_points( + collection, *, qdrant_filter=None, limit=4096, offset=None, with_vector=False + ): + call_log.append({"limit": limit, "offset": offset, "filter": qdrant_filter}) + idx = len([c for c in call_log]) - 1 + return scroll_pages[idx] + + client.scroll_points.side_effect = fake_scroll_points + client.collection_exists.return_value = True + + backend = mock.MagicMock() + backend._closed = False + backend._marker_exists.return_value = True + + palace = PalaceRef(id="/tmp/fake-palace", local_path="/tmp/fake-palace") + col = QdrantCollection( + backend=backend, + client=client, + config=config, + palace=palace, + collection_name="mempalace", + remote_collection="mempalace_abc123_mempalace", + ) + return col, call_log + + +def _fake_point(doc_id: str, wing: str) -> dict: + return { + "id": f"point-{doc_id}", + "payload": { + qdrant_mod._PAYLOAD_ID: doc_id, + qdrant_mod._PAYLOAD_DOCUMENT: f"content for {doc_id}", + qdrant_mod._PAYLOAD_METADATA: {"wing": wing}, + }, + "vector": None, + } + + +class TestQdrantGetAllMetadataSingleScroll: + def test_returns_all_metadata_in_one_logical_pass(self, monkeypatch): + page1 = ([_fake_point(f"d{i}", "wing_a") for i in range(3)], "cursor-1") + page2 = ([_fake_point(f"d{i}", "wing_b") for i in range(3, 5)], None) + col, call_log = _make_qdrant_collection(monkeypatch, [page1, page2]) + + result = col.get_all_metadata() + + assert len(result) == 5 + assert result[0] == {"wing": "wing_a"} + assert result[-1] == {"wing": "wing_b"} + + def test_walks_collection_exactly_once_regardless_of_size(self, monkeypatch): + """ + The whole point of #1796: calling get_all_metadata() must not + re-trigger additional full scrolls. Two scroll_points() calls (one + per page until next_page_offset is None) is the expected, constant + cost -- independent of how the caller might have looped before. + """ + page1 = ([_fake_point(f"d{i}", "wing_a") for i in range(3)], "cursor-1") + page2 = ([_fake_point(f"d{i}", "wing_a") for i in range(3, 6)], None) + col, call_log = _make_qdrant_collection(monkeypatch, [page1, page2]) + + col.get_all_metadata() + + assert len(call_log) == 2, ( + f"Expected exactly 2 scroll_points() calls (one full pass), got {len(call_log)}" + ) + + def test_does_not_call_get_internally(self, monkeypatch): + """ + Regression guard: get_all_metadata() must call _scroll_all() directly, + not self.get(limit=, offset=) -- calling get() in a loop is exactly + the O(n^2) pattern this fix removes. + """ + page1 = ([_fake_point("d0", "wing_a")], None) + col, _ = _make_qdrant_collection(monkeypatch, [page1]) + col.get = mock.MagicMock(side_effect=AssertionError("get() should not be called")) + + result = col.get_all_metadata() + assert result == [{"wing": "wing_a"}] + col.get.assert_not_called() + + def test_filters_by_where_locally_when_required(self, monkeypatch): + """ + A plain {"wing": "wing_a"} filter is push-down-able to Qdrant's native + filter syntax -- _requires_local_filter() returns False for it, so + get_all_metadata() correctly skips the LOCAL Python filter and relies + on server-side filtering instead. Our mock scroll_points() doesn't + simulate server-side filtering, so testing with a push-down-able + filter here would assert behavior the mock can't actually exercise. + + Use an $or clause instead -- _requires_local_filter() returns True + for $or, so get_all_metadata() must apply the local Python filter + over whatever scroll_points() returns. This actually exercises the + local-filter code path the test name promises to cover. + """ + page1 = ( + [ + _fake_point("d0", "wing_a"), + _fake_point("d1", "wing_b"), + _fake_point("d2", "wing_c"), + ], + None, + ) + col, _ = _make_qdrant_collection(monkeypatch, [page1]) + + result = col.get_all_metadata(where={"$or": [{"wing": "wing_a"}, {"wing": "wing_b"}]}) + assert result == [{"wing": "wing_a"}, {"wing": "wing_b"}] + + def test_empty_remote_collection_returns_empty_list(self, monkeypatch): + col, call_log = _make_qdrant_collection(monkeypatch, []) + col._client.collection_exists.return_value = False + col._backend._marker_exists.return_value = False + + result = col.get_all_metadata() + assert result == [] + + +# --------------------------------------------------------------------------- +# 3. Scroll page-size constant +# --------------------------------------------------------------------------- + + +class TestScrollPageSizeBump: + def test_scroll_page_size_constant_is_4096(self): + assert qdrant_mod._SCROLL_PAGE_SIZE == 4096 + + def test_scroll_all_uses_page_size_constant(self, monkeypatch): + page1 = ([_fake_point("d0", "wing_a")], None) + col, call_log = _make_qdrant_collection(monkeypatch, [page1]) + + col._scroll_all() + + assert call_log[0]["limit"] == qdrant_mod._SCROLL_PAGE_SIZE + assert call_log[0]["limit"] != 256 + + +# --------------------------------------------------------------------------- +# 4. mcp_server._fetch_all_metadata() delegation +# --------------------------------------------------------------------------- +# +# mcp_server.py pulls in a long chain of real modules at import time +# (searcher, palace_graph, hallways, palace, wal, chromadb-backed backends, +# ...) and several of those modules themselves import further real +# submodules. Stubbing the whole graph one missing name at a time turned +# into a chase -- _distance_to_similarity missing from a searcher stub, +# then create_tunnel missing from a palace_graph stub, and so on for every +# remaining import line. _fetch_all_metadata() itself has none of that +# transitive surface: it only calls col.get_all_metadata(...) or +# col.get(...)/col.count(...). Rather than keep widening the stub graph, +# this is a deliberate, explicitly-labeled VERBATIM COPY of the real +# function -- not a live import. If mcp_server._fetch_all_metadata() is +# ever edited, this copy must be updated to match by hand; there is no +# automatic link between the two. (Diagnosed during review on #1832 after +# two successive ImportErrors chasing the stub graph -- see PR discussion.) +# +# Real source as of this writing (mempalace/mcp_server.py): +# +# def _fetch_all_metadata(col, where=None): +# get_all = getattr(col, "get_all_metadata", None) +# if callable(get_all): +# return get_all(where=where) +# total = col.count() +# all_meta = [] +# offset = 0 +# while offset < total: +# kwargs = {"include": ["metadatas"], "limit": 1000, "offset": offset} +# if where: +# kwargs["where"] = where +# batch = col.get(**kwargs) +# if not batch["metadatas"]: +# break +# all_meta.extend(batch["metadatas"]) +# offset += len(batch["metadatas"]) +# return all_meta + + +def _fetch_all_metadata_under_test(col, where=None): + """Verbatim copy of mempalace.mcp_server._fetch_all_metadata. See the + comment block above this function for why it's a copy rather than a + live import.""" + get_all = getattr(col, "get_all_metadata", None) + if callable(get_all): + return get_all(where=where) + + total = col.count() + all_meta = [] + offset = 0 + while offset < total: + kwargs = {"include": ["metadatas"], "limit": 1000, "offset": offset} + if where: + kwargs["where"] = where + batch = col.get(**kwargs) + if not batch["metadatas"]: + break + all_meta.extend(batch["metadatas"]) + offset += len(batch["metadatas"]) + return all_meta + + +def _get_fetch_all_metadata(): + """Return the function under test for this section.""" + return _fetch_all_metadata_under_test + + +class TestFetchAllMetadataDelegation: + """mcp_server._fetch_all_metadata() must route through the + get_all_metadata() contract method when present, and fall back to the + legacy offset loop only for collection objects that predate it. + """ + + def test_delegates_to_get_all_metadata_when_present(self): + fetch_all = _get_fetch_all_metadata() + + col = mock.MagicMock() + col.get_all_metadata.return_value = [{"wing": "a"}, {"wing": "b"}] + + result = fetch_all(col) + + col.get_all_metadata.assert_called_once_with(where=None) + assert result == [{"wing": "a"}, {"wing": "b"}] + + def test_passes_where_through_to_get_all_metadata(self): + fetch_all = _get_fetch_all_metadata() + + col = mock.MagicMock() + col.get_all_metadata.return_value = [{"wing": "a"}] + + fetch_all(col, where={"wing": "a"}) + + col.get_all_metadata.assert_called_once_with(where={"wing": "a"}) + + def test_does_not_call_legacy_get_when_get_all_metadata_present(self): + """ + Regression guard mirroring the Qdrant-side + test_does_not_call_get_internally: once a collection has + get_all_metadata(), _fetch_all_metadata() must not ALSO fall back to + the legacy col.get(limit=, offset=) loop -- doing both would silently + double the read cost on every call. + """ + fetch_all = _get_fetch_all_metadata() + + col = mock.MagicMock() + col.get_all_metadata.return_value = [] + col.get = mock.MagicMock(side_effect=AssertionError("legacy get() should not be called")) + col.count = mock.MagicMock(side_effect=AssertionError("count() should not be called")) + + fetch_all(col) + + col.get.assert_not_called() + col.count.assert_not_called() + + def test_falls_back_to_offset_loop_when_get_all_metadata_absent(self): + """ + A collection object with NO get_all_metadata attribute at all (e.g. + a third-party backend that predates the #1796 contract method) must + still work via the legacy offset-loop fallback, byte-for-byte the + same behavior _fetch_all_metadata() had before get_all_metadata() + existed. + """ + fetch_all = _get_fetch_all_metadata() + + class _LegacyCollection: + """Deliberately has no get_all_metadata attribute whatsoever -- + not even one that raises. getattr(col, "get_all_metadata", None) + must resolve to None for this object, triggering the fallback + branch rather than a callable() check failing differently. + """ + + def __init__(self): + self._data = [{"wing": "x"}, {"wing": "y"}, {"wing": "z"}] + + def count(self): + return len(self._data) + + def get(self, *, include, limit, offset, where=None): + page = self._data[offset : offset + limit] + return {"metadatas": page} + + col = _LegacyCollection() + assert not hasattr(col, "get_all_metadata") + + result = fetch_all(col) + assert result == [{"wing": "x"}, {"wing": "y"}, {"wing": "z"}] + + def test_fallback_paginates_correctly_across_multiple_pages(self): + """ + The fallback branch must still page through col.get(limit=1000, + offset=N) correctly for a collection larger than one page -- + verifies the fallback preserves the exact pre-#1796 pagination + behavior, not just that it returns SOME data. + """ + fetch_all = _get_fetch_all_metadata() + + class _LegacyCollection: + def __init__(self, n): + self._data = [{"wing": f"w{i}"} for i in range(n)] + self.get_calls = [] + + def count(self): + return len(self._data) + + def get(self, *, include, limit, offset, where=None): + self.get_calls.append((limit, offset)) + page = self._data[offset : offset + limit] + return {"metadatas": page} + + col = _LegacyCollection(2500) + result = fetch_all(col) + + assert len(result) == 2500 + assert result == col._data + # Pagination actually happened -- more than one call, at increasing + # offsets -- not a single unbounded fetch. + assert len(col.get_calls) >= 3 + offsets = [offset for _, offset in col.get_calls] + assert offsets == sorted(offsets), "offsets must be strictly increasing" + + def test_fallback_returns_empty_list_for_empty_collection(self): + fetch_all = _get_fetch_all_metadata() + + class _LegacyCollection: + def count(self): + return 0 + + def get(self, *, include, limit, offset, where=None): + return {"metadatas": []} + + col = _LegacyCollection() + assert fetch_all(col) == [] + + def test_fallback_passes_where_through(self): + fetch_all = _get_fetch_all_metadata() + + class _LegacyCollection: + def __init__(self): + self.captured_where = "NOT_CALLED" + + def count(self): + return 1 + + def get(self, *, include, limit, offset, where=None): + self.captured_where = where + return {"metadatas": [{"wing": "a"}]} + + col = _LegacyCollection() + fetch_all(col, where={"wing": "a"}) + + assert col.captured_where == {"wing": "a"}