void CPURandomSampler::UniformSample()

in graphlearn_torch/csrc/cpu/random_sampler.cc [155:177]


void CPURandomSampler::UniformSample(const int64_t* col_begin,
                                     const int64_t* col_end,
                                     const int64_t* eid_begin,
                                     const int64_t* eid_end,
                                     const int32_t req_num,
                                     int64_t* out_nbrs,
                                     int64_t* out_eid) {
  // with replacement
  const auto cap = col_end - col_begin;
  if (req_num < cap) {
    uint32_t seed = RandomSeedManager::getInstance().getSeed();
    thread_local static std::mt19937 engine(seed);
    std::uniform_int_distribution<> dist(0, cap-1);
    for (int32_t i = 0; i < req_num; ++i) {
      auto idx = dist(engine);
      out_nbrs[i] = col_begin[idx];
      out_eid[i] = eid_begin[idx];
    }
  } else {
    std::copy(col_begin, col_end, out_nbrs);
    std::copy(eid_begin, eid_end, out_eid);
  }
}