From d206e673a69bb7c994889b8b6c44ec08c9222d50 Mon Sep 17 00:00:00 2001 From: Alonso Guevara Date: Wed, 13 Nov 2024 18:19:06 -0600 Subject: [PATCH] Format --- graphrag/index/emit/table_emitter.py | 1 + graphrag/utils/storage.py | 10 +++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/graphrag/index/emit/table_emitter.py b/graphrag/index/emit/table_emitter.py index 2e63c50fac..82e37d3135 100644 --- a/graphrag/index/emit/table_emitter.py +++ b/graphrag/index/emit/table_emitter.py @@ -10,6 +10,7 @@ class TableEmitter(Protocol): """TableEmitter protocol for emitting tables to a destination.""" + extension: str async def emit(self, name: str, data: pd.DataFrame) -> None: diff --git a/graphrag/utils/storage.py b/graphrag/utils/storage.py index 5863c46440..479d0c0cf3 100644 --- a/graphrag/utils/storage.py +++ b/graphrag/utils/storage.py @@ -51,11 +51,15 @@ async def _load_table_from_storage(name: str, storage: PipelineStorage) -> pd.Da case "parquet": return pd.read_parquet(BytesIO(await storage.get(name, as_bytes=True))) case "json": - return pd.read_json(BytesIO(await storage.get(name, as_bytes=True)), lines=True, orient="records") + return pd.read_json( + BytesIO(await storage.get(name, as_bytes=True)), + lines=True, + orient="records", + ) case "csv": return pd.read_csv(BytesIO(await storage.get(name, as_bytes=True))) - case _: - raise ValueError(f"Unknown file extension for {name}") + case _: + raise ValueError(f"Unknown file extension for {name}") except Exception: log.exception("error loading table from storage: %s", name) raise