updating count_unique_metadata_by_filter to receive multiple metadata fields instead of only one

Author: davidsbatista

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py

@@ -1047,21 +1047,22 @@ async def get_metadata_field_min_max_async(self, metadata_field: str) -> dict[st
         except Exception:
             return {}
 
-    def count_unique_metadata_by_filter(self, metadata_field: str, filters: dict[str, Any] | None = None) -> int:
+    def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]:
         """
-        Returns the number of unique values for the given metadata field among documents that match the filters.
+        Returns the number of unique values for each specified metadata field among documents that match the filters.
 
-        :param metadata_field: The metadata field key (inside ``meta``) to count unique values for.
-        :param filters: Optional filters to restrict the documents considered.
+        :param filters: The filters to restrict the documents considered.
             For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
+        :param metadata_fields: List of metadata field keys (inside ``meta``) to count unique values for.
 
-        :returns: The number of unique values for the field among the filtered documents.
+        :returns: A dictionary mapping each metadata field name to the count of its unique values among the filtered
+                  documents.
         """
         self._initialize_client()
         assert self._client is not None
 
         qdrant_filter = convert_filters_to_qdrant(filters) if filters else None
-        unique_values = set()
+        unique_values_by_field: dict[str, set[Any]] = {field: set() for field in metadata_fields}
 
         try:
             next_offset = None
@@ -1087,37 +1088,38 @@ def count_unique_metadata_by_filter(self, metadata_field: str, filters: dict[str
                 for record in records:
                     if record.payload and "meta" in record.payload:
                         meta = record.payload["meta"]
-                        if metadata_field in meta:
-                            value = meta[metadata_field]
-                            if value is not None:
-                                # Convert to hashable type if needed
-                                if isinstance(value, (list, dict)):
-                                    unique_values.add(str(value))
-                                else:
-                                    unique_values.add(value)
-
-            return len(unique_values)
+                        for field in metadata_fields:
+                            if field in meta:
+                                value = meta[field]
+                                if value is not None:
+                                    if isinstance(value, (list, dict)):
+                                        unique_values_by_field[field].add(str(value))
+                                    else:
+                                        unique_values_by_field[field].add(value)
+
+            return {field: len(unique_values_by_field[field]) for field in metadata_fields}
         except Exception:
-            return 0
+            return dict.fromkeys(metadata_fields, 0)
 
     async def count_unique_metadata_by_filter_async(
-        self, metadata_field: str, filters: dict[str, Any] | None = None
-    ) -> int:
+        self, filters: dict[str, Any], metadata_fields: list[str]
+    ) -> dict[str, int]:
         """
-        Asynchronously returns the number of unique values for the given metadata field among documents that match
-        the filters.
+        Asynchronously returns the number of unique values for each specified metadata field among documents that
+        match the filters.
 
-        :param metadata_field: The metadata field key (inside ``meta``) to count unique values for.
-        :param filters: Optional filters to restrict the documents considered.
+        :param filters: The filters to restrict the documents considered.
             For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
+        :param metadata_fields: List of metadata field keys (inside ``meta``) to count unique values for.
 
-        :returns: The number of unique values for the field among the filtered documents.
+        :returns: A dictionary mapping each metadata field name to the count of its unique values among the filtered
+                  documents.
         """
         await self._initialize_async_client()
         assert self._async_client is not None
 
         qdrant_filter = convert_filters_to_qdrant(filters) if filters else None
-        unique_values = set()
+        unique_values_by_field: dict[str, set[Any]] = {field: set() for field in metadata_fields}
 
         try:
             next_offset = None
@@ -1143,18 +1145,18 @@ async def count_unique_metadata_by_filter_async(
                 for record in records:
                     if record.payload and "meta" in record.payload:
                         meta = record.payload["meta"]
-                        if metadata_field in meta:
-                            value = meta[metadata_field]
-                            if value is not None:
-                                # Convert to hashable type if needed
-                                if isinstance(value, (list, dict)):
-                                    unique_values.add(str(value))
-                                else:
-                                    unique_values.add(value)
-
-            return len(unique_values)
+                        for field in metadata_fields:
+                            if field in meta:
+                                value = meta[field]
+                                if value is not None:
+                                    if isinstance(value, (list, dict)):
+                                        unique_values_by_field[field].add(str(value))
+                                    else:
+                                        unique_values_by_field[field].add(value)
+
+            return {field: len(unique_values_by_field[field]) for field in metadata_fields}
         except Exception:
-            return 0
+            return dict.fromkeys(metadata_fields, 0)
 
     def get_metadata_field_unique_values(
         self, metadata_field: str, filters: dict[str, Any] | None = None, limit: int = 100, offset: int = 0

integrations/qdrant/tests/test_document_store.py

@@ -595,8 +595,20 @@ def test_count_unique_metadata_by_filter(self, document_store: QdrantDocumentSto
         ]
         document_store.write_documents(docs)
 
-        count = document_store.count_unique_metadata_by_filter("category")
-        assert count == 3
+        result = document_store.count_unique_metadata_by_filter(filters={}, metadata_fields=["category"])
+        assert result == {"category": 3}
+
+    def test_count_unique_metadata_by_filter_multiple_fields(self, document_store: QdrantDocumentStore):
+        """Test counting unique values for multiple metadata fields."""
+        docs = [
+            Document(content="Doc 1", meta={"category": "A", "status": "active"}),
+            Document(content="Doc 2", meta={"category": "B", "status": "active"}),
+            Document(content="Doc 3", meta={"category": "A", "status": "inactive"}),
+        ]
+        document_store.write_documents(docs)
+
+        result = document_store.count_unique_metadata_by_filter(filters={}, metadata_fields=["category", "status"])
+        assert result == {"category": 2, "status": 2}
 
     def test_count_unique_metadata_by_filter_with_filter(self, document_store: QdrantDocumentStore):
         """Test counting unique metadata field values with filtering."""
@@ -607,11 +619,11 @@ def test_count_unique_metadata_by_filter_with_filter(self, document_store: Qdran
         ]
         document_store.write_documents(docs)
 
-        # Count unique categories where status is active
-        count = document_store.count_unique_metadata_by_filter(
-            "category", filters={"field": "meta.status", "operator": "==", "value": "active"}
+        result = document_store.count_unique_metadata_by_filter(
+            filters={"field": "meta.status", "operator": "==", "value": "active"},
+            metadata_fields=["category"],
         )
-        assert count == 2
+        assert result == {"category": 2}
 
     def test_get_metadata_field_unique_values(self, document_store: QdrantDocumentStore):
         """Test getting unique metadata field values."""

integrations/qdrant/tests/test_document_store_async.py

@@ -547,8 +547,23 @@ async def test_count_unique_metadata_by_filter_async(self, document_store: Qdran
         ]
         await document_store.write_documents_async(docs)
 
-        count = await document_store.count_unique_metadata_by_filter_async("category")
-        assert count == 3
+        result = await document_store.count_unique_metadata_by_filter_async(filters={}, metadata_fields=["category"])
+        assert result == {"category": 3}
+
+    @pytest.mark.asyncio
+    async def test_count_unique_metadata_by_filter_async_multiple_fields(self, document_store: QdrantDocumentStore):
+        """Test counting unique values for multiple metadata fields (async)."""
+        docs = [
+            Document(content="Doc 1", meta={"category": "A", "status": "active"}),
+            Document(content="Doc 2", meta={"category": "B", "status": "active"}),
+            Document(content="Doc 3", meta={"category": "A", "status": "inactive"}),
+        ]
+        await document_store.write_documents_async(docs)
+
+        result = await document_store.count_unique_metadata_by_filter_async(
+            filters={}, metadata_fields=["category", "status"]
+        )
+        assert result == {"category": 2, "status": 2}
 
     @pytest.mark.asyncio
     async def test_count_unique_metadata_by_filter_async_with_filter(self, document_store: QdrantDocumentStore):
@@ -560,11 +575,11 @@ async def test_count_unique_metadata_by_filter_async_with_filter(self, document_
         ]
         await document_store.write_documents_async(docs)
 
-        # Count unique categories where status is active
-        count = await document_store.count_unique_metadata_by_filter_async(
-            "category", filters={"field": "meta.status", "operator": "==", "value": "active"}
+        result = await document_store.count_unique_metadata_by_filter_async(
+            filters={"field": "meta.status", "operator": "==", "value": "active"},
+            metadata_fields=["category"],
         )
-        assert count == 2
+        assert result == {"category": 2}
 
     @pytest.mark.asyncio
     async def test_get_metadata_field_unique_values_async(self, document_store: QdrantDocumentStore):