diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py b/paimon-python/pypaimon/ray/data_evolution_merge_into.py index e0544f189043..a012883a5bd3 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py @@ -101,7 +101,7 @@ def merge_into( return _execute_and_commit( table, update_ds, delete_ds, insert_ds, update_cols_union, base_snapshot, num_partitions, - ray_remote_args, concurrency, + ray_remote_args, concurrency, is_self_merge=ctx.is_self_merge, ) @@ -424,7 +424,7 @@ def _build_datasets( def _execute_and_commit( table, update_ds, delete_ds, insert_ds, update_cols_union, base_snapshot, num_partitions, - ray_remote_args, concurrency, + ray_remote_args, concurrency, is_self_merge=False, ): collect_action_row_ids = update_ds is not None and delete_ds is not None commit_messages: list = [] @@ -461,6 +461,7 @@ def _execute_and_commit( if base_snapshot is not None else None ), collect_row_ids=collect_action_row_ids, + materialize_before_routing=not is_self_merge, ) ) commit_messages.extend(update_msgs) @@ -475,6 +476,7 @@ def _execute_and_commit( if base_snapshot is not None else None ), collect_row_ids=collect_action_row_ids, + materialize_before_routing=not is_self_merge, ) commit_messages.extend(delete_msgs) diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index af56d7ca302f..39934cc2e4fc 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -659,6 +659,7 @@ def distributed_update_apply( ray_remote_args: Optional[Dict[str, Any]] = None, base_snapshot_id: Optional[int] = None, collect_row_ids: bool = False, + materialize_before_routing: bool = False, ) -> Tuple[list, int, list]: import numpy as np import pickle @@ -697,6 +698,10 @@ def distributed_update_apply( if not sorted_first_row_ids: return [], 0, [] + if materialize_before_routing: + # Keep the matched join and file-routing shuffle in separate Ray jobs. + update_ds = update_ds.materialize() + # Pin commit-time conflict check to the snapshot the join was built on, # so concurrent commits between read and planner are detected. check_from_snapshot = ( @@ -980,6 +985,7 @@ def distributed_delete_apply( ray_remote_args: Optional[Dict[str, Any]] = None, base_snapshot_id: Optional[int] = None, collect_row_ids: bool = False, + materialize_before_routing: bool = False, ) -> Tuple[list, int, list]: import base64 import numpy as np @@ -1003,6 +1009,10 @@ def distributed_delete_apply( if not anchor_info.anchors: return [], 0, [] + if materialize_before_routing: + # Keep the matched join and file-routing shuffle in separate Ray jobs. + delete_ds = delete_ds.materialize() + precomputed_info_ref = ray.put(anchor_info) starts = np.asarray( diff --git a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py index d75dbf4a398f..68d1ed44d75c 100644 --- a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py +++ b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py @@ -419,6 +419,46 @@ def test_matched_delete(self): self.assertEqual(out['name'], ['a']) self.assertEqual(out['age'], [10]) + def test_joined_source_delete_materializes_before_routing(self): + options = dict(self.de_options) + options['deletion-vectors.enabled'] = 'true' + target = self._create_table(options=options) + self._write(target, self._source(ids=(1, 2, 3))) + + source = ray.data.from_arrow(pa.table({ + 'id': pa.array([2, 3], type=pa.int32()), + })).join( + ray.data.from_arrow(pa.table({ + 'id': pa.array([2, 3], type=pa.int32()), + 'selected': [True, True], + })), + join_type='inner', + num_partitions=_TEST_NUM_PARTITIONS, + on=['id'], + ) + real_materialize = ray.data.Dataset.materialize + materialized = [] + + def track_materialize(dataset): + materialized.append(dataset) + return real_materialize(dataset) + + with patch.object( + ray.data.Dataset, 'materialize', new=track_materialize, + ): + metrics = merge_into( + target=target, + source=source, + catalog_options=self.catalog_options, + on=['id'], + when_matched=[WhenMatched.delete()], + num_partitions=_TEST_NUM_PARTITIONS, + ) + + self.assertEqual(metrics['num_matched'], 2) + self.assertEqual(self._read_sorted(target)['id'], [1]) + self.assertEqual(len(materialized), 1) + def test_not_matched_insert_appends_unmatched(self): target = self._create_table() self._write( @@ -829,16 +869,26 @@ def resolve_and_compute(batch): updates = matched.map_batches( resolve_and_compute, batch_format='pyarrow') - metrics = merge_into( - target=name, - source=updates, - catalog_options=self.catalog_options, - on=['id'], - when_matched=[ - WhenMatched.update({'feature': source_col('new_feature')}) - ], - num_partitions=num_partitions, - ) + real_materialize = ray.data.Dataset.materialize + materialized = [] + + def track_materialize(dataset): + materialized.append(dataset) + return real_materialize(dataset) + + with patch.object( + ray.data.Dataset, 'materialize', new=track_materialize, + ): + metrics = merge_into( + target=name, + source=updates, + catalog_options=self.catalog_options, + on=['id'], + when_matched=[ + WhenMatched.update({'feature': source_col('new_feature')}) + ], + num_partitions=num_partitions, + ) table = self.catalog.get_table(name) rb = table.new_read_builder() @@ -848,6 +898,7 @@ def resolve_and_compute(batch): self.assertEqual(out['feature'], [200, 20, 400]) self.assertEqual(out['payload'], [b'aa', b'bbb', b'cccc']) self.assertEqual(metrics['num_matched'], 2) + self.assertEqual(len(materialized), 1) def test_combined_writes_single_snapshot(self): target = self._create_table() @@ -2084,13 +2135,18 @@ def test_self_merge_update_literal(self): ), ) - result = merge_into( - target=target, - source=target, - catalog_options=self.catalog_options, - on=['_ROW_ID'], - when_matched=[WhenMatched.update({'age': lit(99)})], - ) + with patch.object( + ray.data.Dataset, + 'materialize', + side_effect=AssertionError('self-merge must stay streaming'), + ): + result = merge_into( + target=target, + source=target, + catalog_options=self.catalog_options, + on=['_ROW_ID'], + when_matched=[WhenMatched.update({'age': lit(99)})], + ) self.assertEqual(result['num_matched'], 3) out = self._read_sorted(target)