feat: Add Astra document store operations (#2904)

* Add Astra document store operations

* calling static methods from the class and not the instance

---------

Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: Keyur-S-Patel

integrations/astra/src/haystack_integrations/document_stores/astra/astra_client.py

@@ -7,6 +7,7 @@
 from warnings import warn
 
 from astrapy import DataAPIClient as AstraDBClient
+from astrapy.collection import FilterType
 from astrapy.constants import ReturnDocument
 from astrapy.exceptions import CollectionAlreadyExistsException
 from haystack import logging
@@ -216,19 +217,22 @@ def _query(self, vector, top_k, filters=None):
 
         return result
 
-    def find_documents(self, find_query):
+    def find_documents(
+        self, find_query: dict[str, Any], projection: dict[str, Any] | None = None
+    ) -> list[dict[str, Any]]:
         """
         Find documents in the Astra index.
 
         :param find_query: a dictionary with the query options
+        :param projection: optional projection for the returned documents
         :returns: the documents found in the index
         """
         find_cursor = self._astra_db_collection.find(
             filter=find_query.get("filter"),
             sort=find_query.get("sort"),
             limit=find_query.get("limit"),
             include_similarity=find_query.get("includeSimilarity"),
-            projection={"*": 1},
+            projection=projection or {"*": 1},
         )
 
         find_results = []
@@ -354,12 +358,25 @@ def delete_all_documents(self) -> int:
 
         return delete_result.deleted_count
 
-    def count_documents(self, upper_bound: int = 10000) -> int:
+    def count_documents(self, filters: FilterType | None = None, upper_bound: int = 10000) -> int:
         """
         Count the number of documents in the Astra index.
+
+        :param filters: optional filter to restrict the counted documents
+        :param upper_bound: maximum expected count, required by Astra's API
         :returns: the number of documents in the index
         """
-        return self._astra_db_collection.count_documents({}, upper_bound=upper_bound)
+        return self._astra_db_collection.count_documents(filters or {}, upper_bound=upper_bound)
+
+    def distinct(self, key: str, filters: FilterType | None = None) -> list[Any]:
+        """
+        Return the distinct values for a field in the Astra index.
+
+        :param key: field name
+        :param filters: optional filter to restrict the matching documents
+        :returns: distinct values for the field
+        """
+        return self._astra_db_collection.distinct(key, filter=filters)
 
     def update(
         self,

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

@@ -284,6 +284,59 @@ def count_documents(self) -> int:
         """
         return self.index.count_documents()
 
+    @staticmethod
+    def _normalize_new_filter_input(filters: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(filters, dict):
+            msg = "Filters must be a dictionary"
+            raise AstraDocumentStoreFilterError(msg)
+
+        normalized_filters = filters.copy()
+        if "id" in normalized_filters:
+            normalized_filters["_id"] = normalized_filters.pop("id")
+
+        return normalized_filters
+
+    @staticmethod
+    def _infer_metadata_field_type(values: list[Any]) -> str:
+        inferred_types = set()
+        for value in values:
+            if isinstance(value, list):
+                for item in value:
+                    if isinstance(item, bool):
+                        inferred_types.add("boolean")
+                    elif isinstance(item, (int, float)):
+                        inferred_types.add("long")
+                    elif isinstance(item, str):
+                        inferred_types.add("keyword")
+            elif isinstance(value, bool):
+                inferred_types.add("boolean")
+            elif isinstance(value, (int, float)):
+                inferred_types.add("long")
+            elif isinstance(value, str):
+                inferred_types.add("keyword")
+
+        if not inferred_types:
+            return "keyword"
+
+        if len(inferred_types) > 1:
+            logger.warning("Field has mixed metadata types {types}. Defaulting to 'keyword'.", types=inferred_types)
+            return "keyword"
+
+        return next(iter(inferred_types))
+
+    @staticmethod
+    def _normalize_distinct_values(values: list[Any]) -> list[str]:
+        normalized_values: set[str] = set()
+        for value in values:
+            if isinstance(value, list):
+                normalized_values.update(str(item) for item in value)
+            elif value is not None:
+                normalized_values.add(str(value))
+        return sorted(normalized_values)
+
+    def _get_metadata_projection_documents(self) -> list[dict[str, Any]]:
+        return self.index.find_documents({}, projection={"content": 1, "meta": 1})
+
     def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
         """
         Returns at most 1000 documents that match the filter.
@@ -484,3 +537,92 @@ def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int
         logger.info(f"{update_count} documents updated by filter")
 
         return update_count
+
+    def count_documents_by_filter(self, filters: dict[str, Any]) -> int:
+        """
+        Applies a filter and counts the documents that matched it.
+
+        :param filters: The filters to apply to the document list.
+        :returns: The number of documents that match the filter.
+        """
+        normalized_filters = AstraDocumentStore._normalize_new_filter_input(filters)
+        converted_filters = _convert_filters(normalized_filters)
+        return self.index.count_documents(filters=converted_filters, upper_bound=1_000_000_000)
+
+    def count_unique_metadata_by_filter(self, filters: dict[str, Any], metadata_fields: list[str]) -> dict[str, int]:
+        """
+        Applies a filter selecting documents and counts the unique values for each meta field of the matched
+        documents.
+
+        :param filters: The filters to apply to the document list.
+        :param metadata_fields: The metadata fields to count unique values for.
+        :returns: A dictionary where the keys are the metadata field names and the values are the count of unique
+            values.
+        """
+        normalized_filters = AstraDocumentStore._normalize_new_filter_input(filters)
+        converted_filters = _convert_filters(normalized_filters)
+
+        counts = {}
+        for field in metadata_fields:
+            distinct_values = self.index.distinct(f"meta.{field}", filters=converted_filters)
+            counts[field] = len(AstraDocumentStore._normalize_distinct_values(distinct_values))
+        return counts
+
+    def get_metadata_fields_info(self) -> dict[str, dict[str, str]]:
+        """
+        Returns the metadata fields and the corresponding types.
+
+        :returns: A dictionary mapping field names to dictionaries with a `type` key.
+        """
+        documents = self._get_metadata_projection_documents()
+        if not documents:
+            return {}
+
+        fields_info: dict[str, dict[str, str]] = {}
+
+        if any(document.get("content") is not None for document in documents):
+            fields_info["content"] = {"type": "text"}
+
+        field_values: dict[str, list[Any]] = {}
+        for document in documents:
+            for field, value in document.get("meta", {}).items():
+                field_values.setdefault(field, []).append(value)
+
+        for field, values in field_values.items():
+            fields_info[field] = {"type": self._infer_metadata_field_type(values)}
+
+        return fields_info
+
+    def get_metadata_field_min_max(self, metadata_field: str) -> dict[str, Any]:
+        """
+        For a given metadata field, find its max and min value.
+
+        :param metadata_field: The metadata field to inspect.
+        :returns: A dictionary with `min` and `max`.
+        """
+        distinct_values = self.index.distinct(f"meta.{metadata_field}")
+        comparable_values = [value for value in distinct_values if isinstance(value, (str, int, float, bool))]
+        if not comparable_values:
+            return {"min": None, "max": None}
+
+        return {"min": min(comparable_values), "max": max(comparable_values)}
+
+    def get_metadata_field_unique_values(
+        self, metadata_field: str, search_term: str | None = None, from_: int = 0, size: int = 10
+    ) -> tuple[list[str], int]:
+        """
+        Retrieves unique values for a field matching a search term or all possible values if no search term is given.
+
+        :param metadata_field: The metadata field to inspect.
+        :param search_term: Optional case-insensitive substring search term.
+        :param from_: The starting index for pagination.
+        :param size: The number of values to return.
+        :returns: A tuple containing the paginated values and the total count.
+        """
+        values = AstraDocumentStore._normalize_distinct_values(self.index.distinct(f"meta.{metadata_field}"))
+        if search_term:
+            search_term_lower = search_term.lower()
+            values = [value for value in values if search_term_lower in value.lower()]
+
+        total_count = len(values)
+        return values[from_ : from_ + size], total_count

integrations/astra/tests/test_document_store.py

@@ -43,6 +43,93 @@ def test_to_dict(mock_auth):  # noqa
         }
 
 
+@pytest.mark.usefixtures("mock_auth")
+@mock.patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+def test_count_documents_by_filter(mock_astra_client):
+    mock_index = mock_astra_client.return_value
+    mock_index.count_documents.return_value = 2
+
+    store = AstraDocumentStore()
+
+    count = store.count_documents_by_filter({"field": "meta.status", "operator": "==", "value": "draft"})
+
+    assert count == 2
+    mock_index.count_documents.assert_called_once_with(
+        filters={"meta.status": {"$eq": "draft"}}, upper_bound=1_000_000_000
+    )
+
+
+@pytest.mark.usefixtures("mock_auth")
+@mock.patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+def test_count_unique_metadata_by_filter(mock_astra_client):
+    mock_index = mock_astra_client.return_value
+    mock_index.distinct.side_effect = [["news", "docs", ["docs", "faq"], None], [1, 2, 2]]
+
+    store = AstraDocumentStore()
+
+    counts = store.count_unique_metadata_by_filter(
+        {"field": "meta.status", "operator": "==", "value": "published"}, ["category", "priority"]
+    )
+
+    assert counts == {"category": 3, "priority": 2}
+    assert mock_index.distinct.call_args_list == [
+        mock.call("meta.category", filters={"meta.status": {"$eq": "published"}}),
+        mock.call("meta.priority", filters={"meta.status": {"$eq": "published"}}),
+    ]
+
+
+@pytest.mark.usefixtures("mock_auth")
+@mock.patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+def test_get_metadata_fields_info(mock_astra_client):
+    mock_index = mock_astra_client.return_value
+    mock_index.find_documents.return_value = [
+        {"content": "Doc 1", "meta": {"category": "news", "priority": 1, "active": True}},
+        {"content": "Doc 2", "meta": {"category": "docs", "priority": 2.5, "tags": ["a", "b"]}},
+    ]
+
+    store = AstraDocumentStore()
+
+    fields_info = store.get_metadata_fields_info()
+
+    assert fields_info == {
+        "content": {"type": "text"},
+        "category": {"type": "keyword"},
+        "priority": {"type": "long"},
+        "active": {"type": "boolean"},
+        "tags": {"type": "keyword"},
+    }
+    mock_index.find_documents.assert_called_once_with({}, projection={"content": 1, "meta": 1})
+
+
+@pytest.mark.usefixtures("mock_auth")
+@mock.patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+def test_get_metadata_field_min_max(mock_astra_client):
+    mock_index = mock_astra_client.return_value
+    mock_index.distinct.return_value = [10, 3, 7]
+
+    store = AstraDocumentStore()
+
+    result = store.get_metadata_field_min_max("priority")
+
+    assert result == {"min": 3, "max": 10}
+    mock_index.distinct.assert_called_once_with("meta.priority")
+
+
+@pytest.mark.usefixtures("mock_auth")
+@mock.patch("haystack_integrations.document_stores.astra.document_store.AstraClient")
+def test_get_metadata_field_unique_values(mock_astra_client):
+    mock_index = mock_astra_client.return_value
+    mock_index.distinct.return_value = ["Beta", "alpha", ["gamma", "alphabet"], None]
+
+    store = AstraDocumentStore()
+
+    values, total_count = store.get_metadata_field_unique_values("category", search_term="alp", from_=0, size=5)
+
+    assert values == ["alpha", "alphabet"]
+    assert total_count == 2
+    mock_index.distinct.assert_called_once_with("meta.category")
+
+
 @pytest.mark.integration
 @pytest.mark.skipif(
     os.environ.get("ASTRA_DB_APPLICATION_TOKEN", "") == "", reason="ASTRA_DB_APPLICATION_TOKEN env var not set"
@@ -204,6 +291,80 @@ def test_filter_documents_by_in_operator(self, document_store):
         self.assert_documents_are_equal([result[0]], [docs[0]])
         self.assert_documents_are_equal([result[1]], [docs[1]])
 
+    def test_count_documents_by_filter(self, document_store: AstraDocumentStore):
+        docs = [
+            Document(id="1", content="Doc 1", meta={"category": "news", "status": "published", "priority": 3}),
+            Document(id="2", content="Doc 2", meta={"category": "docs", "status": "draft", "priority": 1}),
+            Document(id="3", content="Doc 3", meta={"category": "news", "status": "published", "priority": 5}),
+        ]
+        document_store.write_documents(docs)
+
+        count = document_store.count_documents_by_filter(
+            {"field": "meta.status", "operator": "==", "value": "published"}
+        )
+
+        assert count == 2
+
+    def test_count_unique_metadata_by_filter(self, document_store: AstraDocumentStore):
+        docs = [
+            Document(id="1", content="Doc 1", meta={"category": "news", "status": "published", "priority": 1}),
+            Document(id="2", content="Doc 2", meta={"category": "docs", "status": "published", "priority": 2}),
+            Document(id="3", content="Doc 3", meta={"category": "news", "status": "published", "priority": 2}),
+            Document(id="4", content="Doc 4", meta={"category": "faq", "status": "draft", "priority": 3}),
+        ]
+        document_store.write_documents(docs)
+
+        counts = document_store.count_unique_metadata_by_filter(
+            {"field": "meta.status", "operator": "==", "value": "published"},
+            ["category", "priority"],
+        )
+
+        assert counts == {"category": 2, "priority": 2}
+
+    def test_get_metadata_fields_info(self, document_store: AstraDocumentStore):
+        docs = [
+            Document(id="1", content="Doc 1", meta={"category": "news", "status": "published", "priority": 1}),
+            Document(id="2", content="Doc 2", meta={"category": "docs", "status": "draft", "priority": 2}),
+        ]
+        document_store.write_documents(docs)
+
+        fields_info = document_store.get_metadata_fields_info()
+
+        assert fields_info == {
+            "content": {"type": "text"},
+            "category": {"type": "keyword"},
+            "status": {"type": "keyword"},
+            "priority": {"type": "long"},
+        }
+
+    def test_get_metadata_field_min_max(self, document_store: AstraDocumentStore):
+        docs = [
+            Document(id="1", content="Doc 1", meta={"priority": 3}),
+            Document(id="2", content="Doc 2", meta={"priority": 1}),
+            Document(id="3", content="Doc 3", meta={"priority": 7}),
+        ]
+        document_store.write_documents(docs)
+
+        result = document_store.get_metadata_field_min_max("priority")
+
+        assert result == {"min": 1, "max": 7}
+
+    def test_get_metadata_field_unique_values(self, document_store: AstraDocumentStore):
+        docs = [
+            Document(id="1", content="Doc 1", meta={"category": "alpha"}),
+            Document(id="2", content="Doc 2", meta={"category": "beta"}),
+            Document(id="3", content="Doc 3", meta={"category": "alphabet"}),
+            Document(id="4", content="Doc 4", meta={"category": "gamma"}),
+        ]
+        document_store.write_documents(docs)
+
+        values, total_count = document_store.get_metadata_field_unique_values(
+            "category", search_term="alp", from_=0, size=10
+        )
+
+        assert values == ["alpha", "alphabet"]
+        assert total_count == 2
+
     @pytest.mark.skip(reason="Unsupported filter operator not.")
     def test_not_operator(self, document_store, filterable_docs):
         pass