Class EdgeData

python/dgl/data/csv_dataset_base.py:194–268

Class of edge data which is used for DGLGraph construction. Internal use only.

class EdgeData(BaseData):
    """Class of edge data which is used for DGLGraph construction. Internal use only."""

    def __init__(self, src_id, dst_id, data, type=None, graph_id=None):
        self.src = np.array(src_id)
        self.dst = np.array(dst_id)
        self.data = data
        self.type = type if type is not None else ("_V", "_E", "_V")
        self.graph_id = (
            np.array(graph_id)
            if graph_id is not None
            else np.full(len(src_id), 0)
        )
        _validate_data_length(
            {
                **{"src": self.src, "dst": self.dst, "graph_id": self.graph_id},
                **self.data,
            }
        )

    @staticmethod
    def load_from_csv(
        meta: MetaEdge, data_parser: Callable, base_dir=None, separator=","
    ):
        df = BaseData.read_csv(meta.file_name, base_dir, separator)
        src_ids = BaseData.pop_from_dataframe(df, meta.src_id_field)
        if src_ids is None:
            raise DGLError(
                "Missing src id field [{}] in file [{}].".format(
                    meta.src_id_field, meta.file_name
                )
            )
        dst_ids = BaseData.pop_from_dataframe(df, meta.dst_id_field)
        if dst_ids is None:
            raise DGLError(
                "Missing dst id field [{}] in file [{}].".format(
                    meta.dst_id_field, meta.file_name
                )
            )
        graph_ids = BaseData.pop_from_dataframe(df, meta.graph_id_field)
        etype = tuple(meta.etype)
        edata = data_parser(df)
        return EdgeData(src_ids, dst_ids, edata, type=etype, graph_id=graph_ids)

    @staticmethod
    def to_dict(edge_data: List["EdgeData"], node_dict: dict) -> dict:
        edge_dict = {}
        for e_data in edge_data:
            (src_type, e_type, dst_type) = e_data.type
            graph_ids = np.unique(e_data.graph_id)
            for graph_id in graph_ids:
                if graph_id in edge_dict and e_data.type in edge_dict[graph_id]:
                    raise DGLError(
                        f"Duplicate edge type[{e_data.type}] for same graph[{graph_id}], please place the same edge_type for same graph into single EdgeData."
                    )
                idx = e_data.graph_id == graph_id
                src_mapping = node_dict[graph_id][src_type]["mapping"]
                dst_mapping = node_dict[graph_id][dst_type]["mapping"]