Repository: microsoft/graphrag
File: graphrag/index/operations/build_noun_graph/build_noun_graph.py (Python)

Code before:
        noun_node_df.groupby("title").agg({"text_unit_id": list}).reset_index()
    )
    grouped_node_df = grouped_node_df.rename(columns={"text_unit_id": "text_unit_ids"})
    grouped_node_df["frequency"] = grouped_node_df["text_unit_ids"].apply(len)
    grouped_node_df = grouped_node_df[["title", "frequency", "text_unit_ids"]]
    return grouped_node_df.loc[:, ["title", "frequency", "text_unit_ids"]]


def _extract_edges(
    nodes_df: pd.DataFrame,
    normalize_edge_weights: bool = True,
) -> pd.DataFrame:
    """
    Extract edges from nodes.

    Nodes appear in the same text unit are connected.
    Input: nodes_df with schema [id, title, frequency, text_unit_ids]
    Returns: edges_df with schema [source, target, weight, text_unit_ids]
    """
    text_units_df = nodes_df.explode("text_unit_ids")
    text_units_df = text_units_df.rename(columns={"text_unit_ids": "text_unit_id"})
    text_units_df = (
        text_units_df.groupby("text_unit_id").agg({"title": list}).reset_index()
    )
    text_units_df["edges"] = text_units_df["title"].apply(
        lambda x: list(combinations(x, 2))
    )
    edge_df = text_units_df.explode("edges").loc[:, ["edges", "text_unit_id"]]

    edge_df["source"] = edge_df["edges"].apply(
        lambda x: x[0] if isinstance(x, tuple) else None
    )
    edge_df["target"] = edge_df["edges"].apply(
        lambda x: x[1] if isinstance(x, tuple) else None
    )
    edge_df = edge_df[(edge_df.source.notna()) & (edge_df.target.notna())]
    edge_df = edge_df.drop(columns=["edges"])

    # make sure source is always smaller than target
    edge_df["source"], edge_df["target"] = zip(
        *edge_df.apply(
            lambda x: (x["source"], x["target"])
            if x["source"] < x["target"]
            else (x["target"], x["source"]),
            axis=1,
        ),
        strict=False,
    )

    # group by source and target, count the number of text units
    grouped_edge_df = (

Review comment (suggestion):
I think we could do something like:

```python
titles = text_units_df["title"].tolist()
all_edges = [ list(combinations(t, 2)) for t in titles]

text_units_df = text_units_df.assign(edges=all_edges)
```
to get rid of the apply. May need to testing tho, but I think it could work.

Code after:
    )
    grouped_node_df = grouped_node_df.rename(columns={"text_unit_id": "text_unit_ids"})
    grouped_node_df["frequency"] = grouped_node_df["text_unit_ids"].apply(len)
    grouped_node_df = grouped_node_df[["title", "frequency", "text_unit_ids"]]
    return grouped_node_df.loc[:, ["title", "frequency", "text_unit_ids"]]


def _extract_edges(
    nodes_df: pd.DataFrame,
    normalize_edge_weights: bool = True,
) -> pd.DataFrame:
    """
    Extract edges from nodes.

    Nodes appear in the same text unit are connected.
    Input: nodes_df with schema [id, title, frequency, text_unit_ids]
    Returns: edges_df with schema [source, target, weight, text_unit_ids]
    """
    text_units_df = nodes_df.explode("text_unit_ids")
    text_units_df = text_units_df.rename(columns={"text_unit_ids": "text_unit_id"})

    text_units_df = (
        text_units_df.groupby("text_unit_id")
        .agg({"title": lambda x: list(x) if len(x) > 1 else np.nan})
        .reset_index()
    )
    text_units_df = text_units_df.dropna()
    titles = text_units_df["title"].tolist()
    all_edges: list[list[tuple[str, str]]] = [list(combinations(t, 2)) for t in titles]

    text_units_df = text_units_df.assign(edges=all_edges)  # type: ignore
    edge_df = text_units_df.explode("edges")[["edges", "text_unit_id"]]

    edge_df[["source", "target"]] = edge_df.loc[:, "edges"].to_list()
    edge_df["min_source"] = edge_df[["source", "target"]].min(axis=1)
    edge_df["max_target"] = edge_df[["source", "target"]].max(axis=1)
    edge_df = edge_df.drop(columns=["source", "target"]).rename(
        columns={"min_source": "source", "max_target": "target"}  # type: ignore
    )

    edge_df = edge_df[(edge_df.source.notna()) & (edge_df.target.notna())]
    edge_df = edge_df.drop(columns=["edges"])
    # group by source and target, count the number of text units
    grouped_edge_df = (
        edge_df.groupby(["source", "target"]).agg({"text_unit_id": list}).reset_index()
    )
    grouped_edge_df = grouped_edge_df.rename(columns={"text_unit_id": "text_unit_ids"})
    grouped_edge_df["weight"] = grouped_edge_df["text_unit_ids"].apply(len)
    grouped_edge_df = grouped_edge_df.loc[
        :, ["source", "target", "weight", "text_unit_ids"]
    ]