void CPUWeightedSampler::CSRRowWiseSample()

in graphlearn_torch/csrc/cpu/weighted_sampler.cc [121:144]

• 24 lines of code
• 3 McCabe index (conditional complexity)

void CPUWeightedSampler::CSRRowWiseSample(
    const int64_t* nodes,
    const int64_t* nbrs_offset,
    const int32_t bs,
    const int32_t req_num,
    const int64_t row_count,
    const int64_t* row_ptr,
    const int64_t* col_idx,
    const float* prob,
    const int64_t* edge_ids,
    int64_t* out_nbrs,
    int64_t* out_eid) {
  at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){
    for(int32_t i = start; i < end; ++i) {
      auto v = nodes[i];
      if (v < row_count) {
        WeightedSample(col_idx + row_ptr[v], col_idx + row_ptr[v+1],
            edge_ids + row_ptr[v], edge_ids + row_ptr[v+1],
            req_num, prob + row_ptr[v], prob + row_ptr[v+1],
            out_nbrs + nbrs_offset[i], out_eid + nbrs_offset[i]);
      }
    }
  });
}