pytest_utils.py

import json
import logging
import os

import dgl
import numpy as np
import torch
from distpartitioning import array_readwriter
from distpartitioning.array_readwriter.parquet import ParquetArrayParser
from files import setdir


def _chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt, vector_rows=False):
    paths = []
    offset = 0

    for j, n in enumerate(chunk_sizes):
        path = os.path.abspath(path_fmt % j)
        arr_chunk = arr[offset : offset + n]
        shape = arr_chunk.shape
        logging.info("Chunking %d-%d" % (offset, offset + n))
        # If requested we write multi-column arrays as single-column vector Parquet files
        array_parser = array_readwriter.get_array_parser(**fmt_meta)
        if (
            isinstance(array_parser, ParquetArrayParser)
            and len(shape) > 1
            and shape[1] > 1
        ):
            array_parser.write(path, arr_chunk, vector_rows=vector_rows)
        else:
            array_parser.write(path, arr_chunk)
        offset += n
        paths.append(path)

    return paths


def _initialize_num_chunks(g, num_chunks, kwargs=None):
    """Initialize num_chunks for each node/edge.

    Parameters
    ----------
    g: DGLGraph
        Graph to be chunked.
    num_chunks: int
        Default number of chunks to be applied onto node/edge data.
    kwargs: dict
        Key word arguments to specify details for each node/edge data.

    Returns
    -------
    num_chunks_data: dict
        Detailed number of chunks for each node/edge.
    """

    def _init(g, num_chunks, key, kwargs=None):
        chunks_data = kwargs.get(key, None)
        is_node = "_node" in key
        data_types = g.ntypes if is_node else g.canonical_etypes
        if isinstance(chunks_data, int):
            chunks_data = {data_type: chunks_data for data_type in data_types}
        elif isinstance(chunks_data, dict):
            for data_type in data_types:
                if data_type not in chunks_data:
                    chunks_data[data_type] = num_chunks
        else:
            chunks_data = {data_type: num_chunks for data_type in data_types}
        for _, data in chunks_data.items():
            if isinstance(data, dict):
                n_chunks = list(data.values())
            else:
                n_chunks = [data]
            assert all(
                isinstance(v, int) for v in n_chunks
            ), "num_chunks for each data type should be int."
        return chunks_data

    num_chunks_data = {}
    for key in [
        "num_chunks_nodes",
        "num_chunks_edges",
        "num_chunks_node_data",
        "num_chunks_edge_data",
    ]:
        num_chunks_data[key] = _init(g, num_chunks, key, kwargs=kwargs)
    return num_chunks_data


def _chunk_graph(
    g,
    name,
    ndata_paths,
    edata_paths,
    num_chunks,
    data_fmt,
    edges_format,
    vector_rows=False,
    **kwargs,
):
    # First deal with ndata and edata that are homogeneous
    # (i.e. not a dict-of-dict)
    if len(g.ntypes) == 1 and not isinstance(
        next(iter(ndata_paths.values())), dict
    ):
        ndata_paths = {g.ntypes[0]: ndata_paths}
    if len(g.etypes) == 1 and not isinstance(
        next(iter(edata_paths.values())), dict
    ):
        edata_paths = {g.etypes[0]: ndata_paths}
    # Then convert all edge types to canonical edge types
    etypestrs = {etype: ":".join(etype) for etype in g.canonical_etypes}
    edata_paths = {
        ":".join(g.to_canonical_etype(k)): v for k, v in edata_paths.items()
    }

    metadata = {}

    metadata["graph_name"] = name
    metadata["node_type"] = g.ntypes

    # add node_type_counts
    metadata["num_nodes_per_type"] = [g.num_nodes(ntype) for ntype in g.ntypes]

    # Initialize num_chunks for each node/edge.
    num_chunks_details = _initialize_num_chunks(g, num_chunks, kwargs=kwargs)

    # Compute the number of nodes per chunk per node type
    metadata["num_nodes_per_chunk"] = num_nodes_per_chunk = []
    num_chunks_nodes = num_chunks_details["num_chunks_nodes"]
    for ntype in g.ntypes:
        num_nodes = g.num_nodes(ntype)
        num_nodes_list = []
        n_chunks = num_chunks_nodes[ntype]
        for i in range(n_chunks):
            n = num_nodes // n_chunks + (i < num_nodes % n_chunks)
            num_nodes_list.append(n)
        num_nodes_per_chunk.append(num_nodes_list)

    metadata["edge_type"] = [etypestrs[etype] for etype in g.canonical_etypes]
    metadata["num_edges_per_type"] = [
        g.num_edges(etype) for etype in g.canonical_etypes
    ]

    # Compute the number of edges per chunk per edge type
    metadata["num_edges_per_chunk"] = num_edges_per_chunk = []
    num_chunks_edges = num_chunks_details["num_chunks_edges"]
    for etype in g.canonical_etypes:
        num_edges = g.num_edges(etype)
        num_edges_list = []
        n_chunks = num_chunks_edges[etype]
        for i in range(n_chunks):
            n = num_edges // n_chunks + (i < num_edges % n_chunks)
            num_edges_list.append(n)
        num_edges_per_chunk.append(num_edges_list)
    num_edges_per_chunk_dict = {
        k: v for k, v in zip(g.canonical_etypes, num_edges_per_chunk)
    }

    idxes_etypestr = {
        idx: (etype, etypestrs[etype])
        for idx, etype in enumerate(g.canonical_etypes)
    }
    idxes = np.arange(len(idxes_etypestr))

    # Split edge index
    metadata["edges"] = {}
    with setdir("edge_index"):
        np.random.shuffle(idxes)
        for idx in idxes:
            etype = idxes_etypestr[idx][0]
            etypestr = idxes_etypestr[idx][1]
            logging.info("Chunking edge index for %s" % etypestr)
            edges_meta = {}
            if edges_format == "csv":
                fmt_meta = {"name": edges_format, "delimiter": " "}
            elif edges_format == "parquet":
                fmt_meta = {"name": edges_format}
            else:
                raise RuntimeError(f"Invalid edges_fmt: {edges_format}")
            edges_meta["format"] = fmt_meta

            srcdst = torch.stack(g.edges(etype=etype), 1)
            edges_meta["data"] = _chunk_numpy_array(
                srcdst.numpy(),
                fmt_meta,
                num_edges_per_chunk_dict[etype],
                etypestr + "%d.txt",
            )
            metadata["edges"][etypestr] = edges_meta

    # Chunk node data
    reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt}
    file_suffix = "npy" if data_fmt == "numpy" else "parquet"
    metadata["node_data"] = {}
    num_chunks_node_data = num_chunks_details["num_chunks_node_data"]
    with setdir("node_data"):
        for ntype, ndata_per_type in ndata_paths.items():
            ndata_meta = {}
            with setdir(ntype):
                for key, path in ndata_per_type.items():
                    logging.info(
                        "Chunking node data for type %s key %s" % (ntype, key)
                    )
                    chunk_sizes = []
                    num_nodes = g.num_nodes(ntype)
                    n_chunks = num_chunks_node_data[ntype]
                    if isinstance(n_chunks, dict):
                        n_chunks = n_chunks.get(key, num_chunks)
                    assert isinstance(n_chunks, int), (
                        f"num_chunks for {ntype}/{key} should be int while "
                        f"{type(n_chunks)} is got."
                    )
                    for i in range(n_chunks):
                        n = num_nodes // n_chunks + (i < num_nodes % n_chunks)
                        chunk_sizes.append(n)
                    ndata_key_meta = {}
                    arr = array_readwriter.get_array_parser(
                        **reader_fmt_meta
                    ).read(path)
                    ndata_key_meta["format"] = writer_fmt_meta
                    ndata_key_meta["data"] = _chunk_numpy_array(
                        arr,
                        writer_fmt_meta,
                        chunk_sizes,
                        key + "-%d." + file_suffix,
                        vector_rows=vector_rows,
                    )
                    ndata_meta[key] = ndata_key_meta

            metadata["node_data"][ntype] = ndata_meta

    # Chunk edge data
    metadata["edge_data"] = {}
    num_chunks_edge_data = num_chunks_details["num_chunks_edge_data"]
    with setdir("edge_data"):
        for etypestr, edata_per_type in edata_paths.items():
            edata_meta = {}
            etype = tuple(etypestr.split(":"))
            with setdir(etypestr):
                for key, path in edata_per_type.items():
                    logging.info(
                        "Chunking edge data for type %s key %s"
                        % (etypestr, key)
                    )
                    chunk_sizes = []
                    num_edges = g.num_edges(etype)
                    n_chunks = num_chunks_edge_data[etype]
                    if isinstance(n_chunks, dict):
                        n_chunks = n_chunks.get(key, num_chunks)
                    assert isinstance(n_chunks, int), (
                        f"num_chunks for {etype}/{key} should be int while "
                        f"{type(n_chunks)} is got."
                    )
                    for i in range(n_chunks):
                        n = num_edges // n_chunks + (i < num_edges % n_chunks)
                        chunk_sizes.append(n)
                    edata_key_meta = {}
                    arr = array_readwriter.get_array_parser(
                        **reader_fmt_meta
                    ).read(path)
                    edata_key_meta["format"] = writer_fmt_meta
                    edata_key_meta["data"] = _chunk_numpy_array(
                        arr,
                        writer_fmt_meta,
                        chunk_sizes,
                        key + "-%d." + file_suffix,
                        vector_rows=vector_rows,
                    )
                    edata_meta[key] = edata_key_meta

            metadata["edge_data"][etypestr] = edata_meta

    metadata_path = "metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, sort_keys=True, indent=4)
    logging.info("Saved metadata in %s" % os.path.abspath(metadata_path))


def chunk_graph(
    g,
    name,
    ndata_paths,
    edata_paths,
    num_chunks,
    output_path,
    data_fmt="numpy",
    edges_fmt="csv",
    vector_rows=False,
    **kwargs,
):
    """
    Split the graph into multiple chunks.

    A directory will be created at :attr:`output_path` with the metadata and
    chunked edge list as well as the node/edge data.

    Parameters
    ----------
    g : DGLGraph
        The graph.
    name : str
        The name of the graph, to be used later in DistDGL training.
    ndata_paths : dict[str, pathlike] or dict[ntype, dict[str, pathlike]]
        The dictionary of paths pointing to the corresponding numpy array file
        for each node data key.
    edata_paths : dict[etype, pathlike] or dict[etype, dict[str, pathlike]]
        The dictionary of paths pointing to the corresponding numpy array file
        for each edge data key. ``etype`` could be canonical or non-canonical.
    num_chunks : int
        The number of chunks
    output_path : pathlike
        The output directory saving the chunked graph.
    data_fmt : str
        Format of node/edge data: 'numpy' or 'parquet'.
    edges_fmt : str
        Format of edges files: 'csv' or 'parquet'.
    vector_rows : str
        When true will write parquet files as single-column vector row files.
    kwargs : dict
        Key word arguments to control chunk details.
    """
    for ntype, ndata in ndata_paths.items():
        for key in ndata.keys():
            ndata[key] = os.path.abspath(ndata[key])
    for etype, edata in edata_paths.items():
        for key in edata.keys():
            edata[key] = os.path.abspath(edata[key])
    with setdir(output_path):
        _chunk_graph(
            g,
            name,
            ndata_paths,
            edata_paths,
            num_chunks,
            data_fmt,
            edges_fmt,
            vector_rows,
            **kwargs,
        )


def create_homo_chunked_dataset(
    root_dir,
    num_chunks,
    data_fmt="numpy",
    edges_fmt="csv",
    vector_rows=False,
    **kwargs,
):
    """
    This function creates a sample homo dataset.

    Parameters:
    -----------
    root_dir : string
        directory in which all the files for the chunked dataset will be stored.
    """
    # Step0: prepare chunked graph data format.
    # A synthetic mini MAG240.
    num_N = 1200

    def rand_edges(num_src, num_dst, num_edges):
        eids = np.random.choice(num_src * num_dst, num_edges, replace=False)
        src = torch.from_numpy(eids // num_dst)
        dst = torch.from_numpy(eids % num_dst)

        return src, dst

    num_E = 24 * 1000

    # Structure.
    data_dict = {("_N", "_E", "_N"): rand_edges(num_N, num_N, num_E)}
    src, dst = data_dict[("_N", "_E", "_N")]
    data_dict[("_N", "_E", "_N")] = (dst, src)
    g = dgl.heterograph(data_dict)

    # paper feat, label, year
    num_paper_feats = 3
    _N_feat = np.random.randn(num_N, num_paper_feats)
    num_classes = 4
    _N_label = np.random.choice(num_classes, num_N)
    _N_year = np.random.choice(2022, num_N)
    _N_orig_ids = np.arange(0, num_N)

    # masks.
    _N_train_mask = np.random.choice([True, False], num_N)
    _N_test_mask = np.random.choice([True, False], num_N)
    _N_val_mask = np.random.choice([True, False], num_N)

    # Edge features.
    _E_count = np.random.choice(10, num_E)

    # Save features.
    input_dir = os.path.join(root_dir, "data_test")
    os.makedirs(input_dir)
    for sub_d in ["_N", "_E"]:
        os.makedirs(os.path.join(input_dir, sub_d))

    _N_feat_path = os.path.join(input_dir, "_N/feat.npy")
    with open(_N_feat_path, "wb") as f:
        np.save(f, _N_feat)
    g.nodes["_N"].data["feat"] = torch.from_numpy(_N_feat)

    _N_label_path = os.path.join(input_dir, "_N/label.npy")
    with open(_N_label_path, "wb") as f:
        np.save(f, _N_label)
    g.nodes["_N"].data["label"] = torch.from_numpy(_N_label)

    _N_year_path = os.path.join(input_dir, "_N/year.npy")
    with open(_N_year_path, "wb") as f:
        np.save(f, _N_year)
    g.nodes["_N"].data["year"] = torch.from_numpy(_N_year)

    _N_orig_ids_path = os.path.join(input_dir, "_N/orig_ids.npy")
    with open(_N_orig_ids_path, "wb") as f:
        np.save(f, _N_orig_ids)
    g.nodes["_N"].data["orig_ids"] = torch.from_numpy(_N_orig_ids)

    _E_count_path = os.path.join(input_dir, "_E/count.npy")
    with open(_E_count_path, "wb") as f:
        np.save(f, _E_count)
    g.edges["_E"].data["count"] = torch.from_numpy(_E_count)

    _N_train_mask_path = os.path.join(input_dir, "_N/train_mask.npy")
    with open(_N_train_mask_path, "wb") as f:
        np.save(f, _N_train_mask)
    g.nodes["_N"].data["train_mask"] = torch.from_numpy(_N_train_mask)

    _N_test_mask_path = os.path.join(input_dir, "_N/test_mask.npy")
    with open(_N_test_mask_path, "wb") as f:
        np.save(f, _N_test_mask)
    g.nodes["_N"].data["test_mask"] = torch.from_numpy(_N_test_mask)

    _N_val_mask_path = os.path.join(input_dir, "_N/val_mask.npy")
    with open(_N_val_mask_path, "wb") as f:
        np.save(f, _N_val_mask)
    g.nodes["_N"].data["val_mask"] = torch.from_numpy(_N_val_mask)

    node_data = {
        "_N": {
            "feat": _N_feat_path,
            "train_mask": _N_train_mask_path,
            "test_mask": _N_test_mask_path,
            "val_mask": _N_val_mask_path,
            "label": _N_label_path,
            "year": _N_year_path,
            "orig_ids": _N_orig_ids_path,
        }
    }

    edge_data = {"_E": {"count": _E_count_path}}

    output_dir = os.path.join(root_dir, "chunked-data")
    chunk_graph(
        g,
        "mag240m",
        node_data,
        edge_data,
        num_chunks=num_chunks,
        output_path=output_dir,
        data_fmt=data_fmt,
        edges_fmt=edges_fmt,
        vector_rows=vector_rows,
        **kwargs,
    )
    logging.debug("Done with creating chunked graph")

    return g


def create_hetero_chunked_dataset(
    root_dir,
    num_chunks,
    data_fmt="numpy",
    edges_fmt="csv",
    vector_rows=False,
    **kwargs,
):
    """
    This function creates a sample dataset, based on MAG240 dataset.

    Parameters:
    -----------
    root_dir : string
        directory in which all the files for the chunked dataset will be stored.
    """
    # Step0: prepare chunked graph data format.
    # A synthetic mini MAG240.
    num_institutions = 1200
    num_authors = 1200
    num_papers = 1200

    def rand_edges(num_src, num_dst, num_edges):
        eids = np.random.choice(num_src * num_dst, num_edges, replace=False)
        src = torch.from_numpy(eids // num_dst)
        dst = torch.from_numpy(eids % num_dst)

        return src, dst

    num_cite_edges = 24 * 1000
    num_write_edges = 12 * 1000
    num_affiliate_edges = 2400

    # Structure.
    data_dict = {
        ("paper", "cites", "paper"): rand_edges(
            num_papers, num_papers, num_cite_edges
        ),
        ("author", "writes", "paper"): rand_edges(
            num_authors, num_papers, num_write_edges
        ),
        ("author", "affiliated_with", "institution"): rand_edges(
            num_authors, num_institutions, num_affiliate_edges
        ),
        ("institution", "writes", "paper"): rand_edges(
            num_institutions, num_papers, num_write_edges
        ),
    }
    src, dst = data_dict[("author", "writes", "paper")]
    data_dict[("paper", "rev_writes", "author")] = (dst, src)
    g = dgl.heterograph(data_dict)

    # paper feat, label, year
    num_paper_feats = 3
    paper_feat = np.random.randn(num_papers, num_paper_feats)
    num_classes = 4
    paper_label = np.random.choice(num_classes, num_papers)
    paper_year = np.random.choice(2022, num_papers)
    paper_orig_ids = np.arange(0, num_papers)
    writes_orig_ids = np.arange(0, num_write_edges)

    # masks.
    paper_train_mask = np.random.choice([True, False], num_papers)
    paper_test_mask = np.random.choice([True, False], num_papers)
    paper_val_mask = np.random.choice([True, False], num_papers)

    author_train_mask = np.random.choice([True, False], num_authors)
    author_test_mask = np.random.choice([True, False], num_authors)
    author_val_mask = np.random.choice([True, False], num_authors)

    inst_train_mask = np.random.choice([True, False], num_institutions)
    inst_test_mask = np.random.choice([True, False], num_institutions)
    inst_val_mask = np.random.choice([True, False], num_institutions)

    write_train_mask = np.random.choice([True, False], num_write_edges)
    write_test_mask = np.random.choice([True, False], num_write_edges)
    write_val_mask = np.random.choice([True, False], num_write_edges)

    # Edge features.
    cite_count = np.random.choice(10, num_cite_edges)
    write_year = np.random.choice(2022, num_write_edges)
    write2_year = np.random.choice(2022, num_write_edges)

    # Save features.
    input_dir = os.path.join(root_dir, "data_test")
    os.makedirs(input_dir)
    for sub_d in ["paper", "cites", "writes", "writes2"]:
        os.makedirs(os.path.join(input_dir, sub_d))

    paper_feat_path = os.path.join(input_dir, "paper/feat.npy")
    with open(paper_feat_path, "wb") as f:
        np.save(f, paper_feat)
    g.nodes["paper"].data["feat"] = torch.from_numpy(paper_feat)

    paper_label_path = os.path.join(input_dir, "paper/label.npy")
    with open(paper_label_path, "wb") as f:
        np.save(f, paper_label)
    g.nodes["paper"].data["label"] = torch.from_numpy(paper_label)

    paper_year_path = os.path.join(input_dir, "paper/year.npy")
    with open(paper_year_path, "wb") as f:
        np.save(f, paper_year)
    g.nodes["paper"].data["year"] = torch.from_numpy(paper_year)

    paper_orig_ids_path = os.path.join(input_dir, "paper/orig_ids.npy")
    with open(paper_orig_ids_path, "wb") as f:
        np.save(f, paper_orig_ids)
    g.nodes["paper"].data["orig_ids"] = torch.from_numpy(paper_orig_ids)

    cite_count_path = os.path.join(input_dir, "cites/count.npy")
    with open(cite_count_path, "wb") as f:
        np.save(f, cite_count)
    g.edges["cites"].data["count"] = torch.from_numpy(cite_count)

    write_year_path = os.path.join(input_dir, "writes/year.npy")
    with open(write_year_path, "wb") as f:
        np.save(f, write_year)
    g.edges[("author", "writes", "paper")].data["year"] = torch.from_numpy(
        write_year
    )
    g.edges["rev_writes"].data["year"] = torch.from_numpy(write_year)

    writes_orig_ids_path = os.path.join(input_dir, "writes/orig_ids.npy")
    with open(writes_orig_ids_path, "wb") as f:
        np.save(f, writes_orig_ids)
    g.edges[("author", "writes", "paper")].data["orig_ids"] = torch.from_numpy(
        writes_orig_ids
    )

    write2_year_path = os.path.join(input_dir, "writes2/year.npy")
    with open(write2_year_path, "wb") as f:
        np.save(f, write2_year)
    g.edges[("institution", "writes", "paper")].data["year"] = torch.from_numpy(
        write2_year
    )

    etype = ("author", "writes", "paper")
    write_train_mask_path = os.path.join(input_dir, "writes/train_mask.npy")
    with open(write_train_mask_path, "wb") as f:
        np.save(f, write_train_mask)
    g.edges[etype].data["train_mask"] = torch.from_numpy(write_train_mask)

    write_test_mask_path = os.path.join(input_dir, "writes/test_mask.npy")
    with open(write_test_mask_path, "wb") as f:
        np.save(f, write_test_mask)
    g.edges[etype].data["test_mask"] = torch.from_numpy(write_test_mask)

    write_val_mask_path = os.path.join(input_dir, "writes/val_mask.npy")
    with open(write_val_mask_path, "wb") as f:
        np.save(f, write_val_mask)
    g.edges[etype].data["val_mask"] = torch.from_numpy(write_val_mask)

    for sub_d in ["author", "institution"]:
        os.makedirs(os.path.join(input_dir, sub_d))
    paper_train_mask_path = os.path.join(input_dir, "paper/train_mask.npy")
    with open(paper_train_mask_path, "wb") as f:
        np.save(f, paper_train_mask)
    g.nodes["paper"].data["train_mask"] = torch.from_numpy(paper_train_mask)

    paper_test_mask_path = os.path.join(input_dir, "paper/test_mask.npy")
    with open(paper_test_mask_path, "wb") as f:
        np.save(f, paper_test_mask)
    g.nodes["paper"].data["test_mask"] = torch.from_numpy(paper_test_mask)

    paper_val_mask_path = os.path.join(input_dir, "paper/val_mask.npy")
    with open(paper_val_mask_path, "wb") as f:
        np.save(f, paper_val_mask)
    g.nodes["paper"].data["val_mask"] = torch.from_numpy(paper_val_mask)

    author_train_mask_path = os.path.join(input_dir, "author/train_mask.npy")
    with open(author_train_mask_path, "wb") as f:
        np.save(f, author_train_mask)
    g.nodes["author"].data["train_mask"] = torch.from_numpy(author_train_mask)

    author_test_mask_path = os.path.join(input_dir, "author/test_mask.npy")
    with open(author_test_mask_path, "wb") as f:
        np.save(f, author_test_mask)
    g.nodes["author"].data["test_mask"] = torch.from_numpy(author_test_mask)

    author_val_mask_path = os.path.join(input_dir, "author/val_mask.npy")
    with open(author_val_mask_path, "wb") as f:
        np.save(f, author_val_mask)
    g.nodes["author"].data["val_mask"] = torch.from_numpy(author_val_mask)

    inst_train_mask_path = os.path.join(input_dir, "institution/train_mask.npy")
    with open(inst_train_mask_path, "wb") as f:
        np.save(f, inst_train_mask)
    g.nodes["institution"].data["train_mask"] = torch.from_numpy(
        inst_train_mask
    )

    inst_test_mask_path = os.path.join(input_dir, "institution/test_mask.npy")
    with open(inst_test_mask_path, "wb") as f:
        np.save(f, inst_test_mask)
    g.nodes["institution"].data["test_mask"] = torch.from_numpy(inst_test_mask)

    inst_val_mask_path = os.path.join(input_dir, "institution/val_mask.npy")
    with open(inst_val_mask_path, "wb") as f:
        np.save(f, inst_val_mask)
    g.nodes["institution"].data["val_mask"] = torch.from_numpy(inst_val_mask)

    node_data = {
        "paper": {
            "feat": paper_feat_path,
            "train_mask": paper_train_mask_path,
            "test_mask": paper_test_mask_path,
            "val_mask": paper_val_mask_path,
            "label": paper_label_path,
            "year": paper_year_path,
            "orig_ids": paper_orig_ids_path,
        },
        "author": {
            "train_mask": author_train_mask_path,
            "test_mask": author_test_mask_path,
            "val_mask": author_val_mask_path,
        },
        "institution": {
            "train_mask": inst_train_mask_path,
            "test_mask": inst_test_mask_path,
            "val_mask": inst_val_mask_path,
        },
    }

    edge_data = {
        "cites": {"count": cite_count_path},
        ("author", "writes", "paper"): {
            "year": write_year_path,
            "orig_ids": writes_orig_ids_path,
            "train_mask": write_train_mask_path,
            "test_mask": write_test_mask_path,
            "val_mask": write_val_mask_path,
        },
        "rev_writes": {"year": write_year_path},
        ("institution", "writes", "paper"): {"year": write2_year_path},
    }

    output_dir = os.path.join(root_dir, "chunked-data")
    chunk_graph(
        g,
        "mag240m",
        node_data,
        edge_data,
        num_chunks=num_chunks,
        output_path=output_dir,
        data_fmt=data_fmt,
        edges_fmt=edges_fmt,
        vector_rows=vector_rows,
        **kwargs,
    )
    logging.debug("Done with creating chunked graph")

    return g