cugraph_pyg.data.dask_graph_store.DaskGraphStore#

class cugraph_pyg.data.dask_graph_store.DaskGraphStore(F: FeatureStore, G: Union[Dict[Tuple[str, str, str], Tuple[Union[Tensor, ndarray, Series, Series]]], Dict[Tuple[str, str, str], int]], num_nodes_dict: Dict[str, int], *, multi_gpu: bool = False, order: str = 'CSR')[source]#

Duck-typed version of PyG’s GraphStore and FeatureStore.

Attributes:
edge_types
node_types
order

Methods

create_named_tensor(attr_name, properties, ...)

Create a named tensor that contains a subset of properties in the graph.

filter(format, node_dict, row_dict, ...)

Parameters:

get_all_edge_attrs()

Gets a list of all edge types and indices in this store.

get_all_tensor_attrs()

Obtains all tensor attributes stored in this feature store.

get_edge_index(*args, **kwargs)

Synchronously gets an edge_index tensor from the materialized graph.

get_tensor(*args, **kwargs)

Synchronously obtains a FeatureTensorType object from the feature store.

get_tensor_size(*args, **kwargs)

Obtains the size of a tensor given its attributes, or None if the tensor does not exist.

multi_get_tensor(attrs)

Synchronously obtains a FeatureTensorType object from the feature store for each tensor associated with the attributes in attrs.

put_edge_index(edge_index, edge_attr)

Adds additional edges to the graph.

canonical_edge_type_to_numeric

get_vertex_index

numeric_edge_type_to_canonical

put_tensor

__init__(F: FeatureStore, G: Union[Dict[Tuple[str, str, str], Tuple[Union[Tensor, ndarray, Series, Series]]], Dict[Tuple[str, str, str], int]], num_nodes_dict: Dict[str, int], *, multi_gpu: bool = False, order: str = 'CSR')[source]#

Constructs a new DaskGraphStore from the provided arguments.

Parameters:
F: cugraph.gnn.FeatureStore (Required)

The feature store containing this graph’s features. Typed lexicographic-ordered numbering convention should match that of the graph.

G: dict[str, tuple[TensorType]] or dict[str, int] (Required)

Dictionary of edge indices. Option 1 (graph in memory):

Pass the edge indices: i.e. { (‘author’, ‘writes’, ‘paper’): [[0,1,2],[2,0,1]], (‘author’, ‘affiliated’, ‘institution’): [[0,1],[0,1]] }

Option 2 (graph not in memory):

Pass the number of edges: i.e. { (‘author’, ‘writes’, ‘paper’): 2, (‘author’, ‘affiliated’, ‘institution’): 2 } If the graph is not in memory, manipulating the edge indices or calling sampling is not possible. This is for cases where sampling has already been done and samples were written to disk.

Note: the internal cugraph representation will use offsetted vertex and edge ids.

num_nodes_dict: dict (Required)

A dictionary mapping each node type to the count of nodes of that type in the graph.

multi_gpu: bool (Optional, default = False)

Whether the store should be backed by a multi-GPU graph. Requires dask to have been set up.

order: str (Optional [“CSR”, “CSC”], default = CSR)

The order to use for sampling. CSR corresponds to the standard OGB dataset order that is usually used in PyG. CSC order constructs the same graph as CSR, but with edges in the opposite direction.

Methods

__init__(F, G, num_nodes_dict, *[, ...])

Constructs a new DaskGraphStore from the provided arguments.

canonical_edge_type_to_numeric(etype)

create_named_tensor(attr_name, properties, ...)

Create a named tensor that contains a subset of properties in the graph.

filter(format, node_dict, row_dict, ...)

Parameters:

get_all_edge_attrs()

Gets a list of all edge types and indices in this store.

get_all_tensor_attrs()

Obtains all tensor attributes stored in this feature store.

get_edge_index(*args, **kwargs)

Synchronously gets an edge_index tensor from the materialized graph.

get_tensor(*args, **kwargs)

Synchronously obtains a FeatureTensorType object from the feature store.

get_tensor_size(*args, **kwargs)

Obtains the size of a tensor given its attributes, or None if the tensor does not exist.

get_vertex_index(vtypes)

multi_get_tensor(attrs)

Synchronously obtains a FeatureTensorType object from the feature store for each tensor associated with the attributes in attrs.

numeric_edge_type_to_canonical(etype)

put_edge_index(edge_index, edge_attr)

Adds additional edges to the graph.

put_tensor(tensor, attr)

Attributes

edge_types

node_types

order