docstr: Executes a bulk sampling job immediately upon creation.
  Allows iteration over the returned results.
  
  Parameters
  ----------
  feature_store: CuGraphStore
      The feature store containing features for the graph.
  
  graph_store: CuGraphStore
      The graph store containing the graph structure.
  
  input_nodes: InputNodes
      The input nodes associated with this sampler.
      If None, this loader will load batches
      from disk rather than performing sampling in memory.
  
  batch_size: int
      The number of input nodes per sampling batch.
      Generally required unless loading already-sampled
      data from disk.
  
  shuffle: bool (optional, default=False)
      Whether to shuffle the input indices.
      If True, will shuffle the input indices.
      If False, will create batches in the original order.
  
  edge_types: Sequence[Tuple[str]] (optional, default=None)
      The desired edge types for the subgraph.
      Defaults to all edges in the graph.
  
  directory: str (optional, default=new tempdir)
      The path of the directory to write samples to.
      Defaults to a new generated temporary directory.
  
  input_files: List[str] (optional, default=None)
      The input files to read from the directory containing
      samples.  This argument is only used when loading
      alread-sampled batches from disk.
  
  starting_batch_id: int (optional, default=0)
      The starting id for each batch.  Defaults to 0.
  
  batches_per_partition: int (optional, default=100)
      The number of batches in each output partition.
      Defaults to 100.  Gets passed to the bulk
      sampler if there is one; otherwise, this argument
      is used to determine which files to read.
  
  num_neighbors: Union[List[int],
           Dict[Tuple[str, str, str], List[int]]] (required)
      The number of neighbors to sample for each node in each iteration.
      If an entry is set to -1, all neighbors will be included.
      In heterogeneous graphs, may also take in a dictionary denoting
      the number of neighbors to sample for each individual edge type.
  
      Note: in cuGraph, only one value of num_neighbors is currently supported.
      Passing in a dictionary will result in an exception.
docstr_tok: ["Executes","a","bulk","sampling","job","immediately","upon","creation",".","Allows","iteration","over","the","returned","results",".","Parameters","--","--","--","--","--","feature_store",":","CuGraphStore","The","feature","store","containing","features","for","the","graph",".","graph_store",":","CuGraphStore","The","graph","store","containing","the","graph","structure",".","input_nodes",":","InputNodes","The","input","nodes","associated","with","this","sampler",".","If","None",",","this","loader","will","load","batches","from","disk","rather","than","performing","sampling","in","memory",".","batch_size",":","int","The","number","of","input","nodes","per","sampling","batch",".","Generally","required","unless","loading","already-sampled","data","from","disk",".","shuffle",":","bool","(","optional",",","default=False",")","Whether","to","shuffle","the","input","indices",".","If","True",",","will","shuffle","the","input","indices",".","If","False",",","will","create","batches","in","the","original","order",".","edge_types",":","Sequence","[","Tuple","[","str","]","]","(","optional",",","default=None",")","The","desired","edge","types","for","the","subgraph",".","Defaults","to","all","edges","in","the","graph",".","directory",":","str","(","optional",",","default=new","tempdir",")","The","path","of","the","directory","to","write","samples","to",".","Defaults","to","a","new","generated","temporary","directory",".","input_files",":","List","[","str","]","(","optional",",","default=None",")","The","input","files","to","read","from","the","directory","containing","samples",".","This","argument","is","only","used","when","loading","alread-sampled","batches","from","disk",".","starting_batch_id",":","int","(","optional",",","default=0",")","The","starting","id","for","each","batch",".","Defaults","to","0.","batches_per_partition",":","int","(","optional",",","default=100",")","The","number","of","batches","in","each","output","partition",".","Defaults","to","100",".","Gets","passed","to","the","bulk","sampler","if","there","is","one",";","otherwise",",","this","argument","is","used","to","determine","which","files","to","read",".","num_neighbors",":","Union","[","List","[","int","]",",","Dict","[","Tuple","[","str",",","str",",","str","]",",","List","[","int","]","]","]","(","required",")","The","number","of","neighbors","to","sample","for","each","node","in","each","iteration",".","If","an","entry","is","set","to","-1",",","all","neighbors","will","be","included",".","In","heterogeneous","graphs",",","may","also","take","in","a","dictionary","denoting","the","number","of","neighbors","to","sample","for","each","individual","edge","type",".","Note",":","in","cuGraph",",","only","one","value","of","num_neighbors","is","currently","supported",".","Passing","in","a","dictionary","will","result","in","an","exception","."]
code: def __init__(
      self,
      feature_store: CuGraphStore,
      graph_store: CuGraphStore,
      input_nodes: InputNodes = None,
      batch_size: int = 0,
      *,
      shuffle: bool = False,
      drop_last: bool = True,
      edge_types: Sequence[Tuple[str]] = None,
      directory: Union[str, tempfile.TemporaryDirectory] = None,
      input_files: List[str] = None,
      starting_batch_id: int = 0,
      batches_per_partition: int = 100,
      # Sampler args
      num_neighbors: Union[
          List[int], Dict[Tuple[str, str, str], List[int]]
      ] = None,
      replace: bool = True,
      compression: str = "COO",
      # Other kwargs for the BulkSampler
      **kwargs,
  ):
      """
      Executes a bulk sampling job immediately upon creation.
      Allows iteration over the returned results.
  
      Parameters
      ----------
      feature_store: CuGraphStore
          The feature store containing features for the graph.
  
      graph_store: CuGraphStore
          The graph store containing the graph structure.
  
      input_nodes: InputNodes
          The input nodes associated with this sampler.
          If None, this loader will load batches
          from disk rather than performing sampling in memory.
  
      batch_size: int
          The number of input nodes per sampling batch.
          Generally required unless loading already-sampled
          data from disk.
  
      shuffle: bool (optional, default=False)
          Whether to shuffle the input indices.
          If True, will shuffle the input indices.
          If False, will create batches in the original order.
  
      edge_types: Sequence[Tuple[str]] (optional, default=None)
          The desired edge types for the subgraph.
          Defaults to all edges in the graph.
  
      directory: str (optional, default=new tempdir)
          The path of the directory to write samples to.
          Defaults to a new generated temporary directory.
  
      input_files: List[str] (optional, default=None)
          The input files to read from the directory containing
          samples.  This argument is only used when loading
          alread-sampled batches from disk.
  
      starting_batch_id: int (optional, default=0)
          The starting id for each batch.  Defaults to 0.
  
      batches_per_partition: int (optional, default=100)
          The number of batches in each output partition.
          Defaults to 100.  Gets passed to the bulk
          sampler if there is one; otherwise, this argument
          is used to determine which files to read.
  
      num_neighbors: Union[List[int],
               Dict[Tuple[str, str, str], List[int]]] (required)
          The number of neighbors to sample for each node in each iteration.
          If an entry is set to -1, all neighbors will be included.
          In heterogeneous graphs, may also take in a dictionary denoting
          the number of neighbors to sample for each individual edge type.
  
          Note: in cuGraph, only one value of num_neighbors is currently supported.
          Passing in a dictionary will result in an exception.
      """
  
      self.__feature_store = feature_store
      self.__graph_store = graph_store
      self.__next_batch = -1
      self.__end_exclusive = -1
      self.__batches_per_partition = batches_per_partition
      self.__starting_batch_id = starting_batch_id
  
      self._total_read_time = 0.0
      self._total_convert_time = 0.0
      self._total_feature_time = 0.0
  
      if input_nodes is None:
          # Will be loading from disk
          self.__num_batches = input_nodes
          self.__directory = directory
          if input_files is None:
              if isinstance(self.__directory, str):
                  self.__input_files = iter(
                      os.listdir(self.__directory)
                  )
              else:
                  self.__input_files = iter(
                      os.listdir(self.__directory.name)
                  )
          else:
              self.__input_files = iter(input_files)
          return
  
      # To accommodate DLFW/PyG 2.5
      get_input_nodes = torch_geometric.loader.utils.get_input_nodes
      get_input_nodes_kwargs = {}
      if "input_id" in get_input_nodes.__annotations__:
          get_input_nodes_kwargs["input_id"] = None
      input_node_info = get_input_nodes(
          (feature_store, graph_store),
          input_nodes,
          **get_input_nodes_kwargs,
      )
  
      # PyG 2.4
      if len(input_node_info) == 2:
          input_type, input_nodes = input_node_info
      # PyG 2.5
      elif len(input_node_info) == 3:
          input_type, input_nodes, input_id = input_node_info
      # Invalid
      else:
          raise ValueError("Invalid output from get_input_nodes")
  
      if input_type is not None:
          input_nodes = graph_store._get_sample_from_vertex_groups(
              {input_type: input_nodes}
          )
  
      if batch_size is None or batch_size < 1:
          raise ValueError("Batch size must be >= 1")
  
      self.__directory = (
          tempfile.TemporaryDirectory()
          if directory is None
          else directory
      )
  
      if isinstance(num_neighbors, dict):
          raise ValueError(
              "num_neighbors dict is currently unsupported!"
          )
  
      if "renumber" in kwargs:
          warnings.warn(
              "Setting renumbering manually could result in invalid output,"
              " please ensure you intended to do this."
          )
          renumber = kwargs.pop("renumber")
      else:
          renumber = (
              True
              if (
                  (len(self.__graph_store.node_types) == 1)
                  and (len(self.__graph_store.edge_types) == 1)
              )
              else False
          )
  
      bulk_sampler = BulkSampler(
          batch_size,
          self.__directory
          if isinstance(self.__directory, str)
          else self.__directory.name,
          self.__graph_store._subgraph(edge_types),
          fanout_vals=num_neighbors,
          with_replacement=replace,
          batches_per_partition=self.__batches_per_partition,
          renumber=renumber,
          use_legacy_names=False,
          deduplicate_sources=True,
          prior_sources_behavior="exclude",
          include_hop_column=(compression == "COO"),
          **kwargs,
      )
  
      # Make sure indices are in cupy
      input_nodes = cupy.asarray(input_nodes)
  
      # Shuffle
      if shuffle:
          cupy.random.shuffle(input_nodes)
  
      # Truncate if we can't evenly divide the input array
      stop = (len(input_nodes) // batch_size) * batch_size
      input_nodes, remainder = cupy.array_split(input_nodes, [stop])
  
      # Split into batches
      input_nodes = cupy.split(
          input_nodes, max(len(input_nodes) // batch_size, 1)
      )
  
      if not drop_last:
          input_nodes.append(remainder)
  
      self.__num_batches = 0
      for batch_num, batch_i in enumerate(input_nodes):
          batch_len = len(batch_i)
          if batch_len > 0:
              self.__num_batches += 1
              bulk_sampler.add_batches(
                  cudf.DataFrame(
                      {
                          "start": batch_i,
                          "batch": cupy.full(
                              batch_len,
                              batch_num + starting_batch_id,
                              dtype="int32",
                          ),
                      }
                  ),
                  start_col_name="start",
                  batch_col_name="batch",
              )
  
      bulk_sampler.flush()
      self.__input_files = iter(
          os.listdir(
              self.__directory
              if isinstance(self.__directory, str)
              else self.__directory.name
          )
      )
  
file_path: cugraph/python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py
imports: import tempfile
  import os
  import re
  import warnings
  import cupy
  import cudf
  from cugraph.gnn import BulkSampler
  from cugraph.utilities.utils import import_optional, MissingModule
  from cugraph_pyg.data import CuGraphStore
  from cugraph_pyg.sampler.cugraph_sampler import _sampler_output_from_sampling_results_heterogeneous, _sampler_output_from_sampling_results_homogeneous_csr, _sampler_output_from_sampling_results_homogeneous_coo, filter_cugraph_store_csc
  from typing import Union, Tuple, Sequence, List, Dict
label_desc: Use image node_id 1 to create a new BulkSampleLoader object  with example: obj = BulkSampleLoader(feature_store, graph_store, input_nodes, batch_size)