Add numberOfVectors to get_stats api

src/marqo/tensor_search/tensor_search.py

@@ -267,9 +267,24 @@ def _autofill_index_settings(index_settings: dict):
 
 
 def get_stats(config: Config, index_name: str):
-    doc_count = HttpRequests(config).post(path=F"{index_name}/_count")["count"]
+    """Returns the number of documents and vectors in the index.
+
+    The _count API counts top-level documents.
+    The _stats API includes the count of nested documents, which is __chunks in Marqo.
+
+    Difference between the two gives the numberOfVectors."""
+
+    try:
+        doc_count = HttpRequests(config).post(path=F"{index_name}/_count")["count"]
+        nested_doc_count = HttpRequests(config).get(path=F"{index_name}/_stats/docs")["indices"][index_name]["primaries"][
+            "docs"]["count"]
+    except KeyError as e:
+        return errors.InternalError(f"Marqo encountered an unexpected response from Marqo-os when calling `get_stats(index=)`. "
+                                   f"The expected fields do not exist in the response. Original error message = {e}")
+
     return {
-        "numberOfDocuments": doc_count
+        "numberOfDocuments": doc_count,
+        "numberOfVectors": nested_doc_count - doc_count
     }
 
 

tests/tensor_search/test_get_stats.py

@@ -14,25 +14,82 @@ def setUp(self) -> None:
         except IndexNotFoundError as s:
             pass
 
-    def test_get_stats_empty(self):
+    def tearDown(self) -> None:
         try:
             tensor_search.delete_index(config=self.config, index_name=self.index_name_1)
         except IndexNotFoundError as s:
             pass
+
+    def test_get_stats_empty(self):
         tensor_search.create_vector_index(config=self.config, index_name=self.index_name_1)
         assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfDocuments"] == 0
 
     def test_get_stats_non_empty(self):
-        try:
-            tensor_search.delete_index(config=self.config, index_name=self.index_name_1)
-        except IndexNotFoundError as s:
-            pass
         tensor_search.create_vector_index(config=self.config, index_name=self.index_name_1)
         tensor_search.add_documents(
             config=self.config, add_docs_params=AddDocsParams(
-                docs=[{"1": "2"},{"134": "2"},{"14": "62"}],
+                docs=[{"1": "2"}, {"134": "2"}, {"14": "62"}],
                 index_name=self.index_name_1,
                 auto_refresh=True, device="cpu"
             )
         )
         assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfDocuments"] == 3
+
+    def test_get_stats_number_of_vectors(self):
+
+        tensor_search.create_vector_index(config=self.config, index_name=self.index_name_1,
+                                          index_settings={'index_defaults': {"model": "random/small"}})
+        expected_number_of_vectors = 7
+        expected_number_of_documents = 5
+        tensor_search.add_documents(
+            config=self.config, add_docs_params=AddDocsParams(
+                docs=[
+                    {"description_1": "test-2", "description_2": "test"},  # 2 vectors
+                    {"description_1": "test-2", "description_2": "test", "description_3": "test"},  # 3 vectors
+                    {"description_2": "test"},  # 1 vector
+                    {"my_multi_modal_field": {
+                        "text_1": "test", "text_2": "test"}},  # 1 vector
+                    {"non_tensor_field": "test"}  # 0 vectors
+                ],
+                index_name=self.index_name_1,
+                auto_refresh=True, device="cpu",
+                non_tensor_fields=["non_tensor_field"],
+                mappings={"my_multi_modal_field": {"type": "multimodal_combination", "weights": {
+                    "text_1": 0.5, "text_2": 0.8}}}
+            )
+        )
+
+        assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfDocuments"] \
+               == expected_number_of_documents
+        assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfVectors"] \
+               == expected_number_of_vectors
+
+    def test_get_stats_number_of_vectors(self):
+
+        tensor_search.create_vector_index(config=self.config, index_name=self.index_name_1,
+                                          index_settings={'index_defaults': {"model": "random/small"}})
+        expected_number_of_vectors = 7
+        expected_number_of_documents = 6
+        res = tensor_search.add_documents(
+            config=self.config, add_docs_params=AddDocsParams(
+                docs=[
+                    {"description_1": "test-2", "description_2": "test"},  # 2 vectors
+                    {"description_1": "test-2", "description_2": "test", "description_3": "test"},  # 3 vectors
+                    {"description_2": "test"},  # 1 vector
+                    {"my_multi_modal_field": {
+                        "text_1": "test", "text_2": "test"}},  # 1 vector
+                    {"non_tensor_field": "test"},  # 0 vectors
+                    {"list_field": ["this", "that"]}, # 0 vectors
+                ],
+                index_name=self.index_name_1,
+                auto_refresh=True, device="cpu",
+                non_tensor_fields=["non_tensor_field", "list_field"],
+                mappings={"my_multi_modal_field":
+                              {"type": "multimodal_combination", "weights": {"text_1": 0.5, "text_2": 0.8}}}
+            )
+        )
+
+        assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfDocuments"] \
+               == expected_number_of_documents
+        assert tensor_search.get_stats(config=self.config, index_name=self.index_name_1)["numberOfVectors"] \
+               == expected_number_of_vectors