torch_geometric.data.GraphStore

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

Bases: ABC

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.

  • *args – Arguments passed to EdgeAttr.

  • **kwargs – Keyword arguments passed to EdgeAttr.

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

Synchronously obtains an edge_index tuple from the GraphStore.

Parameters:
  • *args – Arguments passed to EdgeAttr.

  • **kwargs – Keyword arguments passed to EdgeAttr.

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:
  • *args – Arguments passed to EdgeAttr.

  • **kwargs – Keyword arguments passed to EdgeAttr.

abstract get_all_edge_attrs() List[EdgeAttr][source]

Returns all registered edge attributes.

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]

Returns 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]

Returns 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]

Returns 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)