[FEA] Allow sampling per edge type in Uniform Neighbor Sample

[VibhuJawa]

trafficstars

Describe the solution you'd like and any additional context

We should allow sampling per edge type in uniform neighbor sample. After PR 2660 we have graph with type values.

I believe we should include an argument like do_fanout_per_type and sample fan_out number of edges for each seed node for every edge type.

Context:

DGL Behaviour: DGL when sampling neighbors on Heterogeneous Graphs uses type information and by default samples fanout value per type . See below example.

import dgl
import torch as th

graph_data = {
        ('drug', 'interacts', 'drug'): (th.tensor([0,1,2]), th.tensor([3,4,5])),
        ('drug', 'treats', 'disease'): (th.tensor([0,1,2]), th.tensor([6,7,8])),
        ('gene', 'treats', 'disease'): (th.tensor([10,11,12]), th.tensor([6,7,8]))
        }
g = dgl.heterograph(graph_data)

frontier_0 = g.sample_neighbors(nodes={'drug':th.tensor([0]) }, fanout=1, edge_dir='out')
for etype in frontier_0.canonical_etypes:
    print(etype,frontier_0.edges(etype=etype))

('drug', 'interacts', 'drug') (tensor([0]), tensor([3]))
('drug', 'treats', 'disease') (tensor([0]), tensor([6]))
('gene', 'treats', 'disease') (tensor([], dtype=torch.int64), tensor([], dtype=torch.int64))

From DGL API docs (link)

fanout (int or dict[etype, int]) –
The number of edges to be sampled for each node on each edge type.

This argument can take a single int or a dictionary of edge types and ints.
If a single int is given, DGL will sample this number of edges "for each node for every edge type".

If -1 is given for a single edge type, all the neighboring edges with that edge type will be selected.

cuGraph current behavior:

import cudf
import cugraph
df = cudf.DataFrame({'source':[0,1,2,0,1,2,10,11,12],
                     'destination':[3,4,5,6,7,8,6,7,8],
                     'type': [0,0,0,1,1,1,2,2,2]
                     # 0- ('drug', 'interacts', 'drug')
                     # 1-('drug', 'treats', 'disease')
                     # 2-('gene', 'treats', 'disease')
                    })


G = cugraph.MultiGraph(directed=True)
G.from_cudf_edgelist(df, edge_attr='type')

sample_df = cugraph.uniform_neighbor_sample(
    G,
    cudf.Series([0]),
    fanout_vals=[1],
    is_edge_ids=False
)
sample_df = sample_df.rename(columns={'indices':'type'})
 # 0- ('drug', 'interacts', 'drug')
 # 1-('drug', 'treats', 'disease')
 # 2-('gene', 'treats', 'disease')
print(sample_df)

   type  sources  destinations
0     0        0             3

cuGraph requested feature

After PR 2660 we have graph with type values. https://github.com/rapidsai/cugraph/pull/2660, we incude an argument like do_fanout_per_type and sample fanoout number of edges for each seed node for every edge type.

sample_df = cugraph.uniform_neighbor_sample(
    G,
    cudf.Series([0]),
    fanout_vals=[1],
    do_fanout_per_type = True
    is_edge_ids=False
)

   type  sources  destinations
0     0        0             3
1     1        0             6

Alternate work-around

We store n_etypes views of the data and do the sampling call for each edge type individually which will be slow because we will make n_etypes calls.

CC: @ChuckHastings , @seunghwak , @alexbarghi-nv

CC: @rlratzel, @BradReesWork