graphlearn_torch/python/partition/base.py

# Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import os
import pickle
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, Union

import torch

from ..typing import (
  NodeType, EdgeType, as_str, TensorDataType,
  GraphPartitionData, HeteroGraphPartitionData,
  FeaturePartitionData, HeteroFeaturePartitionData,
)
from ..utils import convert_to_tensor, ensure_dir, id2idx, append_tensor_to_file, load_and_concatenate_tensors

class PartitionBook(object):
  @abstractmethod
  def __getitem__(self, indices):
    pass

  @property
  def offset(self):
    return 0

HeteroNodePartitionDict = Dict[NodeType, PartitionBook]
HeteroEdgePartitionDict = Dict[EdgeType, PartitionBook]


def save_meta(
  output_dir: str,
  num_parts: int,
  data_cls: str = 'homo',
  node_types: Optional[List[NodeType]] = None,
  edge_types: Optional[List[EdgeType]] = None,
):
  r""" Save partitioning meta info into the output directory.
  """
  meta = {
    'num_parts': num_parts,
    'data_cls': data_cls,
    'node_types': node_types,
    'edge_types': edge_types
  }
  with open(os.path.join(output_dir, 'META'), 'wb') as outfile:
    pickle.dump(meta, outfile, pickle.HIGHEST_PROTOCOL)


def save_node_pb(
  output_dir: str,
  node_pb: PartitionBook,
  ntype: Optional[NodeType] = None
):
  r""" Save a partition book of graph nodes into the output directory.
  """
  if ntype is not None:
    subdir = os.path.join(output_dir, 'node_pb')
    ensure_dir(subdir)
    fpath = os.path.join(subdir, f'{as_str(ntype)}.pt')
  else:
    fpath = os.path.join(output_dir, 'node_pb.pt')
  torch.save(node_pb, fpath)


def save_edge_pb(
  output_dir: str,
  edge_pb: PartitionBook,
  etype: Optional[EdgeType] = None
):
  r""" Save a partition book of graph edges into the output directory.
  """
  if etype is not None:
    subdir = os.path.join(output_dir, 'edge_pb')
    ensure_dir(subdir)
    fpath = os.path.join(subdir, f'{as_str(etype)}.pt')
  else:
    fpath = os.path.join(output_dir, 'edge_pb.pt')
  torch.save(edge_pb, fpath)

def save_graph_cache(
  output_dir: str,
  graph_partition_list: List[GraphPartitionData],
  etype: Optional[EdgeType] = None,
  with_edge_feat: bool = False
):
  r""" Save full graph topology into the output directory.
  """
  if len(graph_partition_list) == 0:
    return
  subdir = os.path.join(output_dir, 'graph')
  if etype is not None:
    subdir = os.path.join(subdir, as_str(etype))
  ensure_dir(subdir)
  rows = torch.cat([graph_partition.edge_index[0] for graph_partition in graph_partition_list])
  cols = torch.cat([graph_partition.edge_index[1] for graph_partition in graph_partition_list])
  weights = None
  if graph_partition_list[0].weights is not None:
    weights = torch.cat([graph_partition.weights for graph_partition in graph_partition_list])
  torch.save(rows, os.path.join(subdir, 'rows.pt'))
  torch.save(cols, os.path.join(subdir, 'cols.pt'))
  if with_edge_feat:
    edge_ids = torch.cat([graph_partition.eids for graph_partition in graph_partition_list])
    torch.save(edge_ids, os.path.join(subdir, 'eids.pt'))
  if weights is not None:
    torch.save(weights, os.path.join(subdir, 'weights.pt'))


def save_graph_partition(
  output_dir: str,
  partition_idx: int,
  graph_partition: GraphPartitionData,
  etype: Optional[EdgeType] = None
):
  r""" Save a graph topology partition into the output directory.
  """
  subdir = os.path.join(output_dir, f'part{partition_idx}', 'graph')
  if etype is not None:
    subdir = os.path.join(subdir, as_str(etype))
  ensure_dir(subdir)
  torch.save(graph_partition.edge_index[0], os.path.join(subdir, 'rows.pt'))
  torch.save(graph_partition.edge_index[1], os.path.join(subdir, 'cols.pt'))
  torch.save(graph_partition.eids, os.path.join(subdir, 'eids.pt'))
  if graph_partition.weights is not None:
    torch.save(graph_partition.weights, os.path.join(subdir, 'weights.pt'))

def save_feature_partition(
  output_dir: str,
  partition_idx: int,
  feature_partition: FeaturePartitionData,
  group: str = 'node_feat',
  graph_type: Optional[Union[NodeType, EdgeType]] = None
):
  r""" Save a feature partition into the output directory.
  """
  subdir = os.path.join(output_dir, f'part{partition_idx}', group)
  if graph_type is not None:
    subdir = os.path.join(subdir, as_str(graph_type))
  ensure_dir(subdir)
  append_tensor_to_file(os.path.join(subdir, 'feats.pkl'), feature_partition.feats)
  append_tensor_to_file(os.path.join(subdir,'ids.pkl'), feature_partition.ids)
  if feature_partition.cache_feats is not None:
    torch.save(feature_partition.cache_feats, os.path.join(subdir, 'cache_feats.pt'))
    torch.save(feature_partition.cache_ids, os.path.join(subdir, 'cache_ids.pt'))

def save_feature_partition_chunk(
  output_dir: str,
  partition_idx: int,
  feature_partition: FeaturePartitionData,
  group: str = 'node_feat',
  graph_type: Optional[Union[NodeType, EdgeType]] = None
):
  r""" Append a chunk of a feature partition to files in the output directory.
  """
  subdir = os.path.join(output_dir, f'part{partition_idx}', group)
  if graph_type is not None:
    subdir = os.path.join(subdir, as_str(graph_type))
  ensure_dir(subdir)
  append_tensor_to_file(os.path.join(subdir, 'feats.pkl'), feature_partition.feats)
  append_tensor_to_file(os.path.join(subdir,'ids.pkl'), feature_partition.ids)

def save_feature_partition_cache(
  output_dir: str,
  partition_idx: int,
  feature_partition: FeaturePartitionData,
  group: str = 'node_feat',
  graph_type: Optional[Union[NodeType, EdgeType]] = None
):
  r""" Save the feature cache of a partition into the output directory.
  """
  subdir = os.path.join(output_dir, f'part{partition_idx}', group)
  if graph_type is not None:
    subdir = os.path.join(subdir, as_str(graph_type))
  ensure_dir(subdir)
  if feature_partition.cache_feats is not None:
    torch.save(feature_partition.cache_feats, os.path.join(subdir, 'cache_feats.pt'))
    torch.save(feature_partition.cache_ids, os.path.join(subdir, 'cache_ids.pt'))


class PartitionerBase(ABC):
  r""" Base class for partitioning graphs and features.
  """
  def __init__(
    self,
    output_dir: str,
    num_parts: int,
    num_nodes: Union[int, Dict[NodeType, int]],
    edge_index: Union[TensorDataType, Dict[EdgeType, TensorDataType]],
    node_feat: Optional[Union[TensorDataType, Dict[NodeType, TensorDataType]]] = None,
    node_feat_dtype: torch.dtype = torch.float32,
    edge_feat: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None,
    edge_feat_dtype: torch.dtype = torch.float32,
    edge_weights: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None,
    edge_assign_strategy: str = 'by_src',
    chunk_size: int = 10000,
  ):
    self.output_dir = output_dir
    ensure_dir(self.output_dir)

    self.num_parts = num_parts
    assert self.num_parts > 1

    self.num_nodes = num_nodes
    self.edge_index = convert_to_tensor(edge_index, dtype=torch.int64)
    self.node_feat = convert_to_tensor(node_feat, dtype=node_feat_dtype)
    self.edge_feat = convert_to_tensor(edge_feat, dtype=edge_feat_dtype)
    self.edge_weights = convert_to_tensor(edge_weights, dtype=torch.float32)

    if isinstance(self.num_nodes, dict):
      assert isinstance(self.edge_index, dict)
      assert isinstance(self.node_feat, dict) or self.node_feat is None
      assert isinstance(self.edge_feat, dict) or self.edge_feat is None
      self.data_cls = 'hetero'
      self.node_types = list(self.num_nodes.keys())
      self.edge_types = list(self.edge_index.keys())
      self.num_edges = {}
      for etype, index in self.edge_index.items():
        self.num_edges[etype] = len(index[0])
    else:
      self.data_cls = 'homo'
      self.node_types = None
      self.edge_types = None
      self.num_edges = len(self.edge_index[0])

    self.edge_assign_strategy = edge_assign_strategy.lower()
    assert self.edge_assign_strategy in ['by_src', 'by_dst']
    self.chunk_size = chunk_size

  def get_edge_index(self, etype: Optional[EdgeType] = None):
    if 'hetero' == self.data_cls:
      assert etype is not None
      return self.edge_index[etype]
    return self.edge_index

  def get_node_feat(self, ntype: Optional[NodeType] = None):
    if self.node_feat is None:
      return None
    if 'hetero' == self.data_cls:
      assert ntype is not None
      return self.node_feat[ntype]
    return self.node_feat

  def get_edge_feat(self, etype: Optional[EdgeType] = None):
    if self.edge_feat is None:
      return None
    if 'hetero' == self.data_cls:
      assert etype is not None
      return self.edge_feat[etype]
    return self.edge_feat

  @abstractmethod
  def _partition_node(
    self,
    ntype: Optional[NodeType] = None
  ) -> Tuple[List[torch.Tensor], PartitionBook]:
    r""" Partition graph nodes of a specify node type, needs to be overwritten.

    Args:
      ntype (str): The type for input nodes, must be provided for heterogeneous
        graph. (default: ``None``)

    Returns:
      List[torch.Tensor]: The list of partitioned nodes ids.
      PartitionBook: The partition book of graph nodes.
    """

  @abstractmethod
  def _cache_node(
    self,
    ntype: Optional[NodeType] = None
  ) -> List[Optional[torch.Tensor]]:
    r""" Do feature caching and get cached results of a specify
    node type, needs to be overwritten.

    Returns:
      List[Optional[torch.Tensor]]: list of node ids need to be cached on
        each partition.
    """

  def _partition_graph(
    self,
    node_pb: Union[PartitionBook, Dict[NodeType, PartitionBook]],
    etype: Optional[EdgeType] = None
  ) -> Tuple[List[GraphPartitionData], PartitionBook]:
    r""" Partition graph topology of a specified edge type, needs to be
      overwritten.

    Args:
      node_pb (PartitionBook or Dict[NodeType, PartitionBook]): The partition
        books of graph nodes.
      etype (Tuple[str, str, str]): The type for input edges, must be provided
        for heterogeneous graph. (default: ``None``)

    Returns:
      List[GraphPartitionData]: A list of graph data for each partition.
      PartitionBook: The partition book of graph edges.
    """
    edge_index = self.get_edge_index(etype)
    rows, cols = edge_index[0], edge_index[1]
    edge_num = len(rows)
    eids = torch.arange(edge_num, dtype=torch.int64)
    weights = self.edge_weights[etype] if isinstance(self.edge_weights, dict) \
              else self.edge_weights

    if 'hetero' == self.data_cls:
      assert etype is not None
      assert isinstance(node_pb, dict)
      src_ntype, _, dst_ntype = etype

      if 'by_src' == self.edge_assign_strategy:
        target_node_pb = node_pb[src_ntype]
        target_indices = rows
      else:
        target_node_pb = node_pb[dst_ntype]
        target_indices = cols
    else:
      target_node_pb = node_pb
      target_indices = rows if 'by_src' == self.edge_assign_strategy else cols

    chunk_num = (edge_num + self.chunk_size - 1) // self.chunk_size
    chunk_start_pos = 0
    res = [[] for _ in range(self.num_parts)]
    for _ in range(chunk_num):
      chunk_end_pos = min(edge_num, chunk_start_pos + self.chunk_size)
      current_chunk_size = chunk_end_pos - chunk_start_pos
      chunk_idx = torch.arange(current_chunk_size, dtype=torch.long)
      chunk_rows = rows[chunk_start_pos:chunk_end_pos]
      chunk_cols = cols[chunk_start_pos:chunk_end_pos]
      chunk_eids = eids[chunk_start_pos:chunk_end_pos]
      if weights is not None:
        chunk_weights = weights[chunk_start_pos:chunk_end_pos]

      chunk_target_indices = target_indices[chunk_start_pos:chunk_end_pos]
      chunk_partition_idx = target_node_pb[chunk_target_indices]
      for pidx in range(self.num_parts):
        mask = (chunk_partition_idx == pidx)
        idx = torch.masked_select(chunk_idx, mask)
        res[pidx].append(GraphPartitionData(
          edge_index=(chunk_rows[idx], chunk_cols[idx]),
          eids=chunk_eids[idx],
          weights=chunk_weights[idx] if weights is not None else None
        ))
      chunk_start_pos += current_chunk_size

    partition_book = torch.zeros(edge_num, dtype=torch.long)
    partition_results = []
    for pidx in range(self.num_parts):
      p_rows = torch.cat([r.edge_index[0] for r in res[pidx]])
      p_cols = torch.cat([r.edge_index[1] for r in res[pidx]])
      p_eids = torch.cat([r.eids for r in res[pidx]])
      if weights is not None:
        p_weights = torch.cat([r.weights for r in res[pidx]])
      partition_book[p_eids] = pidx
      partition_results.append(GraphPartitionData(
        edge_index=(p_rows, p_cols),
        eids=p_eids,
        weights=p_weights if weights is not None else None
      ))

    return partition_results, partition_book

  def _partition_and_save_node_feat(
    self,
    node_ids_list: List[torch.Tensor],
    ntype: Optional[NodeType] = None,
  ):
    r""" Partition node features by the partitioned node results, and calculate
    the cached nodes if needed.
    """
    node_feat = self.get_node_feat(ntype)
    if node_feat is None:
      return
    cache_node_ids_list = self._cache_node(ntype)
    for pidx in range(self.num_parts):
      # save partitioned node feature cache
      cache_n_ids = cache_node_ids_list[pidx]
      p_node_cache_feat = FeaturePartitionData(
        feats=None, 
        ids=None,
        cache_feats=(node_feat[cache_n_ids] if cache_n_ids is not None else None),
        cache_ids=cache_n_ids
      )
      save_feature_partition_cache(self.output_dir, pidx, p_node_cache_feat,
                                   group='node_feat', graph_type=ntype)
      # save partitioned node feature chunk
      n_ids = node_ids_list[pidx]
      n_ids_chunks = torch.chunk(n_ids, chunks=((n_ids.shape[0] + self.chunk_size - 1) // self.chunk_size))
      for chunk in n_ids_chunks:
        p_node_feat_chunk = FeaturePartitionData(
          feats=node_feat[chunk],
          ids=chunk.clone(),
          cache_feats=None,
          cache_ids=None
        )
        save_feature_partition_chunk(self.output_dir, pidx, p_node_feat_chunk,
                                     group='node_feat', graph_type=ntype)

  def _partition_and_save_edge_feat(
    self,
    graph_list: List[GraphPartitionData],
    etype: Optional[EdgeType] = None
  ):
    r""" Partition edge features by the partitioned edge results.
    """
    edge_feat = self.get_edge_feat(etype)
    if edge_feat is None:
      return
    for pidx in range(self.num_parts):
      eids = graph_list[pidx].eids
      eids_chunks = torch.chunk(
        eids, chunks=((eids.shape[0] + self.chunk_size - 1) // self.chunk_size)
      )
      for chunk in eids_chunks:
        p_edge_feat_chunk = FeaturePartitionData(
          feats=edge_feat[chunk],
          ids=chunk.clone(),
          cache_feats=None,
          cache_ids=None
        )
        save_feature_partition_chunk(self.output_dir, pidx, p_edge_feat_chunk,
                                     group='edge_feat', graph_type=etype)

  def _process_node(self, ntype, with_feature):
    node_ids_list, node_pb = self._partition_node(ntype)
    save_node_pb(self.output_dir, node_pb, ntype)
    self.node_pb_dict[ntype] = node_pb 
    if with_feature:
      self._partition_and_save_node_feat(node_ids_list, ntype)
  
  def _process_edge(self, etype, with_feature):
    graph_list, edge_pb = self._partition_graph(self.node_pb_dict, etype)
    save_edge_pb(self.output_dir, edge_pb, etype)
    for pidx in range(self.num_parts):
      save_graph_partition(self.output_dir, pidx, graph_list[pidx], etype)
    if with_feature:
      self._partition_and_save_edge_feat(graph_list, etype)
  
  def partition(self, with_feature=True, graph_caching=False):
    r""" Partition graph and feature data into different parts. 
    
    Args:
      with_feature (bool): A flag indicating if the feature should be 
        partitioned with the graph (default: ``True``).
      graph_caching (bool): A flag indicating if the full graph topology
        will be saved (default: ``False``).

    The output directory of partitioned graph data will be like:

    * homogeneous

      root_dir/
      |-- META
      |-- node_pb.pt
      |-- edge_pb.pt
      |-- part0/
          |-- graph/
              |-- rows.pt
              |-- cols.pt
              |-- eids.pt
              |-- weights.pt (optional)
          |-- node_feat/
              |-- feats.pkl
              |-- ids.pkl
              |-- cache_feats.pt (optional)
              |-- cache_ids.pt (optional)
          |-- edge_feat/
              |-- feats.pkl
              |-- ids.pkl
              |-- cache_feats.pt (optional)
              |-- cache_ids.pt (optional)
      |-- part1/
          |-- graph/
              ...
          |-- node_feat/
              ...
          |-- edge_feat/
              ...

    * heterogeneous

      root_dir/
      |-- META
      |-- node_pb/
          |-- ntype1.pt
          |-- ntype2.pt
      |-- edge_pb/
          |-- etype1.pt
          |-- etype2.pt
      |-- part0/
          |-- graph/
              |-- etype1/
                  |-- rows.pt
                  |-- cols.pt
                  |-- eids.pt
                  |-- weights.pt
              |-- etype2/
                  ...
          |-- node_feat/
              |-- ntype1/
                  |-- feats.pkl
                  |-- ids.pkl
                  |-- cache_feats.pt (optional)
                  |-- cache_ids.pt (optional)
              |-- ntype2/
                  ...
          |-- edge_feat/
              |-- etype1/
                  |-- feats.pkl
                  |-- ids.pkl
                  |-- cache_feats.pt (optional)
                  |-- cache_ids.pt (optional)
              |-- etype2/
                  ...
      |-- part1/
          |-- graph/
              ...
          |-- node_feat/
              ...
          |-- edge_feat/
              ...

    """
    if 'hetero' == self.data_cls:
      node_pb_dict = {}
      for ntype in self.node_types:
        node_ids_list, node_pb = self._partition_node(ntype)
        save_node_pb(self.output_dir, node_pb, ntype)
        node_pb_dict[ntype] = node_pb
        if with_feature:
          self._partition_and_save_node_feat(node_ids_list, ntype)

      for etype in self.edge_types:
        graph_list, edge_pb = self._partition_graph(node_pb_dict, etype)
        edge_feat = self.get_edge_feat(etype)
        with_edge_feat = (edge_feat != None)
        if graph_caching:
          if with_edge_feat:
            save_edge_pb(self.output_dir, edge_pb, etype)
          save_graph_cache(self.output_dir, graph_list, etype, with_edge_feat)
        else:
          save_edge_pb(self.output_dir, edge_pb, etype)
          for pidx in range(self.num_parts):
            save_graph_partition(self.output_dir, pidx, graph_list[pidx], etype)
        if with_feature:
          self._partition_and_save_edge_feat(graph_list, etype)

    else:
      node_ids_list, node_pb = self._partition_node()
      save_node_pb(self.output_dir, node_pb)
      if with_feature:
        self._partition_and_save_node_feat(node_ids_list)

      graph_list, edge_pb = self._partition_graph(node_pb)
      edge_feat = self.get_edge_feat()
      with_edge_feat = (edge_feat != None)

      if graph_caching:
        if with_edge_feat:
          save_edge_pb(self.output_dir, edge_pb)
        save_graph_cache(self.output_dir, graph_list, with_edge_feat)
      else:
        save_edge_pb(self.output_dir, edge_pb)
        for pidx in range(self.num_parts):
          save_graph_partition(self.output_dir, pidx, graph_list[pidx])
      if with_feature:
        self._partition_and_save_edge_feat(graph_list)

    # save meta.
    save_meta(self.output_dir, self.num_parts, self.data_cls,
              self.node_types, self.edge_types)


def build_partition_feature(
  root_dir: str, 
  partition_idx: int,
  chunk_size: int = 10000,
  node_feat: Optional[Union[TensorDataType, Dict[NodeType, TensorDataType]]] = None,
  node_feat_dtype: torch.dtype = torch.float32,
  edge_feat: Optional[Union[TensorDataType, Dict[EdgeType, TensorDataType]]] = None,
  edge_feat_dtype: torch.dtype = torch.float32):
  
  r""" In the case that the graph topology is partitioned, but the feature
       partitioning is not executed. This method extracts and persist the 
       feature for a specific partition. 
  
  Args:
    root_dir (str): The root directory for saved partition files.
    partition_idx (int): The partition idx.
    chunk_size: The chunk size for partitioning.
    node_feat: The node feature data, should be a dict for hetero data.
    node_feat_dtype: The data type of node features.
    edge_feat: The edge feature data, should be a dict for hetero data.
    edge_feat_dtype: The data type of edge features.

  """
  with open(os.path.join(root_dir, 'META'), 'rb') as infile:
    meta = pickle.load(infile)
  num_partitions = meta['num_parts']
  assert partition_idx >= 0
  assert partition_idx < num_partitions
  partition_dir = os.path.join(root_dir, f'part{partition_idx}')
  assert os.path.exists(partition_dir)
  graph_dir = os.path.join(partition_dir, 'graph')
  device = torch.device('cpu')

  node_feat = convert_to_tensor(node_feat, dtype=node_feat_dtype)
  edge_feat = convert_to_tensor(edge_feat, dtype=edge_feat_dtype)
  
  # homogenous
  if meta['data_cls'] == 'homo':
    # step 1: build and persist the node feature partition
    node_pb = torch.load(os.path.join(root_dir, 'node_pb.pt'), 
      map_location=device)
    node_num = node_pb.size(0)
    ids = torch.arange(node_num, dtype=torch.int64)
    mask = (node_pb == partition_idx)
    n_ids = torch.masked_select(ids, mask)
    # save partitioned node feature chunk
    n_ids_chunks = torch.chunk(n_ids,
      chunks=((n_ids.shape[0] + chunk_size - 1) // chunk_size))
    for chunk in n_ids_chunks:
      p_node_feat_chunk = FeaturePartitionData(
        feats=node_feat[chunk],
        ids=chunk.clone(),
        cache_feats=None,
        cache_ids=None
      )
      save_feature_partition_chunk(root_dir, partition_idx, p_node_feat_chunk,
                                   group='node_feat', graph_type=None)
    
    # step 2: build and persist the edge feature partition
    if edge_feat is None:
      return
    graph = load_graph_partition_data(graph_dir, device)
    eids = graph.eids
    eids_chunks = torch.chunk(
      eids, chunks=((eids.shape[0] + chunk_size - 1) // chunk_size)
    )
    for chunk in eids_chunks:
      p_edge_feat_chunk = FeaturePartitionData(
        feats=edge_feat[chunk],
        ids=chunk.clone(),
        cache_feats=None,
        cache_ids=None
      )
      save_feature_partition_chunk(root_dir, partition_idx, p_edge_feat_chunk,
                                   group='edge_feat', graph_type=None)
  # heterogenous
  else:  
    # step 1: build and persist the node feature partition
    node_pb_dir = os.path.join(root_dir, 'node_pb')
    for ntype in node_feat.keys():
      node_pb = torch.load(
        os.path.join(node_pb_dir, f'{as_str(ntype)}.pt'), map_location=device)
      feat = node_feat[ntype]
      node_num = node_pb.size(0)
      ids = torch.arange(node_num, dtype=torch.int64)
      mask = (node_pb == partition_idx)
      n_ids = torch.masked_select(ids, mask)
      # save partitioned node feature chunk
      n_ids_chunks = torch.chunk(n_ids, 
        chunks=((n_ids.shape[0] + chunk_size - 1) // chunk_size))
      for chunk in n_ids_chunks:
        p_node_feat_chunk = FeaturePartitionData(
          feats=feat[chunk],
          ids=chunk.clone(),
          cache_feats=None,
          cache_ids=None
        )
        save_feature_partition_chunk(root_dir, partition_idx, p_node_feat_chunk,
                                    group='node_feat', graph_type=ntype)
    # step 2: build and persist the edge feature partition
    if edge_feat is None:
        return
    for etype in edge_feat.keys():
      feat = edge_feat[etype]
      graph = load_graph_partition_data(
        os.path.join(graph_dir, as_str(etype)), device)
      eids = graph.eids
      eids_chunks = torch.chunk(
        eids, chunks=((eids.shape[0] + chunk_size - 1) // chunk_size)
      )
      for chunk in eids_chunks:
        p_edge_feat_chunk = FeaturePartitionData(
          feats=feat[chunk],
          ids=chunk.clone(),
          cache_feats=None,
          cache_ids=None
        )
        save_feature_partition_chunk(root_dir, partition_idx, p_edge_feat_chunk,
                                    group='edge_feat', graph_type=etype)

def load_graph_partition_data(
  graph_data_dir: str,
  device: torch.device
) -> GraphPartitionData:
  r""" Load a graph partition data from the specified directory.
  """
  if not os.path.exists(graph_data_dir):
    return None
  rows = torch.load(os.path.join(graph_data_dir, 'rows.pt'),
                    map_location=device)
  cols = torch.load(os.path.join(graph_data_dir, 'cols.pt'),
                    map_location=device)
  eids = None
  eids_dir = os.path.join(graph_data_dir, 'eids.pt')
  if os.path.exists(eids_dir):
    eids = torch.load(eids_dir, map_location=device)

  if os.path.exists(os.path.join(graph_data_dir, 'weights.pt')):
    weights = torch.load(os.path.join(graph_data_dir, 'weights.pt'),
                          map_location=device)
  else:
    weights = None
  pdata = GraphPartitionData(edge_index=(rows, cols), eids=eids, weights=weights)
  return pdata


def load_feature_partition_data(
  feature_data_dir: str,
  device: torch.device
) -> FeaturePartitionData:
  r""" Load a feature partition data from the specified directory.
  """
  if not os.path.exists(feature_data_dir):
    return None

  feats = load_and_concatenate_tensors(os.path.join(feature_data_dir, 'feats.pkl'), device)
  ids = load_and_concatenate_tensors(os.path.join(feature_data_dir, 'ids.pkl'), device)
  cache_feats_path = os.path.join(feature_data_dir, 'cache_feats.pt')
  cache_ids_path = os.path.join(feature_data_dir, 'cache_ids.pt')
  cache_feats = None
  cache_ids = None
  if os.path.exists(cache_feats_path) and os.path.exists(cache_ids_path):
    cache_feats = torch.load(cache_feats_path, map_location=device)
    cache_ids = torch.load(cache_ids_path, map_location=device)
  pdata = FeaturePartitionData(
    feats=feats, ids=ids, cache_feats=cache_feats, cache_ids=cache_ids
  )
  return pdata


def load_partition(
  root_dir: str,
  partition_idx: int,
  graph_caching: bool = False,
  device: torch.device = torch.device('cpu')
) -> Union[Tuple[int, int,
                 GraphPartitionData,
                 Optional[FeaturePartitionData],
                 Optional[FeaturePartitionData],
                 PartitionBook,
                 PartitionBook],
           Tuple[int, int,
                 HeteroGraphPartitionData,
                 Optional[HeteroFeaturePartitionData],
                 Optional[HeteroFeaturePartitionData],
                 HeteroNodePartitionDict,
                 HeteroEdgePartitionDict]]:
  r""" Load a partition from saved directory.

  Args:
    root_dir (str): The root directory for saved files.
    partition_idx (int): The partition idx to load.
    device (torch.device): The device where loaded graph partition data locates.
    graph_caching: (bool): Whether to load entire graph topology

  Returns:
    int: Number of all partitions.
    int: The current partition idx.
    GraphPartitionData/HeteroGraphPartitionData: graph partition data.
    FeaturePartitionData/HeteroFeaturePartitionData: node feature partition
      data, optional.
    FeaturePartitionData/HeteroFeaturePartitionData: edge feature partition
      data, optional.
    PartitionBook/HeteroNodePartitionDict: node partition book.
    PartitionBook/HeteroEdgePartitionDict: edge partition book.
  """
  with open(os.path.join(root_dir, 'META'), 'rb') as infile:
    meta = pickle.load(infile)
  num_partitions = meta['num_parts']
  assert partition_idx >= 0
  assert partition_idx < num_partitions
  partition_dir = os.path.join(root_dir, f'part{partition_idx}')
  assert os.path.exists(partition_dir)

  if graph_caching:
    graph_dir = os.path.join(root_dir, 'graph')
  else:
    graph_dir = os.path.join(partition_dir, 'graph')
  node_feat_dir = os.path.join(partition_dir, 'node_feat')
  edge_feat_dir = os.path.join(partition_dir, 'edge_feat')

  # homogenous

  if meta['data_cls'] == 'homo':
    graph = load_graph_partition_data(graph_dir, device)
    node_feat = load_feature_partition_data(node_feat_dir, device)
    edge_feat = load_feature_partition_data(edge_feat_dir, device)
    node_pb = torch.load(os.path.join(root_dir, 'node_pb.pt'),
                         map_location=device)
    edge_pb = torch.load(os.path.join(root_dir, 'edge_pb.pt'),
                         map_location=device)
    return (
      num_partitions, partition_idx,
      graph, node_feat, edge_feat, node_pb, edge_pb
    )

  # heterogenous

  graph_dict = {}
  for etype in meta['edge_types']:
    graph_dict[etype] = load_graph_partition_data(
      os.path.join(graph_dir, as_str(etype)), device)

  node_feat_dict = {}
  for ntype in meta['node_types']:
    node_feat = load_feature_partition_data(
      os.path.join(node_feat_dir, as_str(ntype)), device)
    if node_feat is not None:
      node_feat_dict[ntype] = node_feat
  if len(node_feat_dict) == 0:
    node_feat_dict = None

  edge_feat_dict = {}
  for etype in meta['edge_types']:
    edge_feat = load_feature_partition_data(
      os.path.join(edge_feat_dir, as_str(etype)), device)
    if edge_feat is not None:
      edge_feat_dict[etype] = edge_feat
  if len(edge_feat_dict) == 0:
    edge_feat_dict = None

  node_pb_dict = {}
  node_pb_dir = os.path.join(root_dir, 'node_pb')
  for ntype in meta['node_types']:
    node_pb_dict[ntype] = torch.load(
      os.path.join(node_pb_dir, f'{as_str(ntype)}.pt'), map_location=device)

  edge_pb_dict = {}
  edge_pb_dir = os.path.join(root_dir, 'edge_pb')
  for etype in meta['edge_types']:
    edge_pb_file = os.path.join(edge_pb_dir, f'{as_str(etype)}.pt')
    if os.path.exists(edge_pb_file):
      edge_pb_dict[etype] = torch.load(
        edge_pb_file, map_location=device)

  return (
    num_partitions, partition_idx,
    graph_dict, node_feat_dict, edge_feat_dict, node_pb_dict, edge_pb_dict
  )


def cat_feature_cache(
  partition_idx: int,
  feat_pdata: FeaturePartitionData,
  feat_pb: PartitionBook
) -> Tuple[float, torch.Tensor, torch.Tensor, PartitionBook]:
  r""" Concatenate and deduplicate partitioned features and its cached
  features into a new feature patition.

  Note that if the input `feat_pdata` does not contain a feature cache, this
  func will do nothing and return the results corresponding to the original
  partition data.

  Returns:
    float: The proportion of cache features.
    torch.Tensor: The new feature tensor, where the cached feature data is
      arranged before the original partition data.
    torch.Tensor: The tensor that indicates the mapping from global node id
      to its local index in new features.
    PartitionBook: The modified partition book for the new feature tensor.
  """
  feats = feat_pdata.feats
  ids = feat_pdata.ids
  cache_feats = feat_pdata.cache_feats
  cache_ids = feat_pdata.cache_ids
  if cache_feats is None or cache_ids is None:
    return 0.0, feats, id2idx(ids), feat_pb
  device = feats.device
  cache_ratio = cache_ids.size(0) / (cache_ids.size(0) + ids.size(0))
  # cat features
  new_feats = torch.cat([cache_feats, feats])
  # compute id2idx
  max_id = max(torch.max(cache_ids).item(), torch.max(ids).item())
  nid2idx = torch.zeros(max_id + 1, dtype=torch.int64, device=device)
  nid2idx[ids] = (torch.arange(ids.size(0), dtype=torch.int64, device=device) +
                  cache_ids.size(0))
  nid2idx[cache_ids] = torch.arange(cache_ids.size(0), dtype=torch.int64,
                                   device=device)
  # modify partition book
  new_feat_pb = feat_pb.clone()
  new_feat_pb[cache_ids] = partition_idx

  return cache_ratio, new_feats, nid2idx, new_feat_pb