Transient entity graph (#1349)

* Make base_entity_graph transient

* Add transient snapshots

* Semver

* Fix unit test

* Fix smoke tests

Author: natoverse

.semversioner/next-release/patch-20241105004012425642.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Transient entity graph and snapshotting."
+}

docs/config/env_vars.md

@@ -199,11 +199,13 @@ This section controls the reporting mechanism used by the pipeline, for common e
 
 ## Data Snapshotting
 
-| Parameter                           | Description                                 | Type   | Required or Optional | Default |
-| ----------------------------------- | ------------------------------------------- | ------ | -------------------- | ------- |
-| `GRAPHRAG_SNAPSHOT_GRAPHML`         | Whether to enable GraphML snapshots.        | `bool` | optional             | False   |
-| `GRAPHRAG_SNAPSHOT_RAW_ENTITIES`    | Whether to enable raw entity snapshots.     | `bool` | optional             | False   |
-| `GRAPHRAG_SNAPSHOT_TOP_LEVEL_NODES` | Whether to enable top-level node snapshots. | `bool` | optional             | False   |
+| Parameter                              | Description                                     | Type   | Required or Optional | Default |
+| -------------------------------------- | ----------------------------------------------- | ------ | -------------------- | ------- |
+| `GRAPHRAG_SNAPSHOT_EMBEDDINGS`         | Whether to enable embeddings snapshots.         | `bool` | optional             | False   |
+| `GRAPHRAG_SNAPSHOT_GRAPHML`            | Whether to enable GraphML snapshots.            | `bool` | optional             | False   |
+| `GRAPHRAG_SNAPSHOT_RAW_ENTITIES`       | Whether to enable raw entity snapshots.         | `bool` | optional             | False   |
+| `GRAPHRAG_SNAPSHOT_TOP_LEVEL_NODES`    | Whether to enable top-level node snapshots.     | `bool` | optional             | False   |
+| `GRAPHRAG_SNAPSHOT_TRANSIENT`          | Whether to enable transient table snapshots.    | `bool` | optional             | False   |
 
 # Miscellaneous Settings
 

docs/config/json_yaml.md

@@ -216,9 +216,11 @@ This is the base LLM configuration section. Other steps may override this config
 
 ### Fields
 
-- `graphml` **bool** - Emit graphml snapshots.
-- `raw_entities` **bool** - Emit raw entity snapshots.
-- `top_level_nodes` **bool** - Emit top-level-node snapshots.
+- `embeddings` **bool** - Emit embeddings snapshots to parquet.
+- `graphml` **bool** - Emit graph snapshots to GraphML.
+- `raw_entities` **bool** - Emit raw entity snapshots to JSON.
+- `top_level_nodes` **bool** - Emit top-level-node snapshots to JSON.
+- `transient` **bool** - Emit transient workflow tables snapshots to parquet.
 
 ## encoding_model
 

graphrag/config/create_graphrag_config.py

@@ -415,6 +415,7 @@ def hydrate_parallelization_params(
                 top_level_nodes=reader.bool("top_level_nodes")
                 or defs.SNAPSHOTS_TOP_LEVEL_NODES,
                 embeddings=reader.bool("embeddings") or defs.SNAPSHOTS_EMBEDDINGS,
+                transient=reader.bool("transient") or defs.SNAPSHOTS_TRANSIENT,
             )
         with reader.envvar_prefix(Section.umap), reader.use(values.get("umap")):
             umap_model = UmapConfig(

graphrag/config/defaults.py

@@ -83,6 +83,7 @@
 SNAPSHOTS_RAW_ENTITIES = False
 SNAPSHOTS_TOP_LEVEL_NODES = False
 SNAPSHOTS_EMBEDDINGS = False
+SNAPSHOTS_TRANSIENT = False
 STORAGE_BASE_DIR = "output"
 STORAGE_TYPE = StorageType.file
 SUMMARIZE_DESCRIPTIONS_MAX_LENGTH = 500

graphrag/config/input_models/snapshots_config_input.py

@@ -9,6 +9,8 @@
 class SnapshotsConfigInput(TypedDict):
     """Configuration section for snapshots."""
 
+    embeddings: NotRequired[bool | str | None]
     graphml: NotRequired[bool | str | None]
     raw_entities: NotRequired[bool | str | None]
     top_level_nodes: NotRequired[bool | str | None]
+    transient: NotRequired[bool | str | None]

graphrag/config/models/snapshots_config.py

@@ -11,6 +11,10 @@
 class SnapshotsConfig(BaseModel):
     """Configuration section for snapshots."""
 
+    embeddings: bool = Field(
+        description="A flag indicating whether to take snapshots of embeddings.",
+        default=defs.SNAPSHOTS_EMBEDDINGS,
+    )
     graphml: bool = Field(
         description="A flag indicating whether to take snapshots of GraphML.",
         default=defs.SNAPSHOTS_GRAPHML,
@@ -23,7 +27,7 @@ class SnapshotsConfig(BaseModel):
         description="A flag indicating whether to take snapshots of top-level nodes.",
         default=defs.SNAPSHOTS_TOP_LEVEL_NODES,
     )
-    embeddings: bool = Field(
-        description="A flag indicating whether to take snapshots of embeddings.",
-        default=defs.SNAPSHOTS_EMBEDDINGS,
+    transient: bool = Field(
+        description="A flag indicating whether to take snapshots of transient tables.",
+        default=defs.SNAPSHOTS_TRANSIENT,
     )

graphrag/index/create_pipeline_config.py

@@ -171,6 +171,7 @@ def _text_unit_workflows(
         PipelineWorkflowReference(
             name=create_base_text_units,
             config={
+                "snapshot_transient": settings.snapshots.transient,
                 "chunk_by": settings.chunks.group_by_columns,
                 "text_chunk": {
                     "strategy": settings.chunks.resolved_strategy(
@@ -215,7 +216,9 @@ def _graph_workflows(settings: GraphRagConfig) -> list[PipelineWorkflowReference
         PipelineWorkflowReference(
             name=create_base_entity_graph,
             config={
-                "graphml_snapshot": settings.snapshots.graphml,
+                "snapshot_graphml": settings.snapshots.graphml,
+                "snapshot_transient": settings.snapshots.transient,
+                "snapshot_raw_entities": settings.snapshots.raw_entities,
                 "entity_extract": {
                     **settings.entity_extraction.parallelization.model_dump(),
                     "async_mode": settings.entity_extraction.async_mode,

graphrag/index/flows/create_base_entity_graph.py

@@ -42,8 +42,9 @@ async def create_base_entity_graph(
     summarization_strategy: dict[str, Any] | None = None,
     summarization_num_threads: int = 4,
     embedding_strategy: dict[str, Any] | None = None,
-    graphml_snapshot_enabled: bool = False,
-    raw_entity_snapshot_enabled: bool = False,
+    snapshot_graphml_enabled: bool = False,
+    snapshot_raw_entities_enabled: bool = False,
+    snapshot_transient_enabled: bool = False,
 ) -> pd.DataFrame:
     """All the steps to create the base entity graph."""
     # this returns a graph for each text unit, to be merged later
@@ -92,15 +93,15 @@ async def create_base_entity_graph(
             strategy=embedding_strategy,
         )
 
-    if raw_entity_snapshot_enabled:
+    if snapshot_raw_entities_enabled:
         await snapshot(
             entities,
             name="raw_extracted_entities",
             storage=storage,
             formats=["json"],
         )
 
-    if graphml_snapshot_enabled:
+    if snapshot_graphml_enabled:
         await snapshot_graphml(
             merged_graph,
             name="merged_graph",
@@ -131,4 +132,14 @@ async def create_base_entity_graph(
     if embedding_strategy:
         final_columns.append("embeddings")
 
-    return cast(pd.DataFrame, clustered[final_columns])
+    output = cast(pd.DataFrame, clustered[final_columns])
+
+    if snapshot_transient_enabled:
+        await snapshot(
+            output,
+            name="create_base_entity_graph",
+            storage=storage,
+            formats=["parquet"],
+        )
+
+    return output

graphrag/index/flows/create_base_text_units.py

@@ -15,16 +15,20 @@
 )
 
 from graphrag.index.operations.chunk_text import chunk_text
+from graphrag.index.operations.snapshot import snapshot
+from graphrag.index.storage import PipelineStorage
 from graphrag.index.utils import gen_md5_hash
 
 
-def create_base_text_units(
+async def create_base_text_units(
     documents: pd.DataFrame,
     callbacks: VerbCallbacks,
+    storage: PipelineStorage,
     chunk_column_name: str,
     n_tokens_column_name: str,
     chunk_by_columns: list[str],
     chunk_strategy: dict[str, Any] | None = None,
+    snapshot_transient_enabled: bool = False,
 ) -> pd.DataFrame:
     """All the steps to transform base text_units."""
     sort = documents.sort_values(by=["id"], ascending=[True])
@@ -73,10 +77,20 @@ def create_base_text_units(
     )
     chunked["id"] = chunked["chunk_id"]
 
-    return cast(
+    output = cast(
         pd.DataFrame, chunked[chunked[chunk_column_name].notna()].reset_index(drop=True)
     )
 
+    if snapshot_transient_enabled:
+        await snapshot(
+            output,
+            name="create_base_text_units",
+            storage=storage,
+            formats=["parquet"],
+        )
+
+    return output
+
 
 # TODO: would be nice to inline this completely in the main method with pandas
 def _aggregate_df(