how to get the all or part of k-hop neighbors of specific nodes in a graph?

[RX28666]

Hello,

Suppose I have a graph formatted in Data, such as Data(num_nodes=169343, x=[169343, 128], y=[169343], adj_t=[169343, 169343, nnz=2484941], train_mask=[169343], val_mask=[169343], test_mask=[169343], n_id=[169343]) where data.n_id = torch.arange(data.num_nodes), I am wondering how to obtain the k-hop neighbors, either all or a subset, of specific nodes identified in a list or tensor by their indices and return their attributes like features, labels, masks and original indices. For example, how to get the 2 hops neighbors of nodes indexed by [1, 5000, 10000, 1024, 12345] and return their original indices in the graph? Can torch_geometric.utils.k_hop_subgraph achieve that? Any help would be appreciated!

[rusty1s]

The most efficient way to do this would be via NeighborLoader, e.g.:

loader = NeighborLoader(data, num_neighbors=[-1, -1])
loader(torch.tensor([1, 5000, 10000, 1024, 12345]))

k_hop_subgraph does the trick as well but it should be a bit slower to compute.