feat: optimize metadata counting using Qdrant server-side facets

mempalace/backends/base.py

@@ -501,6 +501,15 @@ def get_all_metadata(self, where: Optional[dict] = None) -> list[dict]:
             offset += len(batch_meta)
         return all_meta
 
+    def facet_counts(
+        self,
+        field: str,
+        where: Optional[dict] = None,
+        limit: int = 1000,
+    ) -> dict[str, int]:
+        """Return counts for each distinct value of a metadata field."""
+        raise UnsupportedCapabilityError("backend does not support facet_counts")
+
     def maintenance_state(self) -> dict:
         """Return a structured snapshot of this collection's maintenance state.
 

mempalace/backends/qdrant.py

@@ -40,6 +40,7 @@
     PalaceNotFoundError,
     PalaceRef,
     QueryResult,
+    UnsupportedCapabilityError,
     UnsupportedFilterError,
     _IncludeSpec,
 )
@@ -543,6 +544,38 @@ def count_points(self, collection: str) -> int:
         result = response.get("result") or {}
         return int(result.get("count") or 0)
 
+    def facet_counts(
+        self,
+        collection: str,
+        *,
+        field: str,
+        qdrant_filter: Optional[dict] = None,
+        limit: int = 1000,
+    ) -> dict[str, int]:
+        body: dict[str, Any] = {
+            "key": field,
+            "exact": True,
+            "limit": limit,
+        }
+
+        if qdrant_filter:
+            body["filter"] = qdrant_filter
+
+        response = self.request(
+            "POST",
+            f"/collections/{urlparse.quote(collection, safe='')}/facet",
+            body=body,
+        )
+
+        result = response.get("result") or {}
+        hits = result.get("hits") or []
+
+        return {
+            str(hit["value"]): int(hit.get("count") or 0)
+            for hit in hits
+            if hit.get("value") is not None
+        }
+
     def delete_collection(self, collection: str) -> None:
         self.request("DELETE", f"/collections/{urlparse.quote(collection, safe='')}")
 
@@ -1035,6 +1068,34 @@ def get_all_metadata(self, where: Optional[dict] = None) -> list[dict]:
         rows = self._rows(where=where)
         return [row["metadata"] for row in rows]
 
+    def facet_counts(
+        self,
+        field: str,
+        where: Optional[dict] = None,
+        limit: int = 1000,
+    ) -> dict[str, int]:
+        self._ensure_open()
+        # Validate the filter before the existence short-circuit so an
+        # unsupported local-only filter raises regardless of whether the
+        # collection has been materialized yet — matching the order used by
+        # get()/lexical_search() above (#1835 review).
+        _validate_where(where)
+        if _requires_local_filter(where):
+            raise UnsupportedCapabilityError("facet_counts does not support local-only filters")
+        if not self._remote_exists():
+            if self._marker_exists():
+                raise CollectionNotInitializedError(self._collection_name)
+            return {}
+
+        q_filter = _qdrant_filter(where)
+
+        return self._client.facet_counts(
+            self._remote_collection,
+            field=f"{_PAYLOAD_METADATA}.{field}",
+            qdrant_filter=q_filter,
+            limit=limit,
+        )
+
     def delete(self, *, ids=None, where=None):
         _validate_where(where)
         if not self._remote_exists():
@@ -1125,6 +1186,7 @@ class QdrantBackend(BaseBackend):
             "supports_embeddings_out",
             "supports_metadata_filters",
             "supports_lexical_search",
+            "supports_metadata_facets",
             "supports_namespace_isolation",
             "server_mode",
         }

mempalace/mcp_server.py

@@ -1205,6 +1205,15 @@ def _fetch_all_metadata(col, where=None):
     return all_meta
 
 
+def _supports_metadata_facets(col) -> bool:
+    """Return True if the collection's backend implements metadata facets."""
+    backend = getattr(col, "_backend", None)
+    if backend is None:
+        return False
+    capabilities = getattr(backend, "capabilities", None)
+    return isinstance(capabilities, (set, frozenset)) and "supports_metadata_facets" in capabilities
+
+
 _metadata_cache = None
 _metadata_cache_time = 0
 _METADATA_CACHE_TTL = 5.0  # seconds
@@ -1641,13 +1650,47 @@ def tool_status():
         "backend": _selected_backend_name(),
     }
     try:
-        all_meta = _get_cached_metadata(col)
-        for m in all_meta:
-            m = m or {}
-            w = m.get("wing", "unknown")
-            r = m.get("room", "unknown")
-            wings[w] = wings.get(w, 0) + 1
-            rooms[r] = rooms.get(r, 0) + 1
+        if _supports_metadata_facets(col):
+            try:
+                temp_wings = col.facet_counts("wing")
+                wings.update(temp_wings)
+                try:
+                    unknown_wings = count - sum(temp_wings.values())
+                    if unknown_wings > 0:
+                        wings["unknown"] = wings.get("unknown", 0) + unknown_wings
+                except (TypeError, ValueError):
+                    pass
+
+                temp_rooms = col.facet_counts("room")
+                rooms.update(temp_rooms)
+                try:
+                    unknown_rooms = count - sum(temp_rooms.values())
+                    if unknown_rooms > 0:
+                        rooms["unknown"] = rooms.get("unknown", 0) + unknown_rooms
+                except (TypeError, ValueError):
+                    pass
+
+            except Exception as e:
+                logger.warning(
+                    "Failed to fetch metadata facets, falling back to client-side loop: %s", e
+                )
+                rooms.clear()
+                wings.clear()
+                all_meta = _get_cached_metadata(col)
+                for m in all_meta:
+                    m = m or {}
+                    w = m.get("wing", "unknown")
+                    r = m.get("room", "unknown")
+                    wings[w] = wings.get(w, 0) + 1
+                    rooms[r] = rooms.get(r, 0) + 1
+        else:
+            all_meta = _get_cached_metadata(col)
+            for m in all_meta:
+                m = m or {}
+                w = m.get("wing", "unknown")
+                r = m.get("room", "unknown")
+                wings[w] = wings.get(w, 0) + 1
+                rooms[r] = rooms.get(r, 0) + 1
     except Exception as e:
         logger.exception("tool_status metadata fetch failed")
         result["error"] = str(e)
@@ -1702,11 +1745,28 @@ def tool_list_wings():
     wings = {}
     result = {"wings": wings}
     try:
-        all_meta = _get_cached_metadata(col)
-        for m in all_meta:
-            m = m or {}
-            w = m.get("wing", "unknown")
-            wings[w] = wings.get(w, 0) + 1
+        try:
+            if not _supports_metadata_facets(col):
+                raise ValueError("facets not supported")
+            temp_wings = col.facet_counts("wing")
+            wings.update(temp_wings)
+            try:
+                unknown_wings = col.count() - sum(temp_wings.values())
+                if unknown_wings > 0:
+                    wings["unknown"] = wings.get("unknown", 0) + unknown_wings
+            except (TypeError, ValueError):
+                pass
+        except Exception as e:
+            if _supports_metadata_facets(col):
+                logger.warning(
+                    "Failed to fetch metadata facets, falling back to client-side loop: %s", e
+                )
+            wings.clear()
+            all_meta = _get_cached_metadata(col)
+            for m in all_meta:
+                m = m or {}
+                w = m.get("wing", "unknown")
+                wings[w] = wings.get(w, 0) + 1
     except Exception as e:
         logger.exception("tool_list_wings metadata fetch failed")
         result["error"] = str(e)
@@ -1734,13 +1794,34 @@ def tool_list_rooms(wing: str = None):
         return _collection_error_or_no_palace()
     rooms = {}
     result = {"wing": wing or "all", "rooms": rooms}
+    where = {"wing": wing} if wing else None
     try:
-        where = {"wing": wing} if wing else None
-        all_meta = _fetch_all_metadata(col, where=where)
-        for m in all_meta:
-            m = m or {}
-            r = m.get("room", "unknown")
-            rooms[r] = rooms.get(r, 0) + 1
+        try:
+            if not _supports_metadata_facets(col):
+                raise ValueError("facets not supported")
+            temp_rooms = col.facet_counts("room", where=where)
+            rooms.update(temp_rooms)
+            try:
+                if wing:
+                    wing_count = col.facet_counts("wing", where={"wing": wing}).get(wing, 0)
+                    unknown_rooms = wing_count - sum(temp_rooms.values())
+                else:
+                    unknown_rooms = col.count() - sum(temp_rooms.values())
+                if unknown_rooms > 0:
+                    rooms["unknown"] = rooms.get("unknown", 0) + unknown_rooms
+            except (TypeError, ValueError):
+                pass
+        except Exception as e:
+            if _supports_metadata_facets(col):
+                logger.warning(
+                    "Failed to fetch metadata facets, falling back to client-side loop: %s", e
+                )
+            rooms.clear()
+            all_meta = _fetch_all_metadata(col, where=where)
+            for m in all_meta:
+                m = m or {}
+                r = m.get("room", "unknown")
+                rooms[r] = rooms.get(r, 0) + 1
     except Exception as e:
         logger.exception("tool_list_rooms metadata fetch failed")
         result["error"] = str(e)
@@ -1759,14 +1840,42 @@ def tool_get_taxonomy():
     taxonomy = {}
     result = {"taxonomy": taxonomy}
     try:
-        all_meta = _get_cached_metadata(col)
-        for m in all_meta:
-            m = m or {}
-            w = m.get("wing", "unknown")
-            r = m.get("room", "unknown")
-            if w not in taxonomy:
-                taxonomy[w] = {}
-            taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
+        try:
+            if not _supports_metadata_facets(col):
+                raise ValueError("facets not supported")
+            from concurrent.futures import ThreadPoolExecutor
+
+            wing_counts = col.facet_counts("wing")
+            wings = list(wing_counts.keys())
+            temp_taxonomy = {}
+            with ThreadPoolExecutor(max_workers=max(1, min(8, len(wings)))) as executor:
+                futures = {
+                    wing: executor.submit(col.facet_counts, "room", where={"wing": wing})
+                    for wing in wings
+                }
+                for wing, future in futures.items():
+                    room_counts = future.result()
+                    try:
+                        unknown_rooms = wing_counts[wing] - sum(room_counts.values())
+                        if unknown_rooms > 0:
+                            room_counts["unknown"] = room_counts.get("unknown", 0) + unknown_rooms
+                    except (TypeError, ValueError):
+                        pass
+                    temp_taxonomy[wing] = room_counts
+                taxonomy.update(temp_taxonomy)
+        except Exception as e:
+            if _supports_metadata_facets(col):
+                logger.warning(
+                    "Failed to fetch metadata facets, falling back to client-side loop: %s", e
+                )
+            all_meta = _get_cached_metadata(col)
+            for m in all_meta:
+                m = m or {}
+                w = m.get("wing", "unknown")
+                r = m.get("room", "unknown")
+                if w not in taxonomy:
+                    taxonomy[w] = {}
+                taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
     except Exception as e:
         logger.exception("tool_get_taxonomy metadata fetch failed")
         result["error"] = str(e)