Sampling Request
| 425 | |
| 426 | |
| 427 | class SamplingRequest(Request): |
| 428 | """Sampling Request""" |
| 429 | |
| 430 | def __init__( |
| 431 | self, |
| 432 | nodes, |
| 433 | fan_out, |
| 434 | edge_dir="in", |
| 435 | prob=None, |
| 436 | exclude_edges=None, |
| 437 | replace=False, |
| 438 | use_graphbolt=False, |
| 439 | ): |
| 440 | self.seed_nodes = nodes |
| 441 | self.edge_dir = edge_dir |
| 442 | self.prob = prob |
| 443 | self.exclude_edges = exclude_edges |
| 444 | self.replace = replace |
| 445 | self.fan_out = fan_out |
| 446 | self.use_graphbolt = use_graphbolt |
| 447 | |
| 448 | def __setstate__(self, state): |
| 449 | ( |
| 450 | self.seed_nodes, |
| 451 | self.edge_dir, |
| 452 | self.prob, |
| 453 | self.exclude_edges, |
| 454 | self.replace, |
| 455 | self.fan_out, |
| 456 | self.use_graphbolt, |
| 457 | ) = state |
| 458 | |
| 459 | def __getstate__(self): |
| 460 | return ( |
| 461 | self.seed_nodes, |
| 462 | self.edge_dir, |
| 463 | self.prob, |
| 464 | self.exclude_edges, |
| 465 | self.replace, |
| 466 | self.fan_out, |
| 467 | self.use_graphbolt, |
| 468 | ) |
| 469 | |
| 470 | def process_request(self, server_state): |
| 471 | local_g = server_state.graph |
| 472 | partition_book = server_state.partition_book |
| 473 | kv_store = server_state.kv_store |
| 474 | if self.prob is not None and (not self.use_graphbolt): |
| 475 | prob = [kv_store.data_store[self.prob]] |
| 476 | else: |
| 477 | prob = self.prob |
| 478 | res = _sample_neighbors( |
| 479 | self.use_graphbolt, |
| 480 | local_g, |
| 481 | partition_book, |
| 482 | self.seed_nodes, |
| 483 | self.fan_out, |
| 484 | edge_dir=self.edge_dir, |