Sampling Request
| 495 | |
| 496 | |
| 497 | class SamplingRequestEtype(Request): |
| 498 | """Sampling Request""" |
| 499 | |
| 500 | def __init__( |
| 501 | self, |
| 502 | nodes, |
| 503 | fan_out, |
| 504 | edge_dir="in", |
| 505 | prob=None, |
| 506 | exclude_edges=None, |
| 507 | replace=False, |
| 508 | etype_sorted=True, |
| 509 | use_graphbolt=False, |
| 510 | ): |
| 511 | self.seed_nodes = nodes |
| 512 | self.edge_dir = edge_dir |
| 513 | self.prob = prob |
| 514 | self.exclude_edges = exclude_edges |
| 515 | self.replace = replace |
| 516 | self.fan_out = fan_out |
| 517 | self.etype_sorted = etype_sorted |
| 518 | self.use_graphbolt = use_graphbolt |
| 519 | |
| 520 | def __setstate__(self, state): |
| 521 | ( |
| 522 | self.seed_nodes, |
| 523 | self.edge_dir, |
| 524 | self.prob, |
| 525 | self.exclude_edges, |
| 526 | self.replace, |
| 527 | self.fan_out, |
| 528 | self.etype_sorted, |
| 529 | self.use_graphbolt, |
| 530 | ) = state |
| 531 | |
| 532 | def __getstate__(self): |
| 533 | return ( |
| 534 | self.seed_nodes, |
| 535 | self.edge_dir, |
| 536 | self.prob, |
| 537 | self.exclude_edges, |
| 538 | self.replace, |
| 539 | self.fan_out, |
| 540 | self.etype_sorted, |
| 541 | self.use_graphbolt, |
| 542 | ) |
| 543 | |
| 544 | def process_request(self, server_state): |
| 545 | local_g = server_state.graph |
| 546 | partition_book = server_state.partition_book |
| 547 | kv_store = server_state.kv_store |
| 548 | etype_offset = partition_book.local_etype_offset |
| 549 | # See NOTE 1 |
| 550 | if self.prob is not None and (not self.use_graphbolt): |
| 551 | probs = [ |
| 552 | kv_store.data_store[key] if key != "" else None |
| 553 | for key in self.prob |
| 554 | ] |