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 ddddf73dbb [python][ray] Avoid routing shuffle in self-merge updates
(#9335)
ddddf73dbb is described below
commit ddddf73dbbea8463e0d739de97a451de1f16ea14
Author: XiaoHongbo <[email protected]>
AuthorDate: Fri Aug 21 15:43:56 2026 +0800
[python][ray] Avoid routing shuffle in self-merge updates (#9335)
---
.../pypaimon/ray/data_evolution_merge_into.py | 46 +-
.../pypaimon/ray/data_evolution_merge_join.py | 267 +++++++++--
.../tests/ray_data_evolution_merge_into_test.py | 501 +++++++++++++--------
3 files changed, 573 insertions(+), 241 deletions(-)
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
index 3f8e3f3d16..6124ebefd5 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py
@@ -32,8 +32,10 @@ from pypaimon.ray.data_evolution_merge_join import (
build_matched_update_ds,
build_not_matched_insert_ds,
build_self_merge_delete_ds,
- build_self_merge_update_ds,
+ _SelfMergeUpdatePlan,
+ build_self_merge_update_plan,
distributed_delete_apply,
+ distributed_self_merge_update_apply,
distributed_update_apply,
distributed_write_collect_msgs,
)
@@ -92,7 +94,7 @@ def merge_into(
base_snapshot = table.snapshot_manager().get_latest_snapshot()
update_ds, delete_ds, insert_ds, update_cols_union = _build_datasets(
- target, source_ds, matched_specs, not_matched_specs,
+ table, target, source_ds, matched_specs, not_matched_specs,
ctx, base_snapshot, num_partitions, ray_remote_args,
)
@@ -323,7 +325,7 @@ def _is_self_merge(target, source, target_on_cols,
source_on_cols) -> bool:
def _build_datasets(
- target, source_ds, matched_specs, not_matched_specs,
+ table, target, source_ds, matched_specs, not_matched_specs,
ctx: "_PrepareCtx", base_snapshot, num_partitions, ray_remote_args,
):
# Pin every target read to base_snapshot so all branches see the same
@@ -340,18 +342,16 @@ def _build_datasets(
if matched_specs and base_snapshot is not None:
update_cols_union = _union_update_cols(matched_specs)
if update_cols_union:
- update_ds = build_self_merge_update_ds(
- target_identifier=target,
+ update_ds = build_self_merge_update_plan(
+ table=table,
clauses=matched_specs,
target_field_names=ctx.full_target_field_names,
target_pa_schema=ctx.update_pa_schema,
update_cols=update_cols_union,
- catalog_options=ctx.catalog_options,
resolve_target_projection=_resolve_target_projection,
snapshot_id=base_snapshot_id,
scan_predicate=ctx.self_merge_scan_predicate,
read_columns=ctx.read_columns,
- ray_remote_args=ray_remote_args,
)
if any(c.delete for c in matched_specs):
delete_ds = build_self_merge_delete_ds(
@@ -439,16 +439,28 @@ def _execute_and_commit(
try:
if update_ds is not None:
- update_msgs, num_updated, update_row_ids =
distributed_update_apply(
- update_ds, table, update_cols_union,
- num_partitions=num_partitions,
- ray_remote_args=ray_remote_args,
- base_snapshot_id=(
- base_snapshot.id
- if base_snapshot is not None else None
- ),
- collect_row_ids=collect_action_row_ids,
- )
+ if isinstance(update_ds, _SelfMergeUpdatePlan):
+ update_msgs, num_updated, update_row_ids = (
+ distributed_self_merge_update_apply(
+ update_ds,
+ num_partitions=num_partitions,
+ ray_remote_args=ray_remote_args,
+ collect_row_ids=collect_action_row_ids,
+ )
+ )
+ else:
+ update_msgs, num_updated, update_row_ids = (
+ distributed_update_apply(
+ update_ds, table, update_cols_union,
+ num_partitions=num_partitions,
+ ray_remote_args=ray_remote_args,
+ base_snapshot_id=(
+ base_snapshot.id
+ if base_snapshot is not None else None
+ ),
+ collect_row_ids=collect_action_row_ids,
+ )
+ )
commit_messages.extend(update_msgs)
if delete_ds is not None:
diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
index d7e381adc4..af56d7ca30 100644
--- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py
+++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py
@@ -16,6 +16,7 @@
# limitations under the License.
################################################################################
+from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple
import pyarrow as pa
@@ -44,6 +45,39 @@ def _map_kwargs(
return kwargs
+@dataclass(frozen=True)
+class _SelfMergeUpdatePlan:
+ """Pinned target file groups for self-merge update execution."""
+
+ table: Any
+ scan_table: Any
+ file_groups: list
+ predicate: Any
+ read_type: list
+ clauses: List[_NormalizedClause]
+ update_cols: List[str]
+ update_schema: pa.Schema
+ row_id_name: str
+ snapshot_id: int
+ callable_input_columns: Optional[List[str]]
+
+
+@dataclass(frozen=True)
+class _SelfMergeUpdateContext:
+ """Worker state shared by all file groups."""
+
+ table: Any
+ scan_table: Any
+ predicate: Any
+ read_type: list
+ clauses: List[_NormalizedClause]
+ update_cols: List[str]
+ update_schema: pa.Schema
+ row_id_name: str
+ snapshot_id: int
+ callable_input_columns: Optional[List[str]]
+
+
def _resolve_source_projection(
clauses: List[_NormalizedClause],
source_on: Sequence[str],
@@ -189,22 +223,23 @@ def _build_matched_delete_transform(
return _transform
-def build_self_merge_update_ds(
+def build_self_merge_update_plan(
*,
- target_identifier: str,
+ table,
clauses: List[_NormalizedClause],
target_field_names: Sequence[str],
target_pa_schema: pa.Schema,
update_cols: Sequence[str],
- catalog_options: Dict[str, str],
resolve_target_projection,
snapshot_id: Optional[int] = None,
scan_predicate=None,
read_columns: Sequence[str] = (),
- ray_remote_args: Optional[Dict[str, Any]] = None,
-) -> Tuple:
- from pypaimon.ray.ray_paimon import read_paimon
+) -> _SelfMergeUpdatePlan:
+ from pypaimon.common.options.core_options import (
+ CoreOptions, GlobalIndexSearchMode,
+ )
from pypaimon.table.special_fields import SpecialFields
+ from pypaimon.write.table_update import TableUpdate
row_id_name = SpecialFields.ROW_ID.name
needed_cols = set(resolve_target_projection(
@@ -227,59 +262,195 @@ def build_self_merge_update_ds(
c for c in target_field_names if c in needed_cols
]
- read_kwargs = {}
+ dynamic_options = {}
+ if snapshot_id is not None:
+ dynamic_options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id)
if scan_predicate is not None:
- from pypaimon.common.options.core_options import (
- CoreOptions, GlobalIndexSearchMode,
+ dynamic_options[CoreOptions.SCALAR_INDEX_SEARCH_MODE.key()] = (
+ GlobalIndexSearchMode.FULL.value
)
- read_kwargs["filter"] = scan_predicate
- read_kwargs["dynamic_options"] = {
- CoreOptions.SCALAR_INDEX_SEARCH_MODE.key():
- GlobalIndexSearchMode.FULL.value,
- }
- target_ds = read_paimon(
- target_identifier, catalog_options,
- projection=projection, snapshot_id=snapshot_id,
- _preserve_current_schema=True,
- **read_kwargs,
- )
- update_schema = build_update_schema(target_pa_schema, update_cols,
row_id_name)
-
- orig_names = target_ds.schema().names
- target_renamed = target_ds.rename_columns(
- {c: f"t.{c}" for c in orig_names}
+ scan_table = (
+ table.copy_without_time_travel(dynamic_options)
+ if dynamic_options else table
)
-
- def _add_source_aliases(batch: pa.Table) -> pa.Table:
- columns = list(batch.columns)
- names = list(batch.schema.names)
- for orig in orig_names:
- if orig == row_id_name:
- continue
- t_col_name = f"t.{orig}"
- if t_col_name in names:
- idx = names.index(t_col_name)
- columns.append(columns[idx])
- names.append(f"s.{orig}")
- return pa.table(columns, names=names)
-
- aliased = target_renamed.map_batches(
- _add_source_aliases, **_map_kwargs(ray_remote_args),
+ read_builder = scan_table.new_read_builder().with_projection(projection)
+ if scan_predicate is not None:
+ read_builder.with_filter(scan_predicate)
+ scan_plan = read_builder.new_scan().plan_for_write()
+ planned_snapshot_id = (
+ scan_plan.snapshot_id
+ if scan_plan.snapshot_id is not None else -1
)
-
- _transform = _build_matched_transform(
- clauses,
- on_map={row_id_name: row_id_name},
- on_pairs=[(row_id_name, row_id_name)],
+ # A packed scan split may contain multiple logical row-id groups. Keep
+ # each group intact, but do not materialize unrelated groups together.
+ file_groups = list(TableUpdate._predicate_update_file_groups(
+ scan_plan.splits()
+ ))
+ update_schema = build_update_schema(target_pa_schema, update_cols,
row_id_name)
+ return _SelfMergeUpdatePlan(
+ table=table,
+ scan_table=scan_table,
+ file_groups=file_groups,
+ predicate=scan_predicate,
+ read_type=read_builder.read_type(),
+ clauses=clauses,
update_cols=list(update_cols),
- row_id_name=row_id_name,
update_schema=update_schema,
+ row_id_name=row_id_name,
+ snapshot_id=planned_snapshot_id,
callable_input_columns=(
list(read_columns) + [row_id_name]
if read_columns else None
),
)
- return aliased.map_batches(_transform, **_map_kwargs(ray_remote_args))
+
+
+def _self_merge_aliases(batch: pa.Table, row_id_name: str) -> pa.Table:
+ columns = []
+ names = []
+ for name, column in zip(batch.schema.names, batch.columns):
+ columns.append(column)
+ names.append(f"t.{name}")
+ if name != row_id_name:
+ columns.append(column)
+ names.append(f"s.{name}")
+ return pa.table(columns, names=names)
+
+
+def _apply_self_merge_update_group(context, file_group, collect_row_ids):
+ """Read, transform, and stage one complete first-row-id file group."""
+ from pypaimon.read.table_read import TableRead
+ from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
+ from pypaimon.write.file_store_commit import _abort_commit_messages
+ from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+ table_read = TableRead(
+ context.scan_table,
+ context.predicate,
+ context.read_type,
+ )
+ transform = _build_matched_transform(
+ context.clauses,
+ on_map={context.row_id_name: context.row_id_name},
+ on_pairs=[(context.row_id_name, context.row_id_name)],
+ update_cols=context.update_cols,
+ row_id_name=context.row_id_name,
+ update_schema=context.update_schema,
+ callable_input_columns=context.callable_input_columns,
+ )
+ update_parts = []
+ batch_reader = table_read.to_arrow_batch_reader([file_group])
+ for batch in iter(batch_reader.read_next_batch, None):
+ if batch.num_rows == 0:
+ continue
+ matched = pa.Table.from_batches([batch])
+ updates = transform(_self_merge_aliases(
+ matched, context.row_id_name,
+ ))
+ if updates.num_rows > 0:
+ update_parts.append(updates)
+
+ if not update_parts:
+ return [], 0, []
+ updates = (
+ update_parts[0]
+ if len(update_parts) == 1 else pa.concat_tables(update_parts)
+ )
+
+ row_ids = (
+ updates.column(context.row_id_name).to_pylist()
+ if collect_row_ids else []
+ )
+
+ import uuid
+ files_info = TableUpdateByRowId._files_info_from_splits(
+ context.snapshot_id, [file_group],
+ )
+ updater = TableUpdateByRowId(
+ context.table,
+ "_self_merge_group_" + uuid.uuid4().hex[:8],
+ BATCH_COMMIT_IDENTIFIER,
+ _precomputed_files_info=files_info,
+ )
+ try:
+ messages = updater.update_columns(updates, context.update_cols)
+ except Exception:
+ _abort_commit_messages(context.table, updater.commit_messages)
+ raise
+ return messages, updates.num_rows, row_ids
+
+
+def distributed_self_merge_update_apply(
+ plan: _SelfMergeUpdatePlan,
+ *,
+ num_partitions: int,
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+ collect_row_ids: bool = False,
+) -> Tuple[list, int, list]:
+ """Stage self-merge updates in scan tasks without a routing shuffle."""
+ import ray
+
+ if not plan.file_groups:
+ return [], 0, []
+
+ remote_args = dict(ray_remote_args or {})
+ apply_remote = (
+ ray.remote(**remote_args)(_apply_self_merge_update_group)
+ if remote_args else ray.remote(_apply_self_merge_update_group)
+ )
+ context = ray.put(_SelfMergeUpdateContext(
+ table=plan.table,
+ scan_table=plan.scan_table,
+ predicate=plan.predicate,
+ read_type=plan.read_type,
+ clauses=plan.clauses,
+ update_cols=plan.update_cols,
+ update_schema=plan.update_schema,
+ row_id_name=plan.row_id_name,
+ snapshot_id=plan.snapshot_id,
+ callable_input_columns=plan.callable_input_columns,
+ ))
+ group_iter = iter(plan.file_groups)
+ max_in_flight = max(1, min(num_partitions, len(plan.file_groups)))
+ pending = set()
+ messages = []
+ num_updated = 0
+ row_ids = []
+ first_error = None
+
+ def submit_next():
+ try:
+ file_group = next(group_iter)
+ except StopIteration:
+ return False
+ pending.add(apply_remote.remote(
+ context, file_group, collect_row_ids
+ ))
+ return True
+
+ for _ in range(max_in_flight):
+ submit_next()
+
+ while pending:
+ ready, remaining = ray.wait(list(pending), num_returns=1)
+ pending = set(remaining)
+ ref = ready[0]
+ try:
+ split_messages, split_count, split_row_ids = ray.get(ref)
+ messages.extend(split_messages)
+ num_updated += split_count
+ row_ids.extend(split_row_ids)
+ except Exception as error:
+ if first_error is None:
+ first_error = error
+ if first_error is None:
+ submit_next()
+
+ if first_error is not None:
+ from pypaimon.write.file_store_commit import _abort_commit_messages
+ _abort_commit_messages(plan.table, messages)
+ raise first_error
+ return messages, num_updated, row_ids
def build_self_merge_delete_ds(
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 e8a66a6061..8e48773d4b 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
@@ -114,6 +114,26 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
snap = table.snapshot_manager().get_latest_snapshot()
return snap.id if snap is not None else None
+ def _merge_and_capture_self_merge_plan(self, **kwargs):
+ from pypaimon.ray.data_evolution_merge_join import (
+ build_self_merge_update_plan as real_build_plan,
+ )
+
+ captured = {}
+
+ def capture(**plan_kwargs):
+ plan = real_build_plan(**plan_kwargs)
+ captured['plan'] = plan
+ return plan
+
+ with patch(
+ 'pypaimon.ray.data_evolution_merge_into.'
+ 'build_self_merge_update_plan',
+ side_effect=capture,
+ ):
+ result = merge_into(**kwargs)
+ return result, captured['plan']
+
def test_paimon_source_table_pins_snapshot(self):
from pypaimon.ray import data_evolution_merge_into as m
@@ -2038,10 +2058,184 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
self.assertEqual(out['age'], [99, 99, 99])
self.assertEqual(out['name'], ['a', 'b', 'c'])
+ def test_self_merge_update_bypasses_routing_shuffle(self):
+ options = dict(self.de_options)
+ options.update({
+ 'source.split.target-size': '1gb',
+ 'source.split.open-file-cost': '1b',
+ })
+ target = self._create_table(options=options)
+ self._write(target, self._source(ids=(1, 2)))
+ self._write(target, self._source(ids=(3, 4)))
+ table = self.catalog.get_table(target)
+ packed_splits =
table.new_read_builder().new_scan().plan_for_write().splits()
+ self.assertEqual(len(packed_splits), 1)
+
+ with patch.object(
+ ray.data.Dataset,
+ 'groupby',
+ side_effect=AssertionError('routing shuffle is not allowed'),
+ ):
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update({'age': lit(99)})],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
+
+ self.assertEqual(len(plan.file_groups), 2)
+ self.assertEqual(result['num_matched'], 4)
+ self.assertEqual(self._read_sorted(target)['age'], [99, 99, 99, 99])
+
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
- def test_self_merge_callable_assignment(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
+ def test_self_merge_filters_file_group_in_batches(self):
+ from pypaimon.ray import data_evolution_merge_into as merge_module
+ from pypaimon.ray import data_evolution_merge_join as join_module
+ from pypaimon.write.file_store_commit import _abort_commit_messages
+
+ options = dict(self.de_options)
+ options['read.batch-size'] = '2'
+ target = self._create_table(options=options)
+ self._write(
+ target,
+ pa.Table.from_pydict(
+ {
+ 'id': pa.array([1, 2, 3, 4, 5, 6], type=pa.int32()),
+ 'name': ['a', 'b', 'c', 'd', 'e', 'f'],
+ 'age': pa.array([0, 2, 0, 4, 0, 6], type=pa.int32()),
+ },
+ schema=self.pa_schema,
+ ),
+ )
+ clauses = [WhenMatched.update(
+ {'age': lit(99)}, condition='t.age = t.id',
+ )]
+ table, source_ds, matched, not_matched, ctx = merge_module._prepare(
+ target,
+ target,
+ self.catalog_options,
+ clauses,
+ [],
+ ['_ROW_ID'],
+ )
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ plan, _, _, _ = merge_module._build_datasets(
+ table,
+ target,
+ source_ds,
+ matched,
+ not_matched,
+ ctx,
+ snapshot,
+ _TEST_NUM_PARTITIONS,
+ None,
+ )
+ self.assertIsNone(plan.predicate)
+ self.assertEqual(len(plan.file_groups), 1)
+
+ context = join_module._SelfMergeUpdateContext(
+ table=plan.table,
+ scan_table=plan.scan_table,
+ predicate=plan.predicate,
+ read_type=plan.read_type,
+ clauses=plan.clauses,
+ update_cols=plan.update_cols,
+ update_schema=plan.update_schema,
+ row_id_name=plan.row_id_name,
+ snapshot_id=plan.snapshot_id,
+ callable_input_columns=plan.callable_input_columns,
+ )
+ batch_sizes = []
+ real_build_transform = join_module._build_matched_transform
+
+ def tracked_build_transform(*args, **kwargs):
+ real_transform = real_build_transform(*args, **kwargs)
+
+ def tracked_transform(batch):
+ batch_sizes.append(batch.num_rows)
+ return real_transform(batch)
+
+ return tracked_transform
+
+ messages = []
+ try:
+ with patch.object(
+ join_module,
+ '_build_matched_transform',
+ side_effect=tracked_build_transform,
+ ):
+ messages, count, row_ids = (
+ join_module._apply_self_merge_update_group(
+ context, plan.file_groups[0], True,
+ )
+ )
+
+ self.assertEqual(batch_sizes, [2, 2, 2])
+ self.assertEqual(count, 3)
+ self.assertEqual(row_ids, [1, 3, 5])
+ finally:
+ _abort_commit_messages(table, messages)
+
+ def test_self_merge_update_aborts_other_groups_after_failure(self):
+ from pypaimon.ray import data_evolution_merge_into as merge_module
+ from pypaimon.ray.data_evolution_merge_join import (
+ distributed_self_merge_update_apply,
+ )
+
+ options = dict(self.de_options)
+ options.update({
+ 'source.split.target-size': '1b',
+ 'source.split.open-file-cost': '1b',
+ })
+ target = self._create_table(options=options)
+ self._write(target, self._source(ids=(1, 2)))
+ self._write(target, self._source(ids=(3, 4)))
+
+ table, source_ds, matched, not_matched, ctx = merge_module._prepare(
+ target,
+ target,
+ self.catalog_options,
+ [WhenMatched.update({'age': lit(99)})],
+ [],
+ ['_ROW_ID'],
+ )
+ snapshot = table.snapshot_manager().get_latest_snapshot()
+ plan, _, _, _ = merge_module._build_datasets(
+ table,
+ target,
+ source_ds,
+ matched,
+ not_matched,
+ ctx,
+ snapshot,
+ _TEST_NUM_PARTITIONS,
+ None,
+ )
+ self.assertGreaterEqual(len(plan.file_groups), 2)
+
+ before = set()
+ for root, _, files in os.walk(self.warehouse):
+ before.update(os.path.join(root, name) for name in files)
+
+ for data_file in plan.file_groups[-1].files:
+ data_file.file_path += '.missing'
+
+ with self.assertRaises(Exception):
+ distributed_self_merge_update_apply(
+ plan,
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
+
+ after = set()
+ for root, _, files in os.walk(self.warehouse):
+ after.update(os.path.join(root, name) for name in files)
+ self.assertEqual(before, after)
+
+ @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
+ def test_self_merge_callable_assignment(self):
target = self._create_table()
self._write(
target,
@@ -2060,11 +2254,12 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
raise AssertionError(rows.column_names)
return pc.add(rows['age'], 1)
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
+ with patch.object(
+ ray.data.Dataset,
+ 'groupby',
+ side_effect=AssertionError('routing shuffle is not allowed'),
+ ):
+ result, plan = self._merge_and_capture_self_merge_plan(
target=target,
source=target,
catalog_options=self.catalog_options,
@@ -2090,8 +2285,10 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
},
)
self.assertEqual(
- mock_read.call_args[1]['projection'], ['_ROW_ID', 'id', 'age']
+ [field.name for field in plan.read_type],
+ ['_ROW_ID', 'id', 'age'],
)
+ self.assertEqual(plan.callable_input_columns, ['age', '_ROW_ID'])
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_callable_updates_variant_path(self):
@@ -2290,7 +2487,6 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
from pypaimon.common.options.core_options import (
CoreOptions, GlobalIndexSearchMode,
)
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
target = self._create_table()
self._write(
@@ -2304,20 +2500,16 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
schema=self.pa_schema,
),
)
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'age': lit(99)}, condition='t.id IN (1, 3)',
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'age': lit(99)}, condition='t.id IN (1, 3)',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 2)
self.assertEqual(
@@ -2328,18 +2520,23 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
'age': [99, 20, 99],
},
)
- read_kwargs = mock_read.call_args[1]
- predicate = read_kwargs['filter']
+ predicate = plan.predicate
self.assertEqual(predicate.method, 'in')
self.assertEqual(predicate.field, 'id')
self.assertEqual(predicate.literals, [1, 3])
self.assertEqual(
- read_kwargs['dynamic_options'][
- CoreOptions.SCALAR_INDEX_SEARCH_MODE.key()
- ],
+ plan.scan_table.table_schema.options.get(
+ CoreOptions.SCALAR_INDEX_SEARCH_MODE.key()),
GlobalIndexSearchMode.FULL.value,
)
- self.assertTrue(read_kwargs['_preserve_current_schema'])
+ self.assertEqual(
+ plan.scan_table.table_schema.fields,
+ plan.table.table_schema.fields,
+ )
+ self.assertEqual(
+ plan.scan_table.table_schema.id,
+ plan.table.table_schema.id,
+ )
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_pushdown_handles_evolved_file_groups(self):
@@ -2462,33 +2659,27 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_multiple_conditions_push_down_or(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
target = self._create_table()
self._write(target, self._source(ids=(1, 2, 3)))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[
- WhenMatched.update(
- {'age': lit(11)}, condition='t.id = 1',
- ),
- WhenMatched.update(
- {'age': lit(33)}, condition='s.id = 3',
- ),
- ],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[
+ WhenMatched.update(
+ {'age': lit(11)}, condition='t.id = 1',
+ ),
+ WhenMatched.update(
+ {'age': lit(33)}, condition='s.id = 3',
+ ),
+ ],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 2)
- predicate = mock_read.call_args[1]['filter']
+ predicate = plan.predicate
self.assertEqual(predicate.method, 'or')
self.assertEqual(
[(p.field, p.literals) for p in predicate.literals],
@@ -2497,63 +2688,49 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_unconditional_clause_disables_pushdown(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
target = self._create_table()
self._write(target, self._source(ids=(1, 2, 3)))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[
- WhenMatched.update(
- {'age': lit(11)}, condition='t.id = 1',
- ),
- WhenMatched.update({'age': lit(99)}),
- ],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[
+ WhenMatched.update(
+ {'age': lit(11)}, condition='t.id = 1',
+ ),
+ WhenMatched.update({'age': lit(99)}),
+ ],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 3)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_column_comparison_fails_open(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
target = self._create_table()
self._write(target, self._source(ids=(1, 2, 3)))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'name': lit('same')}, condition='t.age = s.age',
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'name': lit('same')}, condition='t.age = s.age',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 3)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
self.assertEqual(self._read_sorted(target)['name'],
['same', 'same', 'same'])
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_pushdown_preserves_field_case(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
case_schema = pa.schema([
('UserID', pa.int32()),
('Value', pa.int32()),
@@ -2568,24 +2745,20 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
schema=case_schema,
))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'Value': lit(99)},
- condition='t.UserID IN (1, 3)',
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'Value': lit(99)},
+ condition='t.UserID IN (1, 3)',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 2)
- predicate = mock_read.call_args[1]['filter']
+ predicate = plan.predicate
self.assertEqual(predicate.field, 'UserID')
table = self.catalog.get_table(target)
read_builder = table.new_read_builder()
@@ -2595,8 +2768,6 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_date_condition_fails_open(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
date_schema = pa.schema([
('id', pa.int32()),
('event_date', pa.date32()),
@@ -2616,24 +2787,20 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
'value': [10, 20],
}, schema=date_schema))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'value': lit(99)},
- condition="t.event_date = '2026-01-01'",
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'value': lit(99)},
+ condition="t.event_date = '2026-01-01'",
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 1)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
table = self.catalog.get_table(target)
read_builder = table.new_read_builder()
splits = read_builder.new_scan().plan().splits()
@@ -2642,8 +2809,6 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_double_condition_fails_open(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
double_schema = pa.schema([
('id', pa.int32()),
('metric', pa.float64()),
@@ -2660,23 +2825,19 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
'value': [10, 20, 30],
}, schema=double_schema))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'value': lit(99)}, condition='t.metric > 0',
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'value': lit(99)}, condition='t.metric > 0',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 2)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
table = self.catalog.get_table(target)
read_builder = table.new_read_builder()
splits = read_builder.new_scan().plan().splits()
@@ -2685,8 +2846,6 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_like_condition_fails_open(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
string_schema = pa.schema([
('id', pa.int32()),
('text', pa.string()),
@@ -2703,23 +2862,19 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
'value': [10, 20, 30, 40, 50],
}, schema=string_schema))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'value': lit(99)}, condition=r"t.text LIKE '%\n%'",
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'value': lit(99)}, condition=r"t.text LIKE '%\n%'",
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 4)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
table = self.catalog.get_table(target)
read_builder = table.new_read_builder()
splits = read_builder.new_scan().plan().splits()
@@ -2728,8 +2883,6 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
@unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
def test_self_merge_out_of_range_integer_fails_open(self):
- from pypaimon.ray.ray_paimon import read_paimon as real_read_paimon
-
int_schema = pa.schema([
('id', pa.int64()),
('value', pa.int32()),
@@ -2744,24 +2897,20 @@ class RayDataEvolutionMergeIntoTest(unittest.TestCase):
'value': [10, 20, 30],
}, schema=int_schema))
- with patch(
- 'pypaimon.ray.ray_paimon.read_paimon',
- wraps=real_read_paimon,
- ) as mock_read:
- result = merge_into(
- target=target,
- source=target,
- catalog_options=self.catalog_options,
- on=['_ROW_ID'],
- when_matched=[WhenMatched.update(
- {'value': lit(99)},
- condition='t.id < 9223372036854775808',
- )],
- num_partitions=_TEST_NUM_PARTITIONS,
- )
+ result, plan = self._merge_and_capture_self_merge_plan(
+ target=target,
+ source=target,
+ catalog_options=self.catalog_options,
+ on=['_ROW_ID'],
+ when_matched=[WhenMatched.update(
+ {'value': lit(99)},
+ condition='t.id < 9223372036854775808',
+ )],
+ num_partitions=_TEST_NUM_PARTITIONS,
+ )
self.assertEqual(result['num_matched'], 3)
- self.assertNotIn('filter', mock_read.call_args[1])
+ self.assertIsNone(plan.predicate)
table = self.catalog.get_table(target)
read_builder = table.new_read_builder()
splits = read_builder.new_scan().plan().splits()
