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]);
}
}
}
}
}