POC: bespoke PalaceStore storage layer (drop-in ChromaDB replacement)

.github/workflows/ci.yml

@@ -2,9 +2,9 @@ name: Tests
 
 on:
   push:
-    branches: [main]
+    branches: [main, develop]
   pull_request:
-    branches: [main]
+    branches: [main, develop]
 
 jobs:
   test-linux:

benchmarks/longmemeval_bench.py

@@ -39,11 +39,22 @@
 from collections import defaultdict
 from datetime import datetime
 
-import chromadb
-
-# Add mempal to path
+# Add mempal to path so palace_store.compat is importable when we
+# opt into it via MEMPAL_STORAGE below.
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
+# Storage backend selector. Default (no env var) uses real chromadb.
+# Setting MEMPAL_STORAGE=palace_store (or any of its aliases) swaps in
+# palace_store.compat, which is API-compatible with chromadb's
+# PersistentClient/EphemeralClient for the narrow surface this
+# benchmark uses. Inlined here rather than imported from mempalace so
+# benchmarks/ stays free of mempalace's internal import dance.
+_BACKEND_NAME = os.environ.get("MEMPAL_STORAGE", "").strip().lower()
+if _BACKEND_NAME in ("palace", "palace_store", "palacestore"):
+    from palace_store import compat as chromadb  # noqa: E402,F401
+else:
+    import chromadb  # noqa: E402
+
 
 # =============================================================================
 # METRICS (reimplemented to avoid LongMemEval dependency)

benchmarks/storage/__init__.py

@@ -0,0 +1,9 @@
+"""Standalone benchmark harness for storage layer candidates.
+
+This harness does not depend on mempalace internals. It measures pure
+storage performance — ingest throughput, query latency, memory usage,
+disk footprint — by exercising a `StoreAdapter` interface that any
+candidate (PalaceStore, ChromaDB, ...) can implement.
+
+See benchmarks/storage/README.md for usage.
+"""

benchmarks/storage/adapters/__init__.py

@@ -0,0 +1 @@
+"""Store adapters. Each module exposes a class implementing StoreAdapter."""

benchmarks/storage/adapters/chroma.py

@@ -0,0 +1,188 @@
+"""ChromaDB baseline adapter.
+
+We bypass Chroma's internal embedder entirely by passing ``embeddings=`` on
+every upsert and ``query_embeddings=`` on every query. A no-op
+EmbeddingFunction is installed so any accidental ``documents=``-only call
+fails loudly instead of silently triggering sentence-transformers.
+
+Cosine space is explicitly requested via collection metadata. Any other
+HNSW tuning is left at Chroma's defaults — we're measuring what mempalace
+ships with today, not a hand-tuned Chroma.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from ..interface import QueryHit
+
+
+class _NoEmbed:
+    """Stub embedder that refuses to run. Forces pre-embedded call sites."""
+
+    def __call__(self, input):  # noqa: A002 — chroma uses this name
+        raise RuntimeError(
+            "ChromaAdapter must be called with embeddings=/query_embeddings=; "
+            "the store's embedder should never run during this benchmark"
+        )
+
+    def name(self) -> str:
+        return "noembed"
+
+
+class ChromaAdapter:
+    name = "chroma"
+    # HNSW is approximate, but tuned search_ef=200 makes it effectively
+    # exact at the scales the gate runs. Mark as exact to apply the
+    # strict threshold.
+    is_exact = True
+
+    def __init__(self, path: str | Path):
+        import chromadb
+        from chromadb.config import Settings
+
+        self._path = Path(path)
+        self._path.mkdir(parents=True, exist_ok=True)
+        self._client = chromadb.PersistentClient(
+            path=str(self._path),
+            settings=Settings(anonymized_telemetry=False),
+        )
+        # HNSW tuned for accuracy, not speed, to make the correctness gate
+        # meaningful. search_ef at the default of 10 is essentially equal
+        # to k=10 which gives HNSW no room to explore alternatives, so the
+        # adapter appears to "miss" tail results that are really just
+        # approximation error. search_ef=200 puts Chroma near-exact at the
+        # scales we test. batch_size + sync_threshold follow the screenshot
+        # PR's guidance to avoid the ingest bloat pathology at 50k+.
+        self._col = self._client.get_or_create_collection(
+            name="bench_drawers",
+            embedding_function=_NoEmbed(),
+            metadata={
+                "hnsw:space": "cosine",
+                "hnsw:search_ef": 200,
+                "hnsw:construction_ef": 200,
+                "hnsw:M": 16,
+                "hnsw:batch_size": 10000,
+                "hnsw:sync_threshold": 50000,
+            },
+        )
+
+    def upsert(
+        self,
+        ids: list[str],
+        vectors: np.ndarray,
+        metadatas: list[dict[str, Any]],
+        texts: list[str],
+    ) -> None:
+        # Chroma wants a list-of-lists for embeddings. .tolist() allocates,
+        # but so does any other marshaling path into their bindings.
+        self._col.upsert(
+            ids=ids,
+            embeddings=vectors.tolist(),
+            metadatas=metadatas,
+            documents=texts,
+        )
+
+    def query(
+        self,
+        query_vector: np.ndarray,
+        k: int,
+        where: dict[str, Any] | None = None,
+    ) -> list[QueryHit]:
+        res = self._col.query(
+            query_embeddings=[query_vector.tolist()],
+            n_results=k,
+            where=self._translate_where(where),
+            include=["metadatas", "distances"],
+        )
+        ids = res["ids"][0]
+        metas = res["metadatas"][0]
+        dists = res["distances"][0]
+        hits: list[QueryHit] = []
+        for i, (rid, meta, dist) in enumerate(zip(ids, metas, dists)):
+            # Cosine distance → similarity. Chroma returns 1 - cos for cosine
+            # space; for unit vectors this is 1 - dot, so score = 1 - dist.
+            hits.append(
+                QueryHit(
+                    id=rid,
+                    score=float(1.0 - dist),
+                    wing=meta.get("wing", ""),
+                    room=meta.get("room", ""),
+                )
+            )
+        return hits
+
+    def get(self, where: dict[str, Any]) -> list[dict[str, Any]]:
+        res = self._col.get(
+            where=self._translate_where(where),
+            include=["metadatas", "documents"],
+        )
+        out: list[dict[str, Any]] = []
+        for rid, meta, doc in zip(res["ids"], res["metadatas"], res["documents"]):
+            out.append(
+                {
+                    "id": rid,
+                    "wing": meta.get("wing", ""),
+                    "room": meta.get("room", ""),
+                    "source_file": meta.get("source_file"),
+                    "chunk_index": meta.get("chunk_index"),
+                    "text": doc,
+                    "metadata": {
+                        k: v
+                        for k, v in meta.items()
+                        if k not in ("wing", "room", "source_file", "chunk_index")
+                    },
+                }
+            )
+        return out
+
+    def delete(self, where: dict[str, Any]) -> int:
+        # Chroma's delete() returns None, so we count first.
+        hits = self._col.get(where=self._translate_where(where), include=[])
+        n = len(hits["ids"])
+        self._col.delete(where=self._translate_where(where))
+        return n
+
+    def count(self) -> int:
+        return self._col.count()
+
+    def disk_bytes(self) -> int:
+        total = 0
+        for p in self._path.rglob("*"):
+            if p.is_file():
+                try:
+                    total += p.stat().st_size
+                except FileNotFoundError:
+                    pass
+        return total
+
+    def warm(self, *, mlock: bool = False) -> None:
+        # Chroma manages its own in-memory HNSW graph after open — there
+        # is no mmap'd region we can touch from Python without digging
+        # into chromadb internals. Treat as no-op.
+        return
+
+    def close(self) -> None:
+        # chromadb PersistentClient doesn't expose a close() we can trust,
+        # so we just drop our references.
+        self._col = None
+        self._client = None
+
+    # ── helpers ───────────────────────────────────────────────────────
+
+    @staticmethod
+    def _translate_where(where: dict[str, Any] | None) -> dict[str, Any] | None:
+        """Translate the adapter's flat where dict to Chroma's $and format.
+
+        ``{wing: X, room: Y}`` must become ``{"$and": [{wing: X}, {room: Y}]}``
+        because Chroma rejects multi-key flat dicts.
+        """
+        if not where:
+            return None
+        keys = list(where.keys())
+        if len(keys) == 1:
+            return {keys[0]: where[keys[0]]}
+        return {"$and": [{k: where[k]} for k in keys]}

benchmarks/storage/adapters/palace.py

@@ -0,0 +1,58 @@
+"""PalaceStore adapter."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from palace_store import PalaceStore
+
+from ..interface import QueryHit
+
+
+class PalaceAdapter:
+    name = "palace"
+    is_exact = True  # flat brute-force f32 cosine
+
+    def __init__(self, path: str | Path, *, parallel_query: bool = False):
+        self._store = PalaceStore(path, parallel_query=parallel_query)
+
+    def upsert(
+        self,
+        ids: list[str],
+        vectors: np.ndarray,
+        metadatas: list[dict[str, Any]],
+        texts: list[str],
+    ) -> None:
+        self._store.upsert(ids, vectors, metadatas, texts)
+
+    def query(
+        self,
+        query_vector: np.ndarray,
+        k: int,
+        where: dict[str, Any] | None = None,
+    ) -> list[QueryHit]:
+        rows = self._store.query(query_vector, k, where=where)
+        return [
+            QueryHit(id=r.id, score=r.score, wing=r.wing, room=r.room) for r in rows
+        ]
+
+    def get(self, where: dict[str, Any]) -> list[dict[str, Any]]:
+        return self._store.get(where)
+
+    def delete(self, where: dict[str, Any]) -> int:
+        return self._store.delete(where)
+
+    def count(self) -> int:
+        return self._store.count()
+
+    def disk_bytes(self) -> int:
+        return self._store.disk_bytes()["total"]
+
+    def warm(self, *, mlock: bool = False) -> None:
+        self._store.warm_pages(mlock=mlock)
+
+    def close(self) -> None:
+        self._store.close()

benchmarks/storage/adapters/palace_i8.py

@@ -0,0 +1,65 @@
+"""PalaceStore int8 adapter.
+
+Same storage engine as the float32 PalaceAdapter, but each shard stores
+per-row quantized int8 vectors plus a float32 per-row scale. On-disk and
+in-RAM size drops by ~4x; query latency rises ~5x because numpy has no
+BLAS int8 matmul path. See palace_store/store.py::VectorShardI8 for the
+quantization math and the expected tradeoff.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from palace_store import PalaceStore
+
+from ..interface import QueryHit
+
+
+class PalaceI8Adapter:
+    name = "palace_i8"
+    is_exact = False  # int8 quantization introduces small numeric error
+
+    def __init__(self, path: str | Path):
+        self._store = PalaceStore(path, dtype="int8")
+
+    def upsert(
+        self,
+        ids: list[str],
+        vectors: np.ndarray,
+        metadatas: list[dict[str, Any]],
+        texts: list[str],
+    ) -> None:
+        self._store.upsert(ids, vectors, metadatas, texts)
+
+    def query(
+        self,
+        query_vector: np.ndarray,
+        k: int,
+        where: dict[str, Any] | None = None,
+    ) -> list[QueryHit]:
+        rows = self._store.query(query_vector, k, where=where)
+        return [
+            QueryHit(id=r.id, score=r.score, wing=r.wing, room=r.room) for r in rows
+        ]
+
+    def get(self, where: dict[str, Any]) -> list[dict[str, Any]]:
+        return self._store.get(where)
+
+    def delete(self, where: dict[str, Any]) -> int:
+        return self._store.delete(where)
+
+    def count(self) -> int:
+        return self._store.count()
+
+    def disk_bytes(self) -> int:
+        return self._store.disk_bytes()["total"]
+
+    def warm(self, *, mlock: bool = False) -> None:
+        self._store.warm_pages(mlock=mlock)
+
+    def close(self) -> None:
+        self._store.close()