torch_geometric.data.GraphStore

class GraphStore(edge_attr_cls: Optional[Any] = None)[source]

Bases: object

An abstract base class to access edges from a remote graph store.

Parameters

edge_attr_cls (EdgeAttr, optional) – A user-defined EdgeAttr class to customize the required attributes and their ordering to uniquely identify edges. (default: None)

put_edge_index(edge_index: Tuple[Tensor, Tensor], *args, **kwargs) → bool[source]

Synchronously adds an edge_index tuple to the GraphStore. Returns whether insertion was successful.

Parameters

  • edge_index (Tuple[torch.Tensor, torch.Tensor]) – The edge_index tuple in a format specified in EdgeAttr.

  • **kwargs (EdgeAttr) – Any relevant edge attributes that correspond to the edge_index tuple. See the EdgeAttr documentation for required and optional attributes.

get_edge_index(*args, **kwargs) → Tuple[Tensor, Tensor][source]

Synchronously obtains an edge_index tuple from the GraphStore.

Parameters

**kwargs (EdgeAttr) – Any relevant edge attributes that correspond to the edge_index tuple. See the EdgeAttr documentation for required and optional attributes.

Raises

KeyError – If the edge_index corresponding to the input EdgeAttr was not found.

remove_edge_index(*args, **kwargs) → bool[source]

Synchronously deletes an edge_index tuple from the GraphStore. Returns whether deletion was successful.

Parameters

**kwargs (EdgeAttr) – Any relevant edge attributes that correspond to the edge_index tuple. See the EdgeAttr documentation for required and optional attributes.

abstract get_all_edge_attrs() → List[EdgeAttr][source]

Obtains all edge attributes stored in the GraphStore.

coo(edge_types: Optional[List[Any]] = None, store: bool = False) → Tuple[Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Optional[Tensor]]][source]

Obtains the edge indices in the GraphStore in COO format.

Parameters

  • edge_types (List[Any], optional) – The edge types of edge indices to obtain. If set to None, will return the edge indices of all existing edge types. (default: None)

  • store (bool, optional) – Whether to store converted edge indices in the GraphStore. (default: False)

csr(edge_types: Optional[List[Any]] = None, store: bool = False) → Tuple[Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Optional[Tensor]]][source]

Obtains the edge indices in the GraphStore in CSR format.

Parameters

  • edge_types (List[Any], optional) – The edge types of edge indices to obtain. If set to None, will return the edge indices of all existing edge types. (default: None)

  • store (bool, optional) – Whether to store converted edge indices in the GraphStore. (default: False)

csc(edge_types: Optional[List[Any]] = None, store: bool = False) → Tuple[Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Optional[Tensor]]][source]

Obtains the edge indices in the GraphStore in CSC format.

Parameters

  • edge_types (List[Any], optional) – The edge types of edge indices to obtain. If set to None, will return the edge indices of all existing edge types. (default: None)

  • store (bool, optional) – Whether to store converted edge indices in the GraphStore. (default: False)