void CPUSubGraphOp::Induce()

in graphlearn_torch/csrc/cpu/subgraph_op.cc [61:89]


void CPUSubGraphOp::Induce(const std::vector<int64_t>& nodes,
                           bool with_edge,
                           std::vector<int64_t>& out_rows,
                           std::vector<int64_t>& out_cols,
                           std::vector<int64_t>& out_eids) {
  const auto indptr = graph_->GetRowPtr();
  const auto indices = graph_->GetColIdx();
  const auto edge_ids = graph_->GetEdgeId();
  const auto row_count = graph_->GetRowCount();
  const auto node_size = nodes.size();
  out_rows.reserve(node_size * node_size);
  out_cols.reserve(node_size * node_size);
  if (with_edge) out_eids.reserve(node_size * node_size);

  for (int32_t i = 0; i < node_size; ++i) {
    auto old_row = nodes[i];
    if (old_row < row_count) {
      for (int32_t j = indptr[old_row]; j < indptr[old_row+1]; ++j) {
        auto old_col = indices[j];
        auto col_iter = glob2local_.find(old_col);
        if (col_iter != glob2local_.end()) {
          out_rows.push_back(i);
          out_cols.push_back(col_iter->second);
          if (with_edge) out_eids.push_back(edge_ids[j]);
        }
      }
    }
  }
}