[python][ray] Avoid MERGE shuffle starvation on small clusters

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)
 

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(

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)