[DistGB] use expand_indptr for csc to coo (#7090)

python/dgl/distributed/graph_services.py

@@ -151,8 +151,11 @@ def _sample_neighbors_graphbolt(
 
     # 3. Map local node IDs to global node IDs.
     local_src = subgraph.indices
-    local_dst = torch.repeat_interleave(
-        subgraph.original_column_node_ids, torch.diff(subgraph.indptr)
+    local_dst = gb.expand_indptr(
+        subgraph.indptr,
+        dtype=local_src.dtype,
+        node_ids=subgraph.original_column_node_ids,
+        output_size=local_src.shape[0],
     )
     global_nid_mapping = g.node_attributes[NID]
     global_src = global_nid_mapping[local_src]

python/dgl/distributed/partition.py

@@ -199,11 +199,10 @@ def _verify_graphbolt_partition(graph, part_id, gpb, ntypes, etypes):
         field in graph.edge_attributes for field in required_edata_fields
     ), "the partition graph should contain edge mapping to global edge ID."
 
-    num_nodes = graph.total_num_nodes
     num_edges = graph.total_num_edges
     local_src_ids = graph.indices
-    local_dst_ids = torch.repeat_interleave(
-        torch.arange(num_nodes), torch.diff(graph.csc_indptr)
+    local_dst_ids = gb.expand_indptr(
+        graph.csc_indptr, dtype=local_src_ids.dtype, output_size=num_edges
     )
     global_src_ids = graph.node_attributes[NID][local_src_ids]
     global_dst_ids = graph.node_attributes[NID][local_dst_ids]