Execute the deduplication operation.
(self, context: ExecutionContext)
| 588 | self.is_materialization = is_materialization |
| 589 | |
| 590 | def execute(self, context: ExecutionContext) -> DAGValue: |
| 591 | """Execute the deduplication operation.""" |
| 592 | input_value = self.get_single_input_value(context) |
| 593 | input_value.assert_format(DAGFormat.RAY) |
| 594 | dataset: Dataset = input_value.data |
| 595 | |
| 596 | join_keys = self.column_info.join_keys |
| 597 | timestamp_col = self.column_info.timestamp_column |
| 598 | |
| 599 | if join_keys: |
| 600 | if self.is_materialization: |
| 601 | # Per-block dedup: streaming-safe, no full shuffle required. |
| 602 | # Cross-block duplicates are handled by the online-store UPSERT. |
| 603 | # |
| 604 | # IMPORTANT: do NOT call dataset.schema() here. For streaming |
| 605 | # datasets backed by slow map_batches actors, .schema() triggers |
| 606 | # eager block execution to |
| 607 | # infer the output type. Those blocks are consumed and LOST — |
| 608 | # they never reach the write stage. We therefore defer the |
| 609 | # column-existence check to inside _dedup_block, which runs in |
| 610 | # a worker per block without interfering with streaming. |
| 611 | _join_keys = list(join_keys) |
| 612 | _ts_col = timestamp_col |
| 613 | |
| 614 | def _dedup_block(block: pd.DataFrame) -> pd.DataFrame: |
| 615 | available = [k for k in _join_keys if k in block.columns] |
| 616 | if not available: |
| 617 | return block |
| 618 | if _ts_col and _ts_col in block.columns: |
| 619 | block = block.sort_values(_ts_col, ascending=False) |
| 620 | return block.drop_duplicates(subset=available) |
| 621 | |
| 622 | dataset = dataset.map_batches(_dedup_block, batch_format="pandas") |
| 623 | else: |
| 624 | # Global dedup via groupby: required for historical retrieval |
| 625 | # where the entity–timestamp join must return exactly one row |
| 626 | # per (entity, query-timestamp) pair. |
| 627 | # NOTE: groupby().map_groups() is a full shuffle and blocks |
| 628 | # until ALL upstream blocks are produced. Use only when |
| 629 | # correctness across partition boundaries is mandatory. |
| 630 | available_join_keys = [ |
| 631 | k for k in join_keys if k in dataset.schema().names |
| 632 | ] |
| 633 | available_ts_col = ( |
| 634 | timestamp_col if timestamp_col in dataset.schema().names else None |
| 635 | ) |
| 636 | |
| 637 | if available_join_keys: |
| 638 | |
| 639 | def _keep_latest_in_group(group: pd.DataFrame) -> pd.DataFrame: |
| 640 | if available_ts_col and available_ts_col in group.columns: |
| 641 | group = group.sort_values(available_ts_col, ascending=False) |
| 642 | return group.head(1) |
| 643 | |
| 644 | dataset = dataset.groupby(available_join_keys).map_groups( |
| 645 | _keep_latest_in_group, batch_format="pandas" |
| 646 | ) |
| 647 |