code: def compact_csc_format(
      csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]],
      dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]],
      dst_timestamps: Optional[
          Union[torch.Tensor, Dict[str, torch.Tensor]]
      ] = None,
  ):
      """
      Relabel the row (source) IDs in the csc formats into a contiguous range from
      0 and return the original row node IDs per type.
  
      Note that
      1. The column (destination) IDs are included in the relabeled row IDs.
      2. If there are repeated row IDs, they would not be uniqued and will be
      treated as different nodes.
      3. If `dst_timestamps` is given, the timestamp of each destination node will
      be broadcasted to its corresponding source nodes.
  
      Parameters
      ----------
      csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]]
          CSC formats representing source-destination edges.
          - If `csc_formats` is a CSCFormatBase: It means the graph is
          homogeneous. Also, indptr and indice in it should be torch.tensor
          representing source and destination pairs in csc format. And IDs inside
          are homogeneous ids.
          - If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
          should be edge type and the values should be csc format node pairs.
          And IDs inside are heterogeneous ids.
      dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]]
          Nodes of all destination nodes in the node pairs.
          - If `dst_nodes` is a tensor: It means the graph is homogeneous.
          - If `dst_nodes` is a dictionary: The keys are node type and the
          values are corresponding nodes. And IDs inside are heterogeneous ids.
  
      dst_timestamps: Optional[Union[torch.Tensor, Dict[str, torch.Tensor]]]
          Timestamps of all destination nodes in the csc formats.
          If given, the timestamp of each destination node will be broadcasted
          to its corresponding source nodes.
  
      Returns
      -------
      Tuple[original_row_node_ids, compacted_csc_formats, ...]
          A tensor of original row node IDs (per type) of all nodes in the input.
          The compacted CSC formats, where node IDs are replaced with mapped node
          IDs ranging from 0 to N.
          The source timestamps (per type) of all nodes in the input if
          `dst_timestamps` is given.
  
      Examples
      --------
      >>> import dgl.graphbolt as gb
      >>> csc_formats = {
      ...     "n2:e2:n1": gb.CSCFormatBase(
      ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
      ...     ),
      ...     "n1:e1:n1": gb.CSCFormatBase(
      ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
      ...     ),
      ... }
      >>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
      >>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
      ...     csc_formats, dst_nodes
      ... )
      >>> original_row_node_ids
      {'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
      >>> compacted_csc_formats
      {'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
                  indices=tensor([0, 1, 2]),
      ), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
                  indices=tensor([2, 3, 4]),
      )}
  
      >>> csc_formats = {
      ...     "n2:e2:n1": gb.CSCFormatBase(
      ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
      ...     ),
      ...     "n1:e1:n1": gb.CSCFormatBase(
      ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
      ...     ),
      ... }
      >>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
      >>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
      ...     csc_formats, dst_nodes
      ... )
      >>> original_row_node_ids
      {'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
      >>> compacted_csc_formats
      {'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
                  indices=tensor([0, 1, 2]),
      ), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
                  indices=tensor([2, 3, 4]),
      )}
  
      >>> dst_timestamps = {"n1": torch.LongTensor([10, 20])}
      >>> (
      ...     original_row_node_ids,
      ...     compacted_csc_formats,
      ...     src_timestamps,
      ... ) = gb.compact_csc_format(csc_formats, dst_nodes, dst_timestamps)
      >>> src_timestamps
      {'n1': tensor([10, 20, 10, 20, 20]), 'n2': tensor([10, 20, 20])}
      """
      is_homogeneous = not isinstance(csc_formats, dict)
      has_timestamp = dst_timestamps is not None
      if is_homogeneous:
          if dst_nodes is not None:
              assert isinstance(
                  dst_nodes, torch.Tensor
              ), "Edge type not supported in homogeneous graph."
              assert len(dst_nodes) + 1 == len(
                  csc_formats.indptr
              ), "The seed nodes should correspond to indptr."
          offset = dst_nodes.size(0)
          original_row_ids = torch.cat((dst_nodes, csc_formats.indices))
          compacted_csc_formats = CSCFormatBase(
              indptr=csc_formats.indptr,
              indices=(
                  torch.arange(
                      0,
                      csc_formats.indices.size(0),
                      device=csc_formats.indices.device,
                  )
                  + offset
              ),
          )
  
          src_timestamps = None
          if has_timestamp:
              src_timestamps = torch.cat(
                  [
                      dst_timestamps,
                      _broadcast_timestamps(
                          compacted_csc_formats, dst_timestamps
                      ),
                  ]
              )
      else:
          compacted_csc_formats = {}
          src_timestamps = None
          original_row_ids = copy.deepcopy(dst_nodes)
          if has_timestamp:
              src_timestamps = copy.deepcopy(dst_timestamps)
          for etype, csc_format in csc_formats.items():
              src_type, _, dst_type = etype_str_to_tuple(etype)
              assert len(dst_nodes.get(dst_type, [])) + 1 == len(
                  csc_format.indptr
              ), "The seed nodes should correspond to indptr."
              device = csc_format.indices.device
              offset = original_row_ids.get(
                  src_type, torch.tensor([], device=device)
              ).size(0)
              original_row_ids[src_type] = torch.cat(
                  (
                      original_row_ids.get(
                          src_type,
                          torch.tensor(
                              [],
                              dtype=csc_format.indices.dtype,
                              device=device,
                          ),
                      ),
                      csc_format.indices,
                  )
              )
              compacted_csc_formats[etype] = CSCFormatBase(
                  indptr=csc_format.indptr,
                  indices=(
                      torch.arange(
                          0,
                          csc_format.indices.size(0),
                          dtype=csc_format.indices.dtype,
                          device=device,
                      )
                      + offset
                  ),
              )
              if has_timestamp:
                  # If destination timestamps are given, broadcast them to the
                  # corresponding source nodes.
                  src_timestamps[src_type] = torch.cat(
                      (
                          src_timestamps.get(
                              src_type,
                              torch.tensor(
                                  [],
                                  dtype=dst_timestamps[dst_type].dtype,
                              ),
                          ),
                          _broadcast_timestamps(
                              csc_format, dst_timestamps[dst_type]
                          ),
                      )
                  )
      if has_timestamp:
          return original_row_ids, compacted_csc_formats, src_timestamps
      return original_row_ids, compacted_csc_formats
  
file_path: dgl/python/dgl/graphbolt/internal/sample_utils.py
imports: import copy
  from collections import defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from ..base import CSCFormatBase, etype_str_to_tuple
label_desc: Use image node_id 5 for calling a global function with example usage: compact_csc_format(csc_formats, dst_nodes, dst_timestamps) and returns: original_row_ids,  compacted_csc_formats, original_row_ids,  compacted_csc_formats,  src_timestamps