add find test

Author: jgbradley1

graphrag/index/run/workflow.py

@@ -33,7 +33,6 @@ async def _inject_workflow_data_dependencies(
     workflow_dependencies: dict[str, list[str]],
     dataset: pd.DataFrame,
     storage: PipelineStorage,
-    extension: str,
 ) -> None:
     """Inject the data dependencies into the workflow."""
     workflow.add_table(DEFAULT_INPUT_NAME, dataset)
@@ -42,7 +41,7 @@ async def _inject_workflow_data_dependencies(
     for id in deps:
         workflow_id = f"workflow:{id}"
         try:
-            table = await _load_table_from_storage(f"{id}.{extension}", storage)
+            table = await _load_table_from_storage(f"{id}.parquet", storage)
         except ValueError:
             # our workflows allow for transient tables, and we avoid putting those in storage
             # however, we need to keep the table in the dependency list for proper execution order.
@@ -98,7 +97,10 @@ async def _process_workflow(
     context.stats.workflows[workflow_name] = {"overall": 0.0}
 
     await _inject_workflow_data_dependencies(
-        workflow, workflow_dependencies, dataset, context.storage, "parquet"
+        workflow,
+        workflow_dependencies,
+        dataset,
+        context.storage,
     )
 
     workflow_start_time = time.time()

graphrag/storage/blob_pipeline_storage.py

@@ -32,7 +32,7 @@ def __init__(
         self,
         connection_string: str | None,
         container_name: str,
-        encoding: str | None = None,
+        encoding: str = "utf-8",
         path_prefix: str | None = None,
         storage_account_blob_url: str | None = None,
     ):
@@ -50,7 +50,7 @@ def __init__(
                 account_url=storage_account_blob_url,
                 credential=DefaultAzureCredential(),
             )
-        self._encoding = encoding or "utf-8"
+        self._encoding = encoding
         self._container_name = container_name
         self._connection_string = connection_string
         self._path_prefix = path_prefix or ""
@@ -198,7 +198,7 @@ async def set(self, key: str, value: Any, encoding: str | None = None) -> None:
             if isinstance(value, bytes):
                 blob_client.upload_blob(value, overwrite=True)
             else:
-                coding = encoding or "utf-8"
+                coding = encoding or self._encoding
                 blob_client.upload_blob(value.encode(coding), overwrite=True)
         except Exception:
             log.exception("Error setting key %s: %s", key)

graphrag/storage/cosmosdb_pipeline_storage.py

@@ -28,7 +28,7 @@ class CosmosDBPipelineStorage(PipelineStorage):
     _cosmosdb_account_url: str | None
     _connection_string: str | None
     _database_name: str
-    _current_container: str | None
+    container_name: str | None
     _encoding: str
 
     def __init__(
@@ -58,7 +58,7 @@ def __init__(
         self._database_name = database_name
         self._connection_string = connection_string
         self._cosmosdb_account_url = cosmosdb_account_url
-        self._current_container = current_container
+        self.container_name = current_container
         self._cosmosdb_account_name = (
             cosmosdb_account_url.split("//")[1].split(".")[0]
             if cosmosdb_account_url
@@ -74,7 +74,7 @@ def __init__(
             self._database_name,
         )
         self._create_database()
-        if self._current_container:
+        if self.container_name:
             self._create_container()
 
     def _create_database(self) -> None:
@@ -117,7 +117,7 @@ def find(
 
         log.info(
             "search container %s for documents matching %s",
-            self._current_container,
+            self.container_name,
             file_pattern.pattern,
         )
 
@@ -130,7 +130,7 @@ def item_filter(item: dict[str, Any]) -> bool:
 
         try:
             container_client = self._database_client.get_container_client(
-                str(self._current_container)
+                str(self.container_name)
             )
             query = "SELECT * FROM c WHERE CONTAINS(c.id, @pattern)"
             parameters: list[dict[str, Any]] = [
@@ -173,9 +173,9 @@ async def get(
     ) -> Any:
         """Get a file in the database for the given key."""
         try:
-            if self._current_container:
+            if self.container_name:
                 container_client = self._database_client.get_container_client(
-                    self._current_container
+                    self.container_name
                 )
                 item = container_client.read_item(item=key, partition_key=key)
                 item_body = item.get("body")
@@ -195,9 +195,9 @@ async def get(
     async def set(self, key: str, value: Any, encoding: str | None = None) -> None:
         """Set a file in the database for the given key."""
         try:
-            if self._current_container:
+            if self.container_name:
                 container_client = self._database_client.get_container_client(
-                    self._current_container
+                    self.container_name
                 )
                 if isinstance(value, bytes):
                     value_df = pd.read_parquet(BytesIO(value))
@@ -217,29 +217,29 @@ async def set(self, key: str, value: Any, encoding: str | None = None) -> None:
 
     async def has(self, key: str) -> bool:
         """Check if the given file exists in the cosmosdb storage."""
-        if self._current_container:
+        if self.container_name:
             container_client = self._database_client.get_container_client(
-                self._current_container
+                self.container_name
             )
             item_names = [item["id"] for item in container_client.read_all_items()]
             return key in item_names
         return False
 
     async def delete(self, key: str) -> None:
         """Delete the given file from the cosmosdb storage."""
-        if self._current_container:
+        if self.container_name:
             container_client = self._database_client.get_container_client(
-                self._current_container
+                self.container_name
             )
             container_client.delete_item(item=key, partition_key=key)
 
     # Function currently deletes all items within the current container, then deletes the container itself.
     # TODO: Decide the granularity of deletion (e.g. delete all items within the current container, delete the current container, delete the current database)
     async def clear(self) -> None:
         """Clear the cosmosdb storage."""
-        if self._current_container:
+        if self.container_name:
             container_client = self._database_client.get_container_client(
-                self._current_container
+                self.container_name
             )
             for item in container_client.read_all_items():
                 item_id = item["id"]
@@ -257,24 +257,24 @@ def child(self, name: str | None) -> PipelineStorage:
 
     def _create_container(self) -> None:
         """Create a container for the current container name if it doesn't exist."""
-        if self._current_container:
+        if self.container_name:
             partition_key = PartitionKey(path="/id", kind="Hash")
             self._database_client.create_container_if_not_exists(
-                id=self._current_container,
+                id=self.container_name,
                 partition_key=partition_key,
             )
 
     def _delete_container(self) -> None:
         """Delete the container with the current container name if it exists."""
-        if self._container_exists() and self._current_container:
-            self._database_client.delete_container(self._current_container)
+        if self._container_exists() and self.container_name:
+            self._database_client.delete_container(self.container_name)
 
     def _container_exists(self) -> bool:
         """Check if the container with the current container name exists."""
         container_names = [
             container["id"] for container in self._database_client.list_containers()
         ]
-        return self._current_container in container_names
+        return self.container_name in container_names
 
 
 # TODO remove this helper function and have the factory instantiate the class directly

graphrag/utils/storage.py

@@ -4,7 +4,7 @@
 """Storage functions for the GraphRAG run module."""
 
 import logging
-from io import BytesIO, StringIO
+from io import BytesIO
 
 import pandas as pd
 
@@ -18,22 +18,8 @@ async def _load_table_from_storage(name: str, storage: PipelineStorage) -> pd.Da
         msg = f"Could not find {name} in storage!"
         raise ValueError(msg)
     try:
-        log.info("read table from storage: %s", name)
-        match name.split(".")[-1]:
-            case "parquet":
-                return pd.read_parquet(BytesIO(await storage.get(name, as_bytes=True)))
-            case "json":
-                return pd.read_json(
-                    StringIO(await storage.get(name)),
-                    lines=False,
-                    orient="records",
-                )
-            case "csv":
-                return pd.read_csv(BytesIO(await storage.get(name, as_bytes=True)))
-            case _:
-                msg = f"Unknown file extension for {name}"
-                log.exception(msg)
-                raise
+        log.info("reading table from storage: %s", name)
+        return pd.read_parquet(BytesIO(await storage.get(name, as_bytes=True)))
     except Exception:
         log.exception("error loading table from storage: %s", name)
         raise

tests/integration/storage/test_cosmosdb_storage.py

@@ -2,20 +2,57 @@
 # Licensed under the MIT License
 """CosmosDB Storage Tests."""
 
+import re
 import sys
 
 import pytest
 
+from graphrag.storage.cosmosdb_pipeline_storage import CosmosDBPipelineStorage
+
+# cspell:disable-next-line well-known-key
+WELL_KNOWN_COSMOS_ACCOUNT_URL = "https://localhost:8081"
+WELL_KNOWN_COSMOS_CONNECTION_STRING = "AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==;"
+
 # the cosmosdb emulator is only available on windows runners at this time
 if not sys.platform.startswith("win"):
     pytest.skip("encountered windows-only tests -- skipping", allow_module_level=True)
 
 
-def test_find():
-    print("test_find")
-    assert True
+async def test_find():
+    storage = CosmosDBPipelineStorage(
+        cosmosdb_account_url=WELL_KNOWN_COSMOS_ACCOUNT_URL,
+        connection_string=WELL_KNOWN_COSMOS_CONNECTION_STRING,
+        database_name="testfind",
+    )
+    try:
+        try:
+            items = list(
+                storage.find(base_dir="input", file_pattern=re.compile(r".*\.txt$"))
+            )
+            items = [item[0] for item in items]
+            assert items == []
+
+            await storage.set("christmas.txt", "Merry Christmas!", encoding="utf-8")
+            items = list(
+                storage.find(base_dir="input", file_pattern=re.compile(r".*\.txt$"))
+            )
+            items = [item[0] for item in items]
+            assert items == ["christmas.txt"]
+
+            await storage.set("test.txt", "Hello, World!", encoding="utf-8")
+            items = list(storage.find(file_pattern=re.compile(r".*\.txt$")))
+            items = [item[0] for item in items]
+            assert items == ["christmas.txt", "test.txt"]
+
+            output = await storage.get("test.txt")
+            assert output == "Hello, World!"
+        finally:
+            await storage.delete("test.txt")
+            output = await storage.get("test.txt")
+            assert output is None
+    finally:
+        storage._delete_container()  # noqa: SLF001
 
 
 def test_child():
-    print("test_child")
     assert True