Merge pull request #2362 from MemPalace/fix/sqlite-exact-cache-and-id-filters

fix(sqlite_exact): cross-handle cache invalidation and filtered id lookups

Author: igorls

mempalace/backends/sqlite_exact.py

@@ -464,8 +464,10 @@ def __init__(
         # collection_id -> (ids, float32 matrix, mini-metadata). Filled lazily
         # by query() so a long-lived hub does not re-read every embedding blob
         # on the next search. Mini-metadata is wing/room/source_file for
-        # in-memory equality filters. Cleared on any write.
+        # in-memory equality filters. Cleared on any write through this handle;
+        # ``_vector_cache_data_version`` detects commits from other handles.
         self._vector_cache: dict[int, tuple[list[str], np.ndarray, list[dict]]] = {}
+        self._vector_cache_data_version: Optional[int] = None
 
 
 class SQLiteExactCollection(BaseCollection):
@@ -958,6 +960,10 @@ def _rank_vectors(
         _validate_where(where_document)
         expected = self._collection_dimension(cur, collection_id)
         empty = np.zeros((0, expected or 0), dtype=np.float32)
+        data_version = int(cur.execute("PRAGMA data_version").fetchone()[0])
+        if self._handle._vector_cache_data_version != data_version:
+            self._handle._vector_cache.clear()
+            self._handle._vector_cache_data_version = data_version
         cached = self._handle._vector_cache.get(collection_id)
         if cached is None:
             cached = self._load_all_vectors(cur, collection_id, expected)
@@ -1054,11 +1060,25 @@ def get(
     ) -> GetResult:
         spec = _IncludeSpec.resolve(include, default_distances=False)
         # get(ids=...) must not scan the collection: look up by primary key.
-        if ids is not None and where is None and where_document is None:
+        if ids is not None:
+            _validate_where(where)
+            _validate_where(where_document)
+            lookup_spec = _IncludeSpec(
+                documents=spec.documents or bool(where_document),
+                metadatas=spec.metadatas or bool(where),
+                distances=False,
+                embeddings=spec.embeddings,
+            )
             with self._cursor() as cur:
                 collection_id = self._collection_id(cur)
-                by_id = self._rows_by_ids(cur, collection_id, list(ids), spec)
-            rows = [by_id[doc_id] for doc_id in ids if doc_id in by_id]
+                by_id = self._rows_by_ids(cur, collection_id, list(ids), lookup_spec)
+            rows = [
+                by_id[doc_id]
+                for doc_id in ids
+                if doc_id in by_id
+                and _matches_where(by_id[doc_id]["metadata"], where)
+                and _matches_where_document(by_id[doc_id]["document"], where_document)
+            ]
             if offset:
                 rows = rows[offset:]
             if limit is not None:

tests/test_sqlite_exact_backend.py

@@ -145,6 +145,37 @@ def test_sqlite_exact_get_preserves_requested_id_order_and_duplicates(tmp_path):
     assert result.documents == ["doc b", "doc a", "doc b"]
 
 
+@pytest.mark.parametrize(
+    "filters",
+    [
+        {"where": {"wing": "keep"}},
+        {"where_document": {"$contains": "needle"}},
+        {
+            "where": {"wing": "keep"},
+            "where_document": {"$contains": "needle"},
+        },
+    ],
+)
+def test_sqlite_exact_get_ids_intersects_filters(tmp_path, filters):
+    _backend, col = _collection(tmp_path)
+    col.add(
+        ids=["requested", "not-requested", "filtered-out"],
+        documents=["needle requested", "needle other", "different"],
+        metadatas=[{"wing": "keep"}, {"wing": "keep"}, {"wing": "drop"}],
+        embeddings=[[1, 0], [0, 1], [0.5, 0.5]],
+    )
+
+    result = col.get(
+        ids=["filtered-out", "requested", "requested"],
+        include=[],
+        **filters,
+    )
+
+    assert result.ids == ["requested", "requested"]
+    assert result.documents == []
+    assert result.metadatas == []
+
+
 def _doc_select_sql(col, action):
     """Run ``action`` while tracing SQL; return (result, [documents SELECTs]).
 
@@ -1211,6 +1242,37 @@ def test_sqlite_exact_query_cache_invalidates_on_add(tmp_path):
     assert second.ids[0] == ["new"]
 
 
+def test_sqlite_exact_query_cache_invalidates_after_external_handle_write(tmp_path):
+    reader_backend, reader = _collection(tmp_path)
+    reader.add(
+        ids=["old"],
+        documents=["old"],
+        metadatas=[{}],
+        embeddings=[[0.0, 1.0]],
+    )
+    first = reader.query(query_embeddings=[[1.0, 0.0]], n_results=1)
+    assert first.ids[0] == ["old"]
+
+    writer_backend = SQLiteExactBackend()
+    writer = writer_backend.get_collection(
+        palace=PalaceRef(id=str(tmp_path), local_path=str(tmp_path)),
+        collection_name="mempalace_drawers",
+        create=False,
+    )
+    writer.add(
+        ids=["new"],
+        documents=["new"],
+        metadatas=[{}],
+        embeddings=[[1.0, 0.0]],
+    )
+
+    second = reader.query(query_embeddings=[[1.0, 0.0]], n_results=2)
+    assert second.ids[0] == ["new", "old"]
+
+    writer_backend.close()
+    reader_backend.close()
+
+
 def test_sqlite_exact_wing_room_counts(tmp_path):
     from mempalace.backends.sqlite_exact import sqlite_wing_room_counts