Collapse create final relationships (#1158)

* Collapse pre/post embedding workflows

* Semver

* Fix smoke tests

---------

Co-authored-by: Alonso Guevara <alonsog@microsoft.com>

Author: natoverse

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

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Collapse create_final_relationships."
+}

graphrag/index/verbs/graph/compute_edge_combined_degree.py

@@ -32,10 +32,34 @@ def compute_edge_combined_degree(
     - to: The name of the column to output the combined degree to. Default="rank"
     """
     edge_df: pd.DataFrame = cast(pd.DataFrame, input.get_input())
-    if to in edge_df.columns:
-        return TableContainer(table=edge_df)
     node_degree_df = _get_node_degree_table(input, node_name_column, node_degree_column)
 
+    output_df = compute_edge_combined_degree_df(
+        edge_df,
+        node_degree_df,
+        to,
+        node_name_column,
+        node_degree_column,
+        edge_source_column,
+        edge_target_column,
+    )
+
+    return TableContainer(table=output_df)
+
+
+def compute_edge_combined_degree_df(
+    edge_df: pd.DataFrame,
+    node_degree_df: pd.DataFrame,
+    to: str,
+    node_name_column: str,
+    node_degree_column: str,
+    edge_source_column: str,
+    edge_target_column: str,
+) -> pd.DataFrame:
+    """Compute the combined degree for each edge in a graph."""
+    if to in edge_df.columns:
+        return edge_df
+
     def join_to_degree(df: pd.DataFrame, column: str) -> pd.DataFrame:
         degree_column = _degree_colname(column)
         result = df.merge(
@@ -48,14 +72,13 @@ def join_to_degree(df: pd.DataFrame, column: str) -> pd.DataFrame:
         result[degree_column] = result[degree_column].fillna(0)
         return result
 
-    edge_df = join_to_degree(edge_df, edge_source_column)
-    edge_df = join_to_degree(edge_df, edge_target_column)
-    edge_df[to] = (
-        edge_df[_degree_colname(edge_source_column)]
-        + edge_df[_degree_colname(edge_target_column)]
+    output_df = join_to_degree(edge_df, edge_source_column)
+    output_df = join_to_degree(output_df, edge_target_column)
+    output_df[to] = (
+        output_df[_degree_colname(edge_source_column)]
+        + output_df[_degree_colname(edge_target_column)]
     )
-
-    return TableContainer(table=edge_df)
+    return output_df
 
 
 def _degree_colname(column: str) -> str:

graphrag/index/workflows/v1/create_final_relationships.py

@@ -16,6 +16,7 @@ def build_steps(
 
     ## Dependencies
     * `workflow:create_base_entity_graph`
+    * `workflow:create_final_nodes`
     """
     base_text_embed = config.get("text_embed", {})
     relationship_description_embed_config = config.get(
@@ -25,25 +26,12 @@ def build_steps(
 
     return [
         {
-            "verb": "unpack_graph",
-            "args": {
-                "column": "clustered_graph",
-                "type": "edges",
-            },
+            "id": "pre_embedding",
+            "verb": "create_final_relationships_pre_embedding",
             "input": {"source": "workflow:create_base_entity_graph"},
         },
         {
-            "verb": "rename",
-            "args": {"columns": {"source_id": "text_unit_ids"}},
-        },
-        {
-            "verb": "filter",
-            "args": {
-                "column": "level",
-                "criteria": [{"type": "value", "operator": "equals", "value": 0}],
-            },
-        },
-        {
+            "id": "description_embedding",
             "verb": "text_embed",
             "enabled": not skip_description_embedding,
             "args": {
@@ -54,41 +42,12 @@ def build_steps(
             },
         },
         {
-            "id": "pruned_edges",
-            "verb": "drop",
-            "args": {"columns": ["level"]},
-        },
-        {
-            "id": "filtered_nodes",
-            "verb": "filter",
-            "args": {
-                "column": "level",
-                "criteria": [{"type": "value", "operator": "equals", "value": 0}],
-            },
-            "input": "workflow:create_final_nodes",
-        },
-        {
-            "verb": "compute_edge_combined_degree",
-            "args": {"to": "rank"},
+            "verb": "create_final_relationships_post_embedding",
             "input": {
-                "source": "pruned_edges",
-                "nodes": "filtered_nodes",
-            },
-        },
-        {
-            "verb": "convert",
-            "args": {
-                "column": "human_readable_id",
-                "type": "string",
-                "to": "human_readable_id",
-            },
-        },
-        {
-            "verb": "convert",
-            "args": {
-                "column": "text_unit_ids",
-                "type": "array",
-                "to": "text_unit_ids",
+                "source": "pre_embedding"
+                if skip_description_embedding
+                else "description_embedding",
+                "nodes": "workflow:create_final_nodes",
             },
         },
     ]

graphrag/index/workflows/v1/subflows/__init__.py

@@ -4,9 +4,17 @@
 """The Indexing Engine workflows -> subflows package root."""
 
 from .create_final_communities import create_final_communities
+from .create_final_relationships_post_embedding import (
+    create_final_relationships_post_embedding,
+)
+from .create_final_relationships_pre_embedding import (
+    create_final_relationships_pre_embedding,
+)
 from .create_final_text_units_pre_embedding import create_final_text_units_pre_embedding
 
 __all__ = [
     "create_final_communities",
+    "create_final_relationships_post_embedding",
+    "create_final_relationships_pre_embedding",
     "create_final_text_units_pre_embedding",
 ]

graphrag/index/workflows/v1/subflows/create_final_relationships_post_embedding.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""All the steps to transform final relationships after they are embedded."""
+
+from typing import Any, cast
+
+import pandas as pd
+from datashaper import (
+    Table,
+    VerbInput,
+    verb,
+)
+from datashaper.table_store.types import VerbResult, create_verb_result
+
+from graphrag.index.utils.ds_util import get_required_input_table
+from graphrag.index.verbs.graph.compute_edge_combined_degree import (
+    compute_edge_combined_degree_df,
+)
+
+
+@verb(
+    name="create_final_relationships_post_embedding",
+    treats_input_tables_as_immutable=True,
+)
+def create_final_relationships_post_embedding(
+    input: VerbInput,
+    **_kwargs: dict,
+) -> VerbResult:
+    """All the steps to transform final relationships after they are embedded."""
+    table = cast(pd.DataFrame, input.get_input())
+    nodes = cast(pd.DataFrame, get_required_input_table(input, "nodes").table)
+
+    pruned_edges = table.drop(columns=["level"])
+
+    filtered_nodes = cast(
+        pd.DataFrame,
+        nodes[nodes["level"] == 0].reset_index(drop=True)[["title", "degree"]],
+    )
+
+    edge_combined_degree = compute_edge_combined_degree_df(
+        pruned_edges,
+        filtered_nodes,
+        to="rank",
+        node_name_column="title",
+        node_degree_column="degree",
+        edge_source_column="source",
+        edge_target_column="target",
+    )
+
+    edge_combined_degree["human_readable_id"] = edge_combined_degree[
+        "human_readable_id"
+    ].astype(str)
+    edge_combined_degree["text_unit_ids"] = _to_array(
+        edge_combined_degree["text_unit_ids"], ","
+    )
+
+    return create_verb_result(cast(Table, edge_combined_degree))
+
+
+# from datashaper, we should be able to inline this
+def _to_array(column, delimiter: str):
+    def convert_value(value: Any) -> list:
+        if pd.isna(value):
+            return []
+        if isinstance(value, list):
+            return value
+        if isinstance(value, str):
+            return value.split(delimiter)
+        return [value]
+
+    return column.apply(convert_value)

graphrag/index/workflows/v1/subflows/create_final_relationships_pre_embedding.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""All the steps to transform final relationships before they are embedded."""
+
+from typing import cast
+
+import pandas as pd
+from datashaper import (
+    Table,
+    VerbCallbacks,
+    VerbInput,
+    verb,
+)
+from datashaper.table_store.types import VerbResult, create_verb_result
+
+from graphrag.index.verbs.graph.unpack import unpack_graph_df
+
+
+@verb(
+    name="create_final_relationships_pre_embedding",
+    treats_input_tables_as_immutable=True,
+)
+def create_final_relationships_pre_embedding(
+    input: VerbInput,
+    callbacks: VerbCallbacks,
+    **_kwargs: dict,
+) -> VerbResult:
+    """All the steps to transform final relationships before they are embedded."""
+    table = cast(pd.DataFrame, input.get_input())
+
+    graph_edges = unpack_graph_df(table, callbacks, "clustered_graph", "edges")
+
+    graph_edges.rename(columns={"source_id": "text_unit_ids"}, inplace=True)
+
+    filtered = graph_edges[graph_edges["level"] == 0].reset_index(drop=True)
+
+    return create_verb_result(cast(Table, filtered))

tests/fixtures/min-csv/config.json

@@ -52,7 +52,7 @@
                 1,
                 2000
             ],
-            "subworkflows": 8,
+            "subworkflows": 2,
             "max_runtime": 100
         },
         "create_final_nodes": {

tests/fixtures/text/config.json

@@ -71,7 +71,7 @@
                 1,
                 2000
             ],
-            "subworkflows": 8,
+            "subworkflows": 2,
             "max_runtime": 100
         },
         "create_final_nodes": {

tests/verbs/test_create_final_relationships.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+from graphrag.index.workflows.v1.create_final_relationships import (
+    build_steps,
+    workflow_name,
+)
+
+from .util import (
+    compare_outputs,
+    get_config_for_workflow,
+    get_workflow_output,
+    load_expected,
+    load_input_tables,
+    remove_disabled_steps,
+)
+
+
+async def test_create_final_relationships():
+    input_tables = load_input_tables([
+        "workflow:create_base_entity_graph",
+        "workflow:create_final_nodes",
+    ])
+    expected = load_expected(workflow_name)
+
+    config = get_config_for_workflow(workflow_name)
+
+    config["skip_description_embedding"] = True
+
+    steps = remove_disabled_steps(build_steps(config))
+
+    actual = await get_workflow_output(
+        input_tables,
+        {
+            "steps": steps,
+        },
+    )
+
+    compare_outputs(actual, expected)