code: def ego_graph_sample(graph, node_ids, samples, edge_types):
      """Heter-EgoGraph
                            ____ n1
                           /----n2
                   etype1 /---- n3
                 --------/-----n4
                /        \-----n5
               /
              /
      start_n
              \             /----n6
               \  etype2 /---- n7
                 --------/-----n8
                        \-----n9
  
      TODO @Yelrose: Speed up and standarize to pgl.distributed.sampling
      """
  
      # All Nodes
      all_new_nodes = [node_ids]
  
      # Node Index
      all_new_nodes_index = [np.zeros_like(node_ids, dtype="int64")]
  
      # The Ego Index for each graph
      all_new_nodes_ego_index = [
          np.arange(0, len(node_ids), dtype="int64")
      ]
  
      unique_nodes = set(node_ids)
  
      ego_graph_list = [
          EgoInfo(node_id=[n], edges=[], edges_type=[], edges_weight=[])
          for n in node_ids
      ]
  
      for sample in samples:
          cur_node_ids = all_new_nodes[-1]
          cur_node_ego_index = all_new_nodes_ego_index[-1]
          cur_node_index = all_new_nodes_index[-1]
  
          nxt_node_ids = []
          nxt_node_ego_index = []
          nxt_node_index = []
          for edge_type_id, edge_type in enumerate(edge_types):
              cur_succs = graph.sample_successor(
                  cur_node_ids, max_degree=sample, edge_type=edge_type
              )
              for succs, ego_index, parent_index in zip(
                  cur_succs, cur_node_ego_index, cur_node_index
              ):
                  if len(succs) == 0:
                      succs = [0]
                  ego = ego_graph_list[ego_index]
  
                  for succ in succs:
                      unique_nodes.add(succ)
                      succ_index = len(ego.node_id)
                      ego.node_id.append(succ)
                      nxt_node_ids.append(succ)
                      nxt_node_ego_index.append(ego_index)
                      nxt_node_index.append(succ_index)
                      if succ == 0:
                          ego.edges.append((succ_index, succ_index))
                          ego.edges_type.append(edge_type_id)
                          ego.edges_weight.append(1.0 / len(succs))
                      else:
                          ego.edges.append((succ_index, parent_index))
                          ego.edges_type.append(edge_type_id)
                          ego.edges_weight.append(1.0 / len(succs))
          all_new_nodes.append(nxt_node_ids)
          all_new_nodes_index.append(nxt_node_index)
          all_new_nodes_ego_index.append(nxt_node_ego_index)
  
      for ego in ego_graph_list:
          pg = pgl.Graph(
              num_nodes=len(ego.node_id),
              edges=ego.edges,
              edge_feat={
                  "edge_type": np.array(ego.edges_type, dtype="int64"),
                  "edge_weight": np.array(
                      ego.edges_weight, dtype="float32"
                  ),
              },
              node_feat={
                  "node_id": np.array(ego.node_id, dtype="int64"),
              },
          )
          ego.graph = pg
      return ego_graph_list, list(unique_nodes)
  
imports: import os
  import re
  import sys
  import time
  import warnings
  import numpy
  from collections import defaultdict
  import pgl
  from pgl.utils.logger import log
  from pgl.distributed import DistGraphClient, DistGraphServer
  from utils.config import prepare_config
  from datasets.helper import stream_shuffle_generator, AsynchronousGenerator
  from datasets.node import NodeGenerator
  import datasets.sampling
label_desc: Use image node_id 1 for calling a global function with example usage: ego_graph_sample(graph, node_ids, samples, edge_types) and returns: ego_graph_list,  list