From 7c7e6b24898195e3268d9d315abc170d5e7840a2 Mon Sep 17 00:00:00 2001 From: Pim Messelink Date: Sun, 28 Jun 2026 14:47:58 +0000 Subject: [PATCH 1/2] fix(pgvector): skip document column for metadata-only fetches (#1840 follow-up) Closes the explicit "separate follow-up to keep this low-risk" callout in PR #1840's description. For remote pgvector deployments (TLS over WAN), `mempalace_status` and every other metadata-only consumer was transferring the full `document` column over the wire even when nothing read it. A single scroll over a 177K-drawer palace on a 175 ms-RTT link moved ~150 MB of document text plus ~50 MB of metadata; this PR drops that to ~50 MB. scroll_rows / _scroll gain `with_document: bool = True`. When False, SELECT projects NULL::text instead of the document column. Positional _row parser unchanged (record[1] stays the document slot, just receives NULL). Existing callers default to True and see byte-for-byte identical behavior. PgVectorCollection.get_all_metadata override: where=None path goes single-scroll with with_document=False. Filtered path falls back to base to keep _matches_where running on array/object metadata values (same correctness contract as #1840's filtered-path decision). Tests: - Update _FakePgVectorClient.scroll_rows to accept with_document; mirror the NULL-becomes-empty-string semantics when False - Update 5 existing scroll_calls assertions to include with_document=True (unchanged intent) - test_pgvector_get_all_metadata_skips_document_column: assert exactly one scroll call with with_document=False - test_pgvector_get_all_metadata_filtered_falls_back_to_base: assert filtered path preserves with_document=True Co-Authored-By: Claude Opus 4.7 Claude-Session: https://claude.ai/code/session_01KC5Qsknh2zFRtRvVyjXiTA --- mempalace/backends/pgvector.py | 40 ++++++++++++- tests/test_pgvector_backend.py | 101 ++++++++++++++++++++++++++++++--- 2 files changed, 131 insertions(+), 10 deletions(-) diff --git a/mempalace/backends/pgvector.py b/mempalace/backends/pgvector.py index cfe3494507..731489c61c 100644 --- a/mempalace/backends/pgvector.py +++ b/mempalace/backends/pgvector.py @@ -674,13 +674,19 @@ def scroll_rows( *, where: Optional[dict] = None, with_embedding: bool = False, + with_document: bool = True, limit: Optional[int] = None, offset: Optional[int] = None, ) -> list[dict]: qi = _quote_identifier(table) params: list = [] where_sql = _where_to_sql(where, params) if where else "TRUE" - cols = "id, document, metadata" + # Project NULL into the document slot when the caller only needs + # metadata (e.g. mempalace_status's wing/room tally). Keeps the + # positional _row parser unchanged — document remains record[1] — + # while avoiding O(n × document_size) bytes over the wire on remote + # pgvector deployments. Follow-up to #1840. + cols = "id, document, metadata" if with_document else "id, NULL::text, metadata" if with_embedding: cols += ", embedding" sql = f"SELECT {cols} FROM {qi} WHERE {where_sql}" @@ -855,7 +861,15 @@ def _ensure_table(self, dimension: int) -> None: ) self._known_dimension = existing_dim or dimension - def _scroll(self, *, where=None, with_embedding=False, limit=None, offset=None) -> list[dict]: + def _scroll( + self, + *, + where=None, + with_embedding=False, + with_document=True, + limit=None, + offset=None, + ) -> list[dict]: self._ensure_open() if not self._table_exists(): if self._marker_exists(): @@ -865,10 +879,32 @@ def _scroll(self, *, where=None, with_embedding=False, limit=None, offset=None) self._table, where=where, with_embedding=with_embedding, + with_document=with_document, limit=limit, offset=offset, ) + def get_all_metadata(self, where=None) -> list[dict]: + """Single-pass metadata-only fetch — projects out the document column. + + The base implementation pages through ``get(include=["metadatas"])``, + which routes here via ``_scroll`` and (pre-this-override) always sent + the ``document`` text over the wire even when nothing consumed it. + For pgvector deployments where the client is remote (TLS over WAN), + that meant ``mempalace_status`` transferred O(n × document_size) + bytes per call, dominating wall time. With ``with_document=False`` + the SELECT replaces document with NULL, dropping the per-row payload + to id + metadata. + + Filtered fetches (``where`` set) fall back to the base implementation + so the ``_matches_where`` post-filter for array/object values keeps + running unchanged — same correctness contract as #1840's filtered + path. Status, the hot caller, passes no filter. + """ + if where is not None: + return super().get_all_metadata(where=where) + return [row["metadata"] for row in self._scroll(with_document=False)] + def _rows( self, *, diff --git a/tests/test_pgvector_backend.py b/tests/test_pgvector_backend.py index f2c591942b..438a0e062c 100644 --- a/tests/test_pgvector_backend.py +++ b/tests/test_pgvector_backend.py @@ -93,8 +93,19 @@ def query_rows(self, table, *, vector, limit, where, with_embedding): out.append(item) return out - def scroll_rows(self, table, *, where=None, with_embedding=False, limit=None, offset=None): - self.scroll_calls.append({"where": where, "limit": limit, "offset": offset}) + def scroll_rows( + self, + table, + *, + where=None, + with_embedding=False, + with_document=True, + limit=None, + offset=None, + ): + self.scroll_calls.append( + {"where": where, "limit": limit, "offset": offset, "with_document": with_document} + ) rows = self._filtered(table, where) if limit is not None or offset: # Mirror the real backend: ORDER BY id, then LIMIT/OFFSET. @@ -108,7 +119,9 @@ def scroll_rows(self, table, *, where=None, with_embedding=False, limit=None, of out.append( { "id": row["id"], - "document": row["document"], + # Match the real backend: NULL document becomes empty string + # via the SELECT NULL::text projection when with_document=False. + "document": row["document"] if with_document else "", "metadata": row.get("metadata") or {}, "embedding": row.get("embedding") if with_embedding else None, "distance": None, @@ -390,7 +403,7 @@ def test_pgvector_get_unfiltered_page_pushes_limit_offset(tmp_path, fake_pgvecto # An unfiltered page is pushed to SQL as LIMIT/OFFSET instead of fetching # the whole table and slicing in Python (the O(rows x pages) path). - assert client.scroll_calls == [{"where": None, "limit": 2, "offset": 1}] + assert client.scroll_calls == [{"where": None, "limit": 2, "offset": 1, "with_document": True}] # ORDER BY id, then OFFSET 1 LIMIT 2 -> b, c. assert page.ids == ["b", "c"] @@ -410,7 +423,9 @@ def test_pgvector_get_filtered_page_stays_on_full_scan(tmp_path, fake_pgvector): # A filtered get keeps the full-scan path (no LIMIT/OFFSET pushed) so the # exact _matches_where re-filter runs before pagination. - assert client.scroll_calls == [{"where": {"wing": "x"}, "limit": None, "offset": None}] + assert client.scroll_calls == [ + {"where": {"wing": "x"}, "limit": None, "offset": None, "with_document": True} + ] assert page.ids == ["c"] @@ -427,13 +442,17 @@ def test_pgvector_get_offset_only_and_limit_only_push(tmp_path, fake_pgvector): # offset-only (limit=None) is pushed. client.scroll_calls.clear() page = col.get(offset=2, include=["metadatas"]) - assert client.scroll_calls == [{"where": None, "limit": None, "offset": 2}] + assert client.scroll_calls == [ + {"where": None, "limit": None, "offset": 2, "with_document": True} + ] assert page.ids == ["c", "d"] # limit-only (offset=None) is pushed. client.scroll_calls.clear() page = col.get(limit=2, include=["metadatas"]) - assert client.scroll_calls == [{"where": None, "limit": 2, "offset": None}] + assert client.scroll_calls == [ + {"where": None, "limit": 2, "offset": None, "with_document": True} + ] assert page.ids == ["a", "b"] @@ -451,7 +470,9 @@ def test_pgvector_get_negative_bounds_use_python_slice(tmp_path, fake_pgvector): # A negative offset must not reach SQL (OFFSET -1 would error); it falls # through to the unchanged full-scan + Python-slice path. page = col.get(offset=-1, include=["metadatas"]) - assert client.scroll_calls == [{"where": None, "limit": None, "offset": None}] + assert client.scroll_calls == [ + {"where": None, "limit": None, "offset": None, "with_document": True} + ] assert page.ids == ["c"] @@ -473,6 +494,70 @@ def test_pgvector_get_pages_tile_without_overlap(tmp_path, fake_pgvector): assert p1 + p2 + p3 == ["a", "b", "c", "d", "e"] +def test_pgvector_get_all_metadata_skips_document_column(tmp_path, fake_pgvector): + """The metadata-only fast path must NOT pull document text over the wire. + + Default base ``get_all_metadata`` pages through ``get(include=["metadatas"])``, + which used to route here via scroll_rows with documents always selected — the + "separate follow-up" #1840 flagged. This override calls scroll_rows with + with_document=False so the SELECT projects NULL into the document slot, + dropping per-row payload for remote (TLS over WAN) clients where status + otherwise dominates wall time. + """ + _backend, col = _collection(tmp_path) + col.add( + ids=["a", "b", "c"], + documents=["doc_a", "doc_b", "doc_c"], + metadatas=[ + {"wing": "p", "room": "backend"}, + {"wing": "p", "room": "frontend"}, + {"wing": "q", "room": "backend"}, + ], + embeddings=[[1, 0], [0, 1], [0.5, 0.5]], + ) + client = fake_pgvector.instances[0] + client.scroll_calls.clear() + + metas = col.get_all_metadata() + + # Exactly one scroll, with_document=False (no document text on the wire). + assert client.scroll_calls == [ + {"where": None, "limit": None, "offset": None, "with_document": False} + ] + # Returns just the metadata dicts (full set, any order — sort by wing+room for stability). + metas_sorted = sorted(metas, key=lambda m: (m["wing"], m["room"])) + assert metas_sorted == [ + {"wing": "p", "room": "backend"}, + {"wing": "p", "room": "frontend"}, + {"wing": "q", "room": "backend"}, + ] + + +def test_pgvector_get_all_metadata_filtered_falls_back_to_base(tmp_path, fake_pgvector): + """Filtered get_all_metadata keeps the base offset-loop + _matches_where path. + + The metadata @> ... pushdown is broader than _matches_where for array/object + values (same correctness contract as #1840's filtered path), so the post-filter + must run. Falling back to super() preserves it. + """ + _backend, col = _collection(tmp_path) + col.add( + ids=["a", "b", "c"], + documents=["doc_a", "doc_b", "doc_c"], + metadatas=[{"wing": "x"}, {"wing": "y"}, {"wing": "x"}], + embeddings=[[1, 0], [0, 1], [0.5, 0.5]], + ) + client = fake_pgvector.instances[0] + client.scroll_calls.clear() + + metas = col.get_all_metadata(where={"wing": "x"}) + + # Filtered path goes through base get() pagination, so documents ARE pulled + # (with_document=True). Correctness > perf for filtered queries. + assert all(call["with_document"] is True for call in client.scroll_calls) + assert sorted(metas, key=lambda m: m["wing"]) == [{"wing": "x"}, {"wing": "x"}] + + def test_pgvector_delete_by_where_pushdown_and_local(tmp_path, fake_pgvector): _backend, col = _collection(tmp_path) col.add( From 3c1d8d0e67414948f815679a69ed082684bd24cb Mon Sep 17 00:00:00 2001 From: Pim Messelink Date: Sun, 28 Jun 2026 14:57:35 +0000 Subject: [PATCH 2/2] fix(pgvector): extend with_document=False fast path to filtered get_all_metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per gemini-code-assist review feedback on #1892: _matches_where only reads metadata, so the where=None vs where=set conditional fall-back was unnecessary. The filtered path can use the same single-scroll with_document=False fast path and apply the post-filter locally on metadata dicts — extending the wire-byte win to every get_all_metadata caller, not just unfiltered ones. Mirrors the pushdown + local _matches_where pattern already used by _rows in the same file: pushdown when _requires_local_filter is False, post-filter in Python otherwise. Same correctness contract as #1840's filtered get path. Renames test_pgvector_get_all_metadata_filtered_falls_back_to_base to test_pgvector_get_all_metadata_filtered_uses_fast_path and asserts the new behavior (with_document=False + pushdown forwards the equality filter to SQL). Co-Authored-By: Claude Opus 4.7 Claude-Session: https://claude.ai/code/session_01KC5Qsknh2zFRtRvVyjXiTA --- mempalace/backends/pgvector.py | 25 ++++++++++++++++--------- tests/test_pgvector_backend.py | 19 +++++++++++-------- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/mempalace/backends/pgvector.py b/mempalace/backends/pgvector.py index 731489c61c..b0632c8c55 100644 --- a/mempalace/backends/pgvector.py +++ b/mempalace/backends/pgvector.py @@ -894,16 +894,23 @@ def get_all_metadata(self, where=None) -> list[dict]: that meant ``mempalace_status`` transferred O(n × document_size) bytes per call, dominating wall time. With ``with_document=False`` the SELECT replaces document with NULL, dropping the per-row payload - to id + metadata. - - Filtered fetches (``where`` set) fall back to the base implementation - so the ``_matches_where`` post-filter for array/object values keeps - running unchanged — same correctness contract as #1840's filtered - path. Status, the hot caller, passes no filter. + to id + metadata for every caller of this method. + + Filtered fetches still need the ``_matches_where`` post-filter for + non-pushdown semantics (array/object values where ``metadata @> ...`` + is broader than the exact match the caller asked for — same + correctness contract as #1840's filtered ``get`` path). Since that + post-filter only reads ``metadata``, we keep the single-scroll + + ``with_document=False`` fast path and just apply the filter locally + on the metadata dicts before returning. This extends the wire-byte + win to filtered callers as well. """ - if where is not None: - return super().get_all_metadata(where=where) - return [row["metadata"] for row in self._scroll(with_document=False)] + _validate_where(where) + pushdown = None if _requires_local_filter(where) else where + rows = self._scroll(where=pushdown, with_document=False) + if where is None: + return [row["metadata"] for row in rows] + return [row["metadata"] for row in rows if _matches_where(row["metadata"], where)] def _rows( self, diff --git a/tests/test_pgvector_backend.py b/tests/test_pgvector_backend.py index 438a0e062c..9505a0bba4 100644 --- a/tests/test_pgvector_backend.py +++ b/tests/test_pgvector_backend.py @@ -533,12 +533,13 @@ def test_pgvector_get_all_metadata_skips_document_column(tmp_path, fake_pgvector ] -def test_pgvector_get_all_metadata_filtered_falls_back_to_base(tmp_path, fake_pgvector): - """Filtered get_all_metadata keeps the base offset-loop + _matches_where path. +def test_pgvector_get_all_metadata_filtered_uses_fast_path(tmp_path, fake_pgvector): + """Filtered get_all_metadata uses the single-pass metadata-only fast path. - The metadata @> ... pushdown is broader than _matches_where for array/object - values (same correctness contract as #1840's filtered path), so the post-filter - must run. Falling back to super() preserves it. + ``_matches_where`` only reads ``metadata``, so we keep ``with_document=False`` + and apply the post-filter locally on the metadata dicts. SQL pushdown still + happens when the filter is pushdownable; the local ``_matches_where`` re-runs + for array/object semantics #1840's filtered path required. """ _backend, col = _collection(tmp_path) col.add( @@ -552,9 +553,11 @@ def test_pgvector_get_all_metadata_filtered_falls_back_to_base(tmp_path, fake_pg metas = col.get_all_metadata(where={"wing": "x"}) - # Filtered path goes through base get() pagination, so documents ARE pulled - # (with_document=True). Correctness > perf for filtered queries. - assert all(call["with_document"] is True for call in client.scroll_calls) + # Exactly one scroll with with_document=False — pushdown forwards the + # equality filter to SQL; no document text on the wire. + assert client.scroll_calls == [ + {"where": {"wing": "x"}, "limit": None, "offset": None, "with_document": False} + ] assert sorted(metas, key=lambda m: m["wing"]) == [{"wing": "x"}, {"wing": "x"}]