This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 5ddd34a96f [python][ray] Fix merge hang with shuffled sources on small
clusters (#9360)
5ddd34a96f is described below
commit 5ddd34a96faf20e44da77744f9c5a35177921ad5
Author: XiaoHongbo <[email protected]>
AuthorDate: Mon Oct 5 22:48:10 2026 +0800
[python][ray] Fix merge hang with shuffled sources on small clusters (#9360)
---
.../pypaimon/ray/data_evolution_merge_into.py | 40 ++++++
.../pypaimon/ray/data_evolution_merge_join.py | 9 ++
.../tests/ray_data_evolution_merge_into_test.py | 157 +++++++++++++++------
3 files changed, 166 insertions(+), 40 deletions(-)
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
index 3a280dd028..6fb9e34b99 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
@@ -97,6 +97,9 @@ def merge_into(
list(when_matched), list(when_not_matched), on,
read_columns,
)
+ materialize_matched_before_routing = (
+ _has_upstream_join_or_shuffle(source_ds)
+ )
base_snapshot = table.snapshot_manager().get_latest_snapshot()
target_empty = _is_target_empty(base_snapshot)
estimated_size_bytes = None
@@ -153,6 +156,7 @@ def merge_into(
ray_remote_args, concurrency,
update_num_partitions=update_num_partitions,
delete_num_partitions=delete_num_partitions,
+ materialize_matched_before_routing=materialize_matched_before_routing,
)
@@ -483,6 +487,7 @@ def _execute_and_commit(
ray_remote_args, concurrency,
update_num_partitions=None,
delete_num_partitions=None,
+ materialize_matched_before_routing=False,
):
collect_action_row_ids = update_ds is not None and delete_ds is not None
commit_messages: list = []
@@ -529,6 +534,9 @@ def _execute_and_commit(
if base_snapshot is not None else None
),
collect_row_ids=collect_action_row_ids,
+ materialize_before_routing=(
+ materialize_matched_before_routing
+ ),
)
)
commit_messages.extend(update_msgs)
@@ -543,6 +551,9 @@ def _execute_and_commit(
if base_snapshot is not None else None
),
collect_row_ids=collect_action_row_ids,
+ materialize_before_routing=(
+ materialize_matched_before_routing
+ ),
)
commit_messages.extend(delete_msgs)
@@ -823,6 +834,35 @@ def _normalize_source(
)
+def _has_upstream_join_or_shuffle(dataset) -> bool:
+ """Best-effort check for a lazy source join or shuffle."""
+ try:
+ logical_plan = getattr(dataset, "_logical_plan", None)
+ operator = getattr(logical_plan, "dag", None)
+ pending = [operator] if operator is not None else []
+ visited = set()
+ while pending:
+ operator = pending.pop()
+ if id(operator) in visited:
+ continue
+ visited.add(id(operator))
+ operator_name = type(operator).__name__
+ if operator_name in {
+ "Aggregate", "Join", "RandomShuffle", "Sort",
+ }:
+ return True
+ if operator_name == "Repartition":
+ shuffle = getattr(
+ operator, "shuffle", getattr(operator, "_shuffle", False)
+ )
+ if shuffle:
+ return True
+ pending.extend(getattr(operator, "input_dependencies", ()))
+ except Exception:
+ pass
+ return False
+
+
def _source_schema_or_raise(source_ds):
"""Get source schema; refuse to proceed if Ray can't tell us the
columns."""
schema = source_ds.schema()
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index 3cccc3dafb..01d4f4228d 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -689,6 +689,7 @@ def distributed_update_apply(
estimated_size_bytes: Optional[int] = None,
estimated_num_rows: Optional[int] = None,
data_context=None,
+ materialize_before_routing: bool = False,
) -> Tuple[list, int, list]:
import numpy as np
import pickle
@@ -728,6 +729,9 @@ def distributed_update_apply(
len(sorted_first_row_ids),
data_context=data_context,
)
+ 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.
@@ -1009,6 +1013,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
@@ -1032,6 +1037,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 f30bd443ca..cea154d5e4 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
@@ -543,14 +543,19 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
schema=self.pa_schema,
)
- merge_into(
- target=target,
- source=source,
- catalog_options=self.catalog_options,
- on=['id'],
- when_matched=[WhenMatched.update('*')],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ with patch.object(
+ ray.data.Dataset,
+ 'materialize',
+ side_effect=AssertionError('plain source must stay streaming'),
+ ):
+ merge_into(
+ target=target,
+ source=source,
+ catalog_options=self.catalog_options,
+ on=['id'],
+ when_matched=[WhenMatched.update('*')],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
out = self._read_sorted(target)
self.assertEqual(out['id'], [1, 2, 3])
@@ -573,21 +578,26 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
),
)
- metrics = merge_into(
- target=target,
- source=pa.Table.from_pydict(
- {
- 'id': pa.array([2, 3], type=pa.int32()),
- 'name': ['ignored', 'ignored'],
- 'age': pa.array([99, 99], type=pa.int32()),
- },
- schema=self.pa_schema,
- ),
- catalog_options=self.catalog_options,
- on=['id'],
- when_matched=[WhenMatched.delete()],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ with patch.object(
+ ray.data.Dataset,
+ 'materialize',
+ side_effect=AssertionError('plain source must stay streaming'),
+ ):
+ metrics = merge_into(
+ target=target,
+ source=pa.Table.from_pydict(
+ {
+ 'id': pa.array([2, 3], type=pa.int32()),
+ 'name': ['ignored', 'ignored'],
+ 'age': pa.array([99, 99], type=pa.int32()),
+ },
+ schema=self.pa_schema,
+ ),
+ catalog_options=self.catalog_options,
+ on=['id'],
+ when_matched=[WhenMatched.delete()],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(metrics, {
'num_matched': 2, 'num_inserted': 0, 'num_unchanged': 0,
@@ -597,6 +607,57 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
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_source_shuffle_detection(self):
+ from pypaimon.ray import data_evolution_merge_into as merge_module
+
+ source = ray.data.from_arrow(pa.table({'id': [1, 2]}))
+ detect = merge_module._has_upstream_join_or_shuffle
+
+ self.assertFalse(detect(source))
+ self.assertFalse(detect(source.repartition(2, shuffle=False)))
+ self.assertTrue(detect(source.repartition(2, shuffle=True)))
+ self.assertTrue(detect(source.random_shuffle()))
+
def test_not_matched_insert_appends_unmatched(self):
target = self._create_table()
self._write(
@@ -1095,16 +1156,26 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
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()
@@ -1114,6 +1185,7 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
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()
@@ -2350,13 +2422,18 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
),
)
- 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)
