Add separate() method to HeteroData for extracting connected components

[jesseangelis]

🚀 The feature, motivation and pitch

Hi PyG team,

I’d like to propose adding a new method to HeteroData for splitting a heterogeneous graph into its connected components. This would extend the existing subgraph and edge_subgraph utilities by using a union-find algorithm to detect and extract disjoint connected subgraphs.

Motivation:

Currently, PyG does not provide an out-of-the-box way to extract connected components from a heterogeneous graph structure (HeteroData).
This would help with tasks like isolating small disconnected graphs, cleaning up noise, preprocessing, and building mini-batches from large sparse graphs.
I believe it’s a natural and general-purpose extension alongside the existing subgraph and node_type_subgraph methods.

Proposed API:

components = hetero_data_object.separate(allowed_edge_types=[...], allowed_node_types=[...])

• Returns: List[HeteroData] — each item is a disconnected subgraph.
• Uses a union-find to detect connected components.
• Optionally restricts which edge/node types are included.

Prototype:

I’ve already written a prototype that:

• Uses local helper functions for union-find.
• Is consistent with PyG’s HeteroData style and idioms.
• Returns new HeteroData objects for each component.
• Includes pytests

If you’re interested, I can polish it up and submit a PR!

Let me know if you’d like any more details. I’m ready to start a PR whenever you give the go-ahead! 🚀

Alternatives

No response

Additional context

No response

[wsad1]

This is a nice proposal. Two questions

1. Would you be open to implementing this for Data (homogenous) graphs too.
2. allowed_edge_types=[...], allowed_node_types=[...]: what do these arguments do? Do we really need them? Instead can we use node_type_subgraph or edge_type_subgraph combined with this new function you'll add?

In general this is a good idea, and ok to accept a PR.

[jesseangelis]

Sure, I can take care of that for Data too.

I was thinking it could be useful to specify which node or edge types to consider when finding connected components, but combining it with the subgraph methods might be the cleaner approach.