| 17 | |
| 18 | |
| 19 | class GsGraphStore(GraphStore): |
| 20 | def __init__(self, config) -> None: |
| 21 | super().__init__() |
| 22 | self.config = config |
| 23 | self.edge_attrs: Dict[Tuple[Tuple[str, str, str], str, bool], EdgeAttr] = {} |
| 24 | |
| 25 | assert config is not None |
| 26 | config = json.loads( |
| 27 | base64.b64decode(config.encode("utf-8", errors="ignore")).decode( |
| 28 | "utf-8", errors="ignore" |
| 29 | ) |
| 30 | ) |
| 31 | self.edges = config["edges"] |
| 32 | self.edge_dir = config["edge_dir"] |
| 33 | |
| 34 | assert self.edges is not None |
| 35 | for edge in self.edges: |
| 36 | edge = tuple(edge) |
| 37 | # Only support COO layout |
| 38 | layout = "coo" |
| 39 | new_edge_attr = EdgeAttr(edge, layout, True) |
| 40 | self.edge_attrs[(edge, layout, True)] = new_edge_attr |
| 41 | |
| 42 | @staticmethod |
| 43 | def key(attr: EdgeAttr) -> Tuple: |
| 44 | return (attr.edge_type, attr.layout.value, attr.is_sorted, attr.size) |
| 45 | |
| 46 | def _put_edge_index( |
| 47 | self, |
| 48 | edge_index: EdgeTensorType, |
| 49 | edge_attr: EdgeAttr, |
| 50 | ) -> bool: |
| 51 | r"""To be implemented by :class:`GsFeatureStore`.""" |
| 52 | raise NotImplementedError |
| 53 | |
| 54 | def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[EdgeTensorType]: |
| 55 | r"""Obtains a :class:`EdgeTensorType` from the remote server with :class:`EdgeAttr`. |
| 56 | |
| 57 | Args: |
| 58 | edge_attr(`EdgeAttr`): Uniquely corresponds to a topology of subgraph . |
| 59 | |
| 60 | Returns: |
| 61 | edge_index(`EdgeTensorType`): The edge index tensor, which is a :class:`tuple` of\ |
| 62 | (row indice tensor, column indice tensor) |
| 63 | """ |
| 64 | group_name, layout, _, _ = self.key(edge_attr) |
| 65 | num_servers, _, _, _ = request_server(0, DistServer.get_dataset_meta) |
| 66 | rows = [] |
| 67 | cols = [] |
| 68 | for server_id in range(num_servers): |
| 69 | (row, col) = request_server( |
| 70 | server_id, DistServer.get_edge_index, group_name, layout |
| 71 | ) |
| 72 | rows.append(row) |
| 73 | cols.append(col) |
| 74 | |
| 75 | global_row = torch.cat(rows, dim=0) |
| 76 | global_row, perm = index_sort(global_row, max_value=int(global_row.max()) + 1) |