| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ################################################################################ |
| |
| from dataclasses import dataclass |
| from typing import Any, Dict, List, Optional, Sequence, Tuple |
| |
| import pyarrow as pa |
| |
| from pypaimon.ray.data_evolution_merge_transform import ( |
| SourceColumnRef, |
| _NormalizedClause, |
| build_delete_schema, |
| build_update_schema, |
| cast_to_schema, |
| vectorized_delete_transform, |
| vectorized_insert_transform, |
| vectorized_matched_transform, |
| ) |
| from pypaimon.ray.partitioning import ( |
| _default_hash_shuffle_parallelism, |
| _resolve_num_partitions, |
| _resolve_row_id_num_partitions, |
| ) |
| |
| |
| def _map_kwargs( |
| ray_remote_args: Optional[Dict[str, Any]], |
| ) -> Dict[str, Any]: |
| """Build kwargs for map_batches/map_groups; spread ray_remote_args because |
| those APIs take remote options as **kwargs, not under a 'ray_remote_args' |
| key.""" |
| kwargs: Dict[str, Any] = {"batch_format": "pyarrow"} |
| if ray_remote_args: |
| kwargs.update(ray_remote_args) |
| return kwargs |
| |
| |
| def _resolve_matched_num_partitions( |
| num_partitions: Optional[int], |
| estimated_size_bytes: Optional[int], |
| target_ds, |
| ) -> int: |
| """Resolve a target-left join from the context sealed on its left input.""" |
| if num_partitions is not None: |
| return num_partitions |
| |
| data_context = getattr(target_ds, "context", None) |
| default_shuffle = _default_hash_shuffle_parallelism(data_context) |
| return _resolve_num_partitions( |
| num_partitions, |
| estimated_size_bytes, |
| min_partitions=default_shuffle, |
| unknown_num_partitions=default_shuffle, |
| data_context=data_context, |
| ) |
| |
| |
| @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], |
| source_field_names: Sequence[str], |
| ) -> list: |
| needed = set(source_on) |
| source_set = set(source_field_names) |
| |
| for clause in clauses: |
| for value in clause.spec.values(): |
| if isinstance(value, SourceColumnRef): |
| needed.add(value.column) |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import extract_columns |
| for ref in extract_columns(clause.condition): |
| prefix, col = ref.split(".", 1) |
| if prefix == "s" and col in source_set: |
| needed.add(col) |
| |
| return [c for c in source_field_names if c in needed] |
| |
| |
| def _build_matched_transform( |
| clauses: List[_NormalizedClause], |
| on_map: Dict[str, str], |
| on_pairs: List[Tuple[str, str]], |
| update_cols: List[str], |
| row_id_name: str, |
| update_schema: pa.Schema, |
| callable_input_columns: Optional[Sequence[str]] = None, |
| ): |
| prepared_clauses = [] |
| for clause in clauses: |
| rewritten = None |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import ( |
| remap_source_on_keys, rewrite_condition, |
| ) |
| rewritten = remap_source_on_keys( |
| rewrite_condition(clause.condition), on_map, |
| ) |
| prepared_clauses.append((clause.spec, rewritten, clause.delete)) |
| |
| _filter_batch = None |
| if any(r is not None for _, r, _ in prepared_clauses): |
| from pypaimon.ray.merge_condition import filter_batch as _filter_batch |
| |
| def _transform(batch: pa.Table) -> pa.Table: |
| remaining = batch |
| parts = [] |
| for spec, rewritten, is_delete in prepared_clauses: |
| if remaining.num_rows == 0: |
| break |
| if rewritten is not None: |
| matched = _filter_batch( |
| remaining, rewritten, _pre_rewritten=True, |
| ) |
| else: |
| matched = remaining |
| if matched.num_rows == 0: |
| continue |
| if not is_delete: |
| callable_input = None |
| if (callable_input_columns is not None |
| and any(callable(value) and not isinstance(value, type) |
| for value in spec.values())): |
| callable_input = pa.Table.from_arrays( |
| [matched.column(f"t.{col}") |
| for col in callable_input_columns], |
| names=list(callable_input_columns), |
| ) |
| parts.append(vectorized_matched_transform( |
| matched, spec, on_pairs, |
| update_cols, row_id_name, |
| update_schema, |
| callable_input=callable_input, |
| )) |
| if rewritten is not None and matched.num_rows < remaining.num_rows: |
| not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" |
| remaining = _filter_batch( |
| remaining, not_cond, _pre_rewritten=True, |
| ) |
| else: |
| remaining = remaining.slice(0, 0) |
| if not parts: |
| return update_schema.empty_table() |
| return pa.concat_tables(parts) |
| |
| return _transform |
| |
| |
| def _build_matched_delete_transform( |
| clauses: List[_NormalizedClause], |
| on_map: Dict[str, str], |
| row_id_name: str, |
| delete_schema: pa.Schema, |
| ): |
| prepared_clauses = [] |
| for clause in clauses: |
| rewritten = None |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import ( |
| remap_source_on_keys, rewrite_condition, |
| ) |
| rewritten = remap_source_on_keys( |
| rewrite_condition(clause.condition), on_map, |
| ) |
| prepared_clauses.append((rewritten, clause.delete)) |
| |
| _filter_batch = None |
| if any(r is not None for r, _ in prepared_clauses): |
| from pypaimon.ray.merge_condition import filter_batch as _filter_batch |
| |
| def _transform(batch: pa.Table) -> pa.Table: |
| remaining = batch |
| parts = [] |
| for rewritten, is_delete in prepared_clauses: |
| if remaining.num_rows == 0: |
| break |
| if rewritten is not None: |
| matched = _filter_batch( |
| remaining, rewritten, _pre_rewritten=True, |
| ) |
| else: |
| matched = remaining |
| if matched.num_rows > 0 and is_delete: |
| parts.append( |
| vectorized_delete_transform( |
| matched, row_id_name, delete_schema, |
| ) |
| ) |
| if rewritten is not None and matched.num_rows < remaining.num_rows: |
| not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" |
| remaining = _filter_batch( |
| remaining, not_cond, _pre_rewritten=True, |
| ) |
| else: |
| remaining = remaining.slice(0, 0) |
| if not parts: |
| return delete_schema.empty_table() |
| return pa.concat_tables(parts) |
| |
| return _transform |
| |
| |
| def build_self_merge_update_plan( |
| *, |
| table, |
| clauses: List[_NormalizedClause], |
| target_field_names: Sequence[str], |
| target_pa_schema: pa.Schema, |
| update_cols: Sequence[str], |
| resolve_target_projection, |
| snapshot_id: Optional[int] = None, |
| scan_predicate=None, |
| read_columns: Sequence[str] = (), |
| ) -> _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( |
| clauses, [row_id_name], update_cols, target_field_names, |
| )) |
| needed_cols.update(read_columns) |
| for clause in clauses: |
| for value in clause.spec.values(): |
| if isinstance(value, SourceColumnRef): |
| needed_cols.add(value.column) |
| target_set = set(target_field_names) |
| for clause in clauses: |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import extract_columns |
| for ref in extract_columns(clause.condition): |
| prefix, col = ref.split(".", 1) |
| if prefix == "s" and col in target_set: |
| needed_cols.add(col) |
| projection = [row_id_name] + [ |
| c for c in target_field_names if c in needed_cols |
| ] |
| |
| dynamic_options = {} |
| if snapshot_id is not None: |
| dynamic_options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) |
| if scan_predicate is not None: |
| dynamic_options[CoreOptions.SCALAR_INDEX_SEARCH_MODE.key()] = ( |
| GlobalIndexSearchMode.FULL.value |
| ) |
| scan_table = ( |
| table.copy_without_time_travel(dynamic_options) |
| if dynamic_options else table |
| ) |
| 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 |
| ) |
| # 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), |
| 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 |
| ), |
| ) |
| |
| |
| 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( |
| *, |
| target_identifier: str, |
| clauses: List[_NormalizedClause], |
| target_field_names: Sequence[str], |
| catalog_options: Dict[str, str], |
| resolve_target_projection, |
| snapshot_id: Optional[int] = None, |
| scan_predicate=None, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| ) -> Tuple: |
| from pypaimon.ray.ray_paimon import read_paimon |
| from pypaimon.table.special_fields import SpecialFields |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| needed_cols = set(resolve_target_projection( |
| clauses, [row_id_name], [], target_field_names, |
| )) |
| target_set = set(target_field_names) |
| for clause in clauses: |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import extract_columns |
| for ref in extract_columns(clause.condition): |
| prefix, col = ref.split(".", 1) |
| if prefix == "s" and col in target_set: |
| needed_cols.add(col) |
| projection = [row_id_name] + [ |
| c for c in target_field_names if c in needed_cols |
| ] |
| |
| read_kwargs = {} |
| if scan_predicate is not None: |
| from pypaimon.common.options.core_options import ( |
| CoreOptions, GlobalIndexSearchMode, |
| ) |
| 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, |
| ) |
| delete_schema = build_delete_schema(row_id_name) |
| |
| orig_names = target_ds.schema().names |
| target_renamed = target_ds.rename_columns( |
| {c: f"t.{c}" for c in orig_names} |
| ) |
| |
| 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), |
| ) |
| |
| _transform = _build_matched_delete_transform( |
| clauses, |
| on_map={row_id_name: row_id_name}, |
| row_id_name=row_id_name, |
| delete_schema=delete_schema, |
| ) |
| return aliased.map_batches(_transform, **_map_kwargs(ray_remote_args)) |
| |
| |
| def build_matched_update_ds( |
| *, |
| target_identifier: str, |
| source_ds, |
| target_on: Sequence[str], |
| source_on: Sequence[str], |
| clauses: List[_NormalizedClause], |
| target_field_names: Sequence[str], |
| target_pa_schema: pa.Schema, |
| update_cols: Sequence[str], |
| catalog_options: Dict[str, str], |
| num_partitions: Optional[int], |
| resolve_target_projection, |
| estimated_size_bytes: Optional[int] = None, |
| snapshot_id: Optional[int] = None, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| ) -> Tuple: |
| from pypaimon.ray.ray_paimon import read_paimon |
| from pypaimon.table.special_fields import SpecialFields |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| needed_cols = resolve_target_projection( |
| clauses, target_on, update_cols, target_field_names, |
| ) |
| projection = [row_id_name] + [c for c in needed_cols if c != row_id_name] |
| |
| target_ds = read_paimon( |
| target_identifier, catalog_options, |
| projection=projection, snapshot_id=snapshot_id, |
| ) |
| update_schema = build_update_schema(target_pa_schema, update_cols, row_id_name) |
| |
| target_renamed = target_ds.rename_columns( |
| {c: f"t.{c}" for c in target_ds.schema().names} |
| ) |
| num_partitions = _resolve_matched_num_partitions( |
| num_partitions, estimated_size_bytes, target_renamed, |
| ) |
| source_cols = _resolve_source_projection( |
| clauses, source_on, source_ds.schema().names, |
| ) |
| source_ds = source_ds.select_columns(source_cols) |
| source_renamed = source_ds.rename_columns( |
| {c: f"s.{c}" for c in source_cols} |
| ) |
| |
| joined = target_renamed.join( |
| source_renamed, |
| join_type="inner", |
| num_partitions=num_partitions, |
| on=tuple(f"t.{c}" for c in target_on), |
| right_on=tuple(f"s.{c}" for c in source_on), |
| ) |
| |
| _transform = _build_matched_transform( |
| clauses, |
| on_map=dict(zip(source_on, target_on)), |
| on_pairs=list(zip(source_on, target_on)), |
| update_cols=list(update_cols), |
| row_id_name=row_id_name, |
| update_schema=update_schema, |
| ) |
| return joined.map_batches(_transform, **_map_kwargs(ray_remote_args)) |
| |
| |
| def build_matched_delete_ds( |
| *, |
| target_identifier: str, |
| source_ds, |
| target_on: Sequence[str], |
| source_on: Sequence[str], |
| clauses: List[_NormalizedClause], |
| target_field_names: Sequence[str], |
| catalog_options: Dict[str, str], |
| num_partitions: Optional[int], |
| resolve_target_projection, |
| estimated_size_bytes: Optional[int] = None, |
| snapshot_id: Optional[int] = None, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| ) -> Tuple: |
| from pypaimon.ray.ray_paimon import read_paimon |
| from pypaimon.table.special_fields import SpecialFields |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| needed_cols = resolve_target_projection( |
| clauses, |
| target_on, |
| [], |
| target_field_names, |
| ) |
| projection = [row_id_name] + [c for c in needed_cols if c != row_id_name] |
| |
| target_ds = read_paimon( |
| target_identifier, catalog_options, |
| projection=projection, snapshot_id=snapshot_id, |
| ) |
| delete_schema = build_delete_schema(row_id_name) |
| |
| target_renamed = target_ds.rename_columns( |
| {c: f"t.{c}" for c in target_ds.schema().names} |
| ) |
| num_partitions = _resolve_matched_num_partitions( |
| num_partitions, estimated_size_bytes, target_renamed, |
| ) |
| source_cols = list(source_ds.schema().names) |
| source_renamed = source_ds.rename_columns( |
| {c: f"s.{c}" for c in source_cols} |
| ) |
| |
| joined = target_renamed.join( |
| source_renamed, |
| join_type="inner", |
| num_partitions=num_partitions, |
| on=tuple(f"t.{c}" for c in target_on), |
| right_on=tuple(f"s.{c}" for c in source_on), |
| ) |
| |
| _transform = _build_matched_delete_transform( |
| clauses, |
| on_map=dict(zip(source_on, target_on)), |
| row_id_name=row_id_name, |
| delete_schema=delete_schema, |
| ) |
| return joined.map_batches(_transform, **_map_kwargs(ray_remote_args)) |
| |
| |
| def distributed_update_apply( |
| update_ds, |
| table, |
| write_update_cols: Sequence[str], |
| *, |
| num_partitions: Optional[int], |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| base_snapshot_id: Optional[int] = None, |
| collect_row_ids: bool = False, |
| estimated_size_bytes: Optional[int] = None, |
| estimated_num_rows: Optional[int] = None, |
| data_context=None, |
| ) -> Tuple[list, int, list]: |
| import numpy as np |
| import pickle |
| import uuid |
| |
| import pyarrow.compute as pc |
| import ray |
| |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER |
| from pypaimon.table.special_fields import SpecialFields |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| cols = list(write_update_cols) |
| |
| for col in cols: |
| if col not in table.field_names: |
| raise ValueError( |
| f"Column '{col}' is not in target table schema." |
| ) |
| |
| # Pin the planner to the caller's base snapshot so row-id routing and the |
| # commit-time conflict check agree even if a concurrent commit lands (mirrors |
| # the delete path). |
| from pypaimon.common.options.core_options import CoreOptions |
| scan_table = ( |
| table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) |
| if base_snapshot_id is not None else table |
| ) |
| planner = TableUpdateByRowId( |
| scan_table, |
| "_merge_into_planner_" + uuid.uuid4().hex[:8], |
| BATCH_COMMIT_IDENTIFIER, |
| ) |
| sorted_first_row_ids = list(planner.first_row_ids) |
| if not sorted_first_row_ids: |
| return [], 0, [] |
| |
| num_partitions = _resolve_row_id_num_partitions( |
| num_partitions, |
| estimated_size_bytes, |
| estimated_num_rows, |
| len(sorted_first_row_ids), |
| data_context=data_context, |
| ) |
| |
| # 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 = ( |
| base_snapshot_id if base_snapshot_id is not None |
| else planner.snapshot_id |
| ) |
| |
| # Put file metadata into Ray's object store and pass a single ref to |
| # workers. Avoids per-task manifest re-scans (Jingsong review #6) and |
| # avoids serializing the metadata into every task's closure. Override |
| # snapshot_id with the join's base snapshot so commit-time conflict |
| # detection covers the read→planner window. |
| from dataclasses import replace |
| files_info = replace( |
| planner._snapshot_files_info(), |
| snapshot_id=check_from_snapshot, |
| ) |
| precomputed_info_ref = ray.put(files_info) |
| |
| frid_col = "_FIRST_ROW_ID" |
| sentinel_row_id = -1 |
| captured_sorted = sorted_first_row_ids |
| captured_sorted_arr = np.asarray(captured_sorted, dtype=np.int64) |
| valid_ranges = planner.valid_row_id_ranges |
| range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) |
| range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) |
| |
| def _assign_frid(batch: pa.Table) -> pa.Table: |
| if batch.num_rows == 0: |
| return batch.append_column( |
| frid_col, pa.array([], type=pa.int64()) |
| ) |
| rid_col = batch.column(row_id_name) |
| if rid_col.null_count: |
| raise ValueError( |
| "_ROW_ID is null; planner snapshot is stale " |
| "or matched rows come from a different table." |
| ) |
| rids = rid_col.to_numpy(zero_copy_only=False) |
| # Check each row_id belongs to a valid range (vectorized). |
| in_range = np.zeros(len(rids), dtype=bool) |
| for s, e in zip(range_starts, range_ends): |
| in_range |= (rids >= s) & (rids <= e) |
| if not in_range.all(): |
| bad = rids[~in_range][0] |
| raise ValueError( |
| f"_ROW_ID {bad} does not belong to any valid range " |
| f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; " |
| f"planner snapshot is stale or matched rows come " |
| f"from a different table." |
| ) |
| idx = np.searchsorted( |
| captured_sorted_arr, rids, side="right" |
| ) - 1 |
| frids = captured_sorted_arr[idx] |
| return batch.append_column( |
| frid_col, pa.array(frids, type=pa.int64()) |
| ) |
| |
| map_kwargs = _map_kwargs(ray_remote_args) |
| with_frid = update_ds.map_batches(_assign_frid, **map_kwargs) |
| |
| from pypaimon.schema.data_types import PyarrowFieldParser |
| target_pa = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields) |
| update_schema = build_update_schema(target_pa, cols, row_id_name) |
| # Ray drops the schema of an empty shuffle input. A negative row id keeps |
| # the group-by schema in a separate group that is never written. |
| sentinel = pa.Table.from_arrays( |
| [pa.array([sentinel_row_id], type=pa.int64())] |
| + [pa.nulls(1, type=update_schema.field(col).type) for col in cols], |
| schema=update_schema, |
| ).append_column( |
| frid_col, pa.array([sentinel_row_id], type=pa.int64()) |
| ) |
| with_frid = with_frid.union(ray.data.from_arrow(sentinel)) |
| |
| captured_table = table |
| captured_cols = cols |
| |
| def _apply_group(group: pa.Table) -> pa.Table: |
| if group.column(frid_col)[0].as_py() == sentinel_row_id: |
| return pa.Table.from_pydict({ |
| "msgs_blob": pa.array([], type=pa.binary()), |
| "n_updated": pa.array([], type=pa.int64()), |
| "row_ids_blob": pa.array([], type=pa.binary()), |
| }) |
| |
| if ( |
| pc.count_distinct(group.column(row_id_name)).as_py() |
| != group.num_rows |
| ): |
| raise ValueError( |
| "MERGE matched multiple source rows to the same " |
| "target _ROW_ID. Deduplicate the source before " |
| "merging." |
| ) |
| |
| for_update = group.drop_columns([frid_col]) |
| row_ids = ( |
| for_update.column(row_id_name).to_pylist() |
| if collect_row_ids else [] |
| ) |
| worker = TableUpdateByRowId( |
| captured_table, |
| "_merge_into_shard_" + uuid.uuid4().hex[:8], |
| BATCH_COMMIT_IDENTIFIER, |
| _precomputed_files_info=ray.get(precomputed_info_ref), |
| ) |
| msgs = worker.update_columns(for_update, list(captured_cols)) |
| return pa.Table.from_pydict({ |
| "msgs_blob": [pickle.dumps(msgs)], |
| "n_updated": pa.array( |
| [for_update.num_rows], type=pa.int64() |
| ), |
| "row_ids_blob": pa.array( |
| [pickle.dumps(row_ids)], type=pa.binary() |
| ), |
| }) |
| |
| # One group per target data file; bounded by file count and num_partitions. |
| group_partitions = max( |
| 1, min(len(captured_sorted), num_partitions) |
| ) |
| msgs_ds = with_frid.groupby( |
| frid_col, num_partitions=group_partitions |
| ).map_groups(_apply_group, **map_kwargs) |
| |
| all_msgs: list = [] |
| num_updated = 0 |
| action_row_ids = [] |
| for batch in msgs_ds.iter_batches(batch_format="pyarrow"): |
| for blob in batch.column("msgs_blob").to_pylist(): |
| all_msgs.extend(pickle.loads(blob)) |
| for n in batch.column("n_updated").to_pylist(): |
| num_updated += n |
| if collect_row_ids: |
| for blob in batch.column("row_ids_blob").to_pylist(): |
| action_row_ids.extend(pickle.loads(blob)) |
| return all_msgs, num_updated, action_row_ids |
| |
| |
| def _read_output_schema(table, read_cols: Sequence[str]) -> "pa.Schema": |
| """Result schema: each projected column's type plus int64 ``_ROW_ID``, in |
| ``read_cols`` order. Shared by the empty-result paths so they can't drift.""" |
| from pypaimon.schema.data_types import PyarrowFieldParser |
| from pypaimon.table.special_fields import SpecialFields |
| |
| rid = SpecialFields.ROW_ID.name |
| full = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields) |
| # Keep each field's nullability so an empty result matches a non-empty read. |
| return pa.schema([ |
| pa.field(rid, pa.int64(), nullable=False) if col == rid else full.field(col) |
| for col in read_cols |
| ]) |
| |
| |
| def distributed_read_by_row_id( |
| row_ids_ds, |
| table, |
| projection: Sequence[str], |
| *, |
| num_partitions: Optional[int], |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| base_snapshot_id: Optional[int] = None, |
| estimated_size_bytes: Optional[int] = None, |
| estimated_num_rows: Optional[int] = None, |
| data_context=None, |
| ): |
| """Read ``projection`` for the ``_ROW_ID``s in ``row_ids_ds``, routing each to its |
| owning file and reading only the matched rows via ``IndexedSplit`` slicing (blob |
| resolved). Returns a ``ray.data.Dataset`` of ``(*projection, _ROW_ID)``, or ``None`` |
| if the target is empty. Read-side mirror of ``distributed_update_apply``. |
| """ |
| import numpy as np |
| import uuid |
| |
| import ray |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.globalindex.indexed_split import IndexedSplit |
| from pypaimon.read.split import DataSplit |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER |
| from pypaimon.table.special_fields import SpecialFields |
| from pypaimon.utils.range import Range |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| read_cols = list(projection) |
| if row_id_name not in read_cols: |
| read_cols.append(row_id_name) |
| |
| # Typed empty block so all output blocks share one schema. |
| empty_out = _read_output_schema(table, read_cols).empty_table() |
| |
| # Read-only planner (only scans the manifest); pinned to the base snapshot for stable routing. |
| scan_table = ( |
| table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) |
| if base_snapshot_id is not None else table |
| ) |
| planner = TableUpdateByRowId( |
| scan_table, |
| "_read_by_row_id_planner_" + uuid.uuid4().hex[:8], |
| BATCH_COMMIT_IDENTIFIER, |
| ) |
| sorted_first_row_ids = list(planner.first_row_ids) |
| if not sorted_first_row_ids: |
| return None |
| |
| num_partitions = _resolve_row_id_num_partitions( |
| num_partitions, |
| estimated_size_bytes, |
| estimated_num_rows, |
| len(sorted_first_row_ids), |
| data_context=data_context, |
| ) |
| |
| precomputed_info_ref = ray.put(planner._snapshot_files_info()) |
| frid_col = "_FIRST_ROW_ID" |
| sorted_arr = np.asarray(sorted_first_row_ids, dtype=np.int64) |
| valid_ranges = planner.valid_row_id_ranges |
| range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) |
| range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) |
| |
| def _assign_frid(batch: pa.Table) -> pa.Table: |
| if batch.num_rows == 0: |
| return batch.append_column(frid_col, pa.array([], type=pa.int64())) |
| rid_col = batch.column(row_id_name) |
| if rid_col.null_count: |
| raise ValueError( |
| "_ROW_ID is null; the planner snapshot is stale or the row ids " |
| "come from a different table." |
| ) |
| rids = rid_col.to_numpy(zero_copy_only=False) |
| # Foreign-id check: valid_ranges are sorted+merged, so one searchsorted finds |
| # the candidate range (O(rows log ranges), like distributed_delete_apply). |
| ridx = np.searchsorted(range_starts, rids, side="right") - 1 |
| safe = np.clip(ridx, 0, len(range_starts) - 1) |
| in_range = ( |
| (ridx >= 0) |
| & (rids >= range_starts[safe]) |
| & (rids <= range_ends[safe]) |
| ) |
| if not in_range.all(): |
| bad = rids[~in_range][0] |
| raise ValueError( |
| f"_ROW_ID {bad} does not belong to any valid range " |
| f"{[f'[{r.from_}, {r.to}]' for r in valid_ranges]}; the planner " |
| f"snapshot is stale or the row ids come from a different table." |
| ) |
| idx = np.searchsorted(sorted_arr, rids, side="right") - 1 |
| return batch.append_column( |
| frid_col, pa.array(sorted_arr[idx], type=pa.int64()) |
| ) |
| |
| captured_table = scan_table # read at the same pinned snapshot the planner routed on |
| captured_read_cols = read_cols |
| captured_empty = empty_out |
| |
| def _read_group(group: pa.Table) -> pa.Table: |
| if group.num_rows == 0: |
| return captured_empty |
| frid = int(group.column(frid_col)[0].as_py()) |
| info = ray.get(precomputed_info_ref) |
| owning_split, target_files = info.first_row_id_index[frid] |
| origin_split = DataSplit( |
| files=target_files, |
| partition=owning_split.partition, |
| bucket=owning_split.bucket, |
| raw_convertible=True, |
| ) |
| # Only matched rows (deduped, contiguous ids -> ranges); blob gets row-index pushdown. |
| wanted = set(group.column(row_id_name).to_pylist()) |
| indexed = IndexedSplit(origin_split, Range.to_ranges(list(wanted))) |
| read = captured_table.new_read_builder().with_projection( |
| captured_read_cols |
| ).new_read() |
| return read.to_arrow([indexed]) |
| |
| map_kwargs = _map_kwargs(ray_remote_args) |
| with_frid = row_ids_ds.map_batches(_assign_frid, **map_kwargs) |
| group_partitions = max(1, min(len(sorted_first_row_ids), num_partitions)) |
| return with_frid.groupby(frid_col, num_partitions=group_partitions).map_groups( |
| _read_group, **map_kwargs |
| ) |
| |
| |
| def distributed_delete_apply( |
| delete_ds, |
| table, |
| *, |
| num_partitions: int, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| base_snapshot_id: Optional[int] = None, |
| collect_row_ids: bool = False, |
| ) -> Tuple[list, int, list]: |
| import base64 |
| import numpy as np |
| import pickle |
| |
| import pyarrow.compute as pc |
| import ray |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.table.special_fields import SpecialFields |
| from pypaimon.write.table_delete import TableDeleteByRowId |
| |
| row_id_name = SpecialFields.ROW_ID.name |
| scan_table = ( |
| table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) |
| if base_snapshot_id is not None else table |
| ) |
| |
| planner = TableDeleteByRowId(scan_table) |
| anchor_info = planner._snapshot_anchor_ranges() |
| if not anchor_info.anchors: |
| return [], 0, [] |
| |
| precomputed_info_ref = ray.put(anchor_info) |
| |
| starts = np.asarray( |
| [a.row_range.from_ for a in anchor_info.anchors], dtype=np.int64 |
| ) |
| ends = np.asarray( |
| [a.row_range.to for a in anchor_info.anchors], dtype=np.int64 |
| ) |
| |
| def _group_key(anchor) -> str: |
| partition_blob = base64.b64encode( |
| pickle.dumps(tuple(anchor.partition.values)) |
| ).decode("ascii") |
| return f"{anchor.bucket}:{partition_blob}" |
| |
| group_keys = [_group_key(a) for a in anchor_info.anchors] |
| unique_group_count = len(set(group_keys)) |
| group_col = "_DELETE_GROUP_KEY" |
| valid_ranges = [ |
| f"[{a.row_range.from_}, {a.row_range.to}]" |
| for a in anchor_info.anchors |
| ] |
| |
| def _assign_group(batch: pa.Table) -> pa.Table: |
| if batch.num_rows == 0: |
| return batch.append_column( |
| group_col, pa.array([], type=pa.string()) |
| ) |
| rid_col = batch.column(row_id_name) |
| if rid_col.null_count: |
| raise ValueError( |
| "_ROW_ID is null; planner snapshot is stale " |
| "or matched rows come from a different table." |
| ) |
| rids = rid_col.to_numpy(zero_copy_only=False) |
| idx = np.searchsorted(starts, rids, side="right") - 1 |
| safe_idx = np.clip(idx, 0, len(starts) - 1) |
| in_range = ( |
| (idx >= 0) |
| & (idx < len(starts)) |
| & (rids >= starts[safe_idx]) |
| & (rids <= ends[safe_idx]) |
| ) |
| if not in_range.all(): |
| bad = rids[~in_range][0] |
| raise ValueError( |
| f"_ROW_ID {bad} does not belong to any valid range " |
| f"{valid_ranges}; planner snapshot is stale or matched " |
| f"rows come from a different table." |
| ) |
| return batch.append_column( |
| group_col, |
| pa.array([group_keys[i] for i in safe_idx], type=pa.string()), |
| ) |
| |
| map_kwargs = _map_kwargs(ray_remote_args) |
| with_group = delete_ds.map_batches(_assign_group, **map_kwargs) |
| captured_table = scan_table |
| |
| def _apply_group(group: pa.Table) -> pa.Table: |
| if group.num_rows == 0: |
| return pa.Table.from_pydict({ |
| "msgs_blob": pa.array([], type=pa.binary()), |
| "n_deleted": pa.array([], type=pa.int64()), |
| "row_ids_blob": pa.array([], type=pa.binary()), |
| }) |
| |
| if ( |
| pc.count_distinct(group.column(row_id_name)).as_py() |
| != group.num_rows |
| ): |
| raise ValueError( |
| "MERGE matched multiple source rows to the same " |
| "target _ROW_ID. Deduplicate the source before " |
| "merging." |
| ) |
| |
| row_ids = group.column(row_id_name).to_pylist() |
| worker = TableDeleteByRowId( |
| captured_table, |
| _precomputed_anchor_ranges=ray.get(precomputed_info_ref), |
| ) |
| msgs = worker.delete(row_ids) |
| return pa.Table.from_pydict({ |
| "msgs_blob": pa.array([pickle.dumps(msgs)], type=pa.binary()), |
| "n_deleted": pa.array([len(row_ids)], type=pa.int64()), |
| "row_ids_blob": pa.array( |
| [pickle.dumps(row_ids if collect_row_ids else [])], |
| type=pa.binary(), |
| ), |
| }) |
| |
| group_partitions = max(1, min(unique_group_count, num_partitions)) |
| msgs_ds = with_group.groupby( |
| group_col, num_partitions=group_partitions |
| ).map_groups(_apply_group, **map_kwargs) |
| |
| all_msgs: list = [] |
| num_deleted = 0 |
| action_row_ids = [] |
| for batch in msgs_ds.iter_batches(batch_format="pyarrow"): |
| for blob in batch.column("msgs_blob").to_pylist(): |
| all_msgs.extend(pickle.loads(blob)) |
| for n in batch.column("n_deleted").to_pylist(): |
| num_deleted += n |
| if collect_row_ids: |
| for blob in batch.column("row_ids_blob").to_pylist(): |
| action_row_ids.extend(pickle.loads(blob)) |
| return all_msgs, num_deleted, action_row_ids |
| |
| |
| def build_not_matched_insert_ds( |
| *, |
| target_identifier: str, |
| source_ds, |
| target_on: Sequence[str], |
| source_on: Sequence[str], |
| clauses: List[_NormalizedClause], |
| target_field_names: Sequence[str], |
| target_pa_schema: pa.Schema, |
| catalog_options: Dict[str, str], |
| num_partitions: int, |
| target_empty: bool = False, |
| snapshot_id: Optional[int] = None, |
| ray_remote_args: Optional[Dict[str, Any]] = None, |
| ): |
| from pypaimon.ray.ray_paimon import read_paimon |
| |
| captured_field_names = list(target_field_names) |
| out_schema = target_pa_schema |
| |
| source_cols = _resolve_source_projection( |
| clauses, source_on, source_ds.schema().names, |
| ) |
| source_ds = source_ds.select_columns(source_cols) |
| source_renamed = source_ds.rename_columns( |
| {c: f"s.{c}" for c in source_cols} |
| ) |
| |
| if target_empty: |
| unmatched = source_renamed.repartition(num_partitions) |
| else: |
| target_ds = read_paimon( |
| target_identifier, catalog_options, |
| projection=list(target_on), snapshot_id=snapshot_id, |
| ) |
| target_renamed = target_ds.rename_columns( |
| {c: f"t.{c}" for c in target_on} |
| ) |
| unmatched = source_renamed.join( |
| target_renamed, |
| join_type="left_anti", |
| num_partitions=num_partitions, |
| on=tuple(f"s.{c}" for c in source_on), |
| right_on=tuple(f"t.{c}" for c in target_on), |
| ) |
| |
| prepared_clauses = [] |
| for clause in clauses: |
| rewritten = None |
| if clause.condition is not None: |
| from pypaimon.ray.merge_condition import rewrite_condition |
| rewritten = rewrite_condition(clause.condition) |
| prepared_clauses.append((clause.spec, rewritten)) |
| |
| _filter_batch_nm = None |
| if any(r is not None for _, r in prepared_clauses): |
| from pypaimon.ray.merge_condition import filter_batch as _filter_batch_nm |
| |
| def _transform(batch: pa.Table) -> pa.Table: |
| remaining = batch |
| parts = [] |
| for spec, rewritten in prepared_clauses: |
| if remaining.num_rows == 0: |
| break |
| if rewritten is not None: |
| matched = _filter_batch_nm( |
| remaining, rewritten, _pre_rewritten=True, |
| ) |
| if matched.num_rows > 0: |
| parts.append(vectorized_insert_transform( |
| matched, spec, captured_field_names, out_schema |
| )) |
| if matched.num_rows < remaining.num_rows: |
| not_cond = f"COALESCE(NOT ({rewritten}), TRUE)" |
| remaining = _filter_batch_nm( |
| remaining, not_cond, _pre_rewritten=True, |
| ) |
| else: |
| remaining = remaining.slice(0, 0) |
| else: |
| parts.append(vectorized_insert_transform( |
| remaining, spec, captured_field_names, out_schema |
| )) |
| remaining = remaining.slice(0, 0) |
| if not parts: |
| return out_schema.empty_table() |
| return cast_to_schema(pa.concat_tables(parts), out_schema) |
| |
| return unmatched.map_batches( |
| _transform, **_map_kwargs(ray_remote_args) |
| ) |
| |
| |
| def distributed_write_collect_msgs( |
| insert_ds, |
| table, |
| *, |
| ray_remote_args: Optional[Dict[str, Any]], |
| concurrency: Optional[int], |
| ) -> list: |
| from pypaimon.write.ray_datasink import PaimonDatasink |
| |
| class _CollectingDatasink(PaimonDatasink): |
| def __init__(self, t): |
| super().__init__(t, overwrite=False) |
| self.collected: list = [] |
| |
| def on_write_complete(self, write_result): |
| self.collected = [ |
| m |
| for batch in self._extract_write_returns(write_result) |
| for m in batch |
| if not m.is_empty() |
| ] |
| |
| sink = _CollectingDatasink(table) |
| write_kwargs: Dict[str, Any] = {} |
| if ray_remote_args is not None: |
| write_kwargs["ray_remote_args"] = ray_remote_args |
| if concurrency is not None: |
| write_kwargs["concurrency"] = concurrency |
| insert_ds.write_datasink(sink, **write_kwargs) |
| return sink.collected |
