| # 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. | |
| """Select update execution before mutations, preserving public owner hooks.""" | |
| from typing import Any, Iterable, List, Mapping, Optional, Sequence | |
| import pyarrow as pa | |
| from pypaimon.common.predicate import Predicate | |
| from pypaimon.write.backend_dispatch import ( | |
| create_python_merge_into_operation, create_python_update_operation, create_python_upsert_operation, | |
| prepare_native_write_operation, | |
| ) | |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER | |
| from pypaimon.table.special_fields import SpecialFields | |
| from pypaimon.write.commit_message import CommitMessage | |
| from pypaimon.write.row_utils import row_to_named_values | |
| from pypaimon.write.row_input import RowIdInput, normalize_row_ids | |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId | |
| from pypaimon.write.table_upsert_by_key import TableUpsertByKey | |
| from pypaimon.write.update_validation import _authorize_update | |
| class TableUpdateOperations: | |
| def __init__(self, context): | |
| self._context = context | |
| def _new_row_id_updater(self, commit_identifier: int, _precomputed_files_info=None) -> TableUpdateByRowId: | |
| context = self._context | |
| # The public constructor and builders share one selection boundary, | |
| # including complete plans transported to distributed workers. | |
| return TableUpdateByRowId(context.table, context.commit_user, commit_identifier, | |
| _precomputed_files_info=_precomputed_files_info) | |
| def new_shard_updator(self, shard_num: int, total_shard_count: int): | |
| context = self._context | |
| from pypaimon.write.table_update import StreamTableUpdate | |
| """Create a shard updater for scan and rewrite style updates.""" | |
| if total_shard_count <= 0 or shard_num < 0 or shard_num >= total_shard_count: | |
| raise ValueError('Shard index must be smaller than a positive shard count') | |
| if context.update_cols == []: | |
| raise ValueError('Shard update columns must not be empty') | |
| _authorize_update(context.table) | |
| native = prepare_native_write_operation( | |
| 'shard_update', | |
| context.table, | |
| context.commit_user, | |
| context.projection, | |
| context.update_cols, | |
| shard_num, | |
| total_shard_count, | |
| stream=isinstance( | |
| context, | |
| StreamTableUpdate)) | |
| if native is not None: | |
| return native | |
| from pypaimon.write.table_update import ShardTableUpdator | |
| return ShardTableUpdator._from_python( | |
| context.table, context.projection, context.update_cols, context.commit_user, shard_num, total_shard_count) | |
| def _update_by_arrow_with_row_id(self, table: pa.Table, commit_identifier: int) -> List[CommitMessage]: | |
| context = self._context | |
| cols = context.update_cols if context.update_cols is not None else [ | |
| c for c in table.column_names if c != SpecialFields.ROW_ID.name] | |
| operation = None | |
| if table.num_rows and cols and (SpecialFields.ROW_ID.name in table.column_names): | |
| operation = prepare_native_write_operation('update', context.table, context.commit_user, cols) | |
| if operation is None: | |
| operation = create_python_update_operation(context, commit_identifier) | |
| return operation.update_by_arrow_with_row_id(table) | |
| def _upsert_by_arrow_with_key( | |
| self, table: pa.Table, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: | |
| """Shared implementation for ``upsert_by_arrow_with_key``. | |
| For each row in the input Arrow table: | |
| * If one or more rows with the same composite ``upsert_keys`` value | |
| already exist → update all of them in-place. | |
| * Otherwise → append as a new row. | |
| The public method lives on the concrete subclasses so each can | |
| expose the signature appropriate to its mode (batch vs stream). | |
| Args: | |
| table: Input Arrow table containing rows to upsert. | |
| upsert_keys: One or more column names forming the composite match key. | |
| commit_identifier: Identifier to tag the produced commit messages with. | |
| Returns: | |
| List of :class:`CommitMessage` objects to be committed. | |
| """ | |
| context = self._context | |
| return TableUpsertByKey(context.table, context.commit_user, commit_identifier).upsert( | |
| table, upsert_keys, context.update_cols) | |
| def _upsert_by_key(self, rows, upsert_keys: List[str], commit_identifier: int) -> List[CommitMessage]: | |
| context = self._context | |
| return TableUpsertByKey(context.table, context.commit_user, commit_identifier).upsert_rows( | |
| rows, upsert_keys, context.update_cols) | |
| def _merge_into(self, source: Any, on, when_matched: Sequence, when_not_matched: Sequence, | |
| commit_identifier: int) -> List[CommitMessage]: | |
| context = self._context | |
| from pypaimon.table.data_evolution_merge_into import merge_into | |
| return merge_into(context.table, source, on=on, when_matched=when_matched, when_not_matched=when_not_matched, | |
| commit_user=context.commit_user, commit_identifier=commit_identifier) | |
| def _update_by_predicate( | |
| self, predicate: Optional[Predicate], assignments: Mapping[str, Any], | |
| commit_identifier: int, read_columns: Optional[Sequence[str]] = None) -> List[CommitMessage]: | |
| """Shared implementation for SQL-like ``UPDATE ... WHERE ...``. | |
| ``predicate`` identifies the target rows. Assignment values may be | |
| literals or callables receiving one matched logical file group as an | |
| Arrow table. | |
| """ | |
| context = self._context | |
| has_callable = any((callable(value) for value in assignments.values())) | |
| has_array = any((isinstance(value, (pa.Array, pa.ChunkedArray)) for value in assignments.values())) | |
| read_columns = tuple(read_columns or ()) | |
| context._validate_predicate_update(assignments, read_columns, has_callable, has_array) | |
| operation = None | |
| if context.table.options.native_write_enabled(): | |
| operation = prepare_native_write_operation( | |
| 'predicate_update', | |
| context.table, | |
| context.commit_user, | |
| predicate, | |
| list(assignments)) | |
| if operation is None: | |
| operation = create_python_update_operation( | |
| context, commit_identifier, predicate, has_callable=has_callable, has_array=has_array) | |
| return operation.update(assignments, read_columns) | |
| def _delete_by_predicate(self, predicate: Optional[Predicate], commit_identifier: int) -> List[CommitMessage]: | |
| context = self._context | |
| operation = prepare_native_write_operation('predicate_delete', context.table, context.commit_user, predicate) | |
| if operation is None: | |
| operation = create_python_update_operation(context, commit_identifier, predicate) | |
| return operation.delete() | |
| def _delete_by_row_id(self, row_ids: RowIdInput, commit_identifier: int) -> List[CommitMessage]: | |
| context = self._context | |
| row_ids = normalize_row_ids(row_ids) | |
| operation = prepare_native_write_operation('delete', context.table, context.commit_user) | |
| if operation is None: | |
| operation = create_python_update_operation(context, commit_identifier) | |
| return operation.delete_by_row_id(row_ids) | |
| def update_by_arrow_batches_with_row_id(self, tables: Iterable[pa.Table]) -> List[CommitMessage]: | |
| """Apply row-id updates from batches using one target-file index. | |
| For each updated column, batches must target disjoint ``first_row_id`` | |
| file groups. Conflicting overlap is rejected. Already prepared files | |
| are retained because they may have been published by the caller. | |
| """ | |
| context = self._context | |
| operation = prepare_native_write_operation('update', context.table, context.commit_user, context.update_cols) | |
| if operation is None: | |
| operation = create_python_update_operation(context, BATCH_COMMIT_IDENTIFIER) | |
| return operation.update_by_arrow_batches_with_row_id(tables) | |
| class UpsertOperations: | |
| def __init__(self, context): | |
| self._context = context | |
| def upsert(self, data: pa.Table, upsert_keys: List[str], | |
| update_cols: Optional[List[str]] = None) -> List[CommitMessage]: | |
| """ | |
| Upsert rows into an append-only table by the specified key columns. | |
| Execution is driven partition-by-partition: | |
| 1. Group input rows by their partition values. | |
| 2. For each partition, scan only that partition to build the | |
| key -> _ROW_ID map. | |
| 3. Split the partition's input rows into matched (update) and | |
| unmatched (append). | |
| 4. Perform updates and appends. | |
| This keeps memory usage proportional to a single partition's key set | |
| and avoids scanning the entire table. | |
| Args: | |
| data: Input Arrow table containing rows to upsert. | |
| Must contain all upsert_keys columns. | |
| upsert_keys: One or more column names used together as a composite | |
| match key. | |
| update_cols: Columns to update for matched rows. | |
| If None, all columns in the table schema are updated. | |
| Returns: | |
| List of CommitMessages to be committed. | |
| """ | |
| context = self._context | |
| context._validate_inputs(data, upsert_keys, update_cols) | |
| if update_cols is None or len(update_cols) == len(context.table.field_names): | |
| effective_update_cols = None | |
| else: | |
| effective_update_cols = update_cols | |
| columns = list(effective_update_cols) if effective_update_cols else list(context.table.field_names) | |
| operation = context._create_native_upsert(data, upsert_keys, columns) | |
| if operation is None: | |
| operation = create_python_upsert_operation(context, data, upsert_keys, effective_update_cols) | |
| return operation.upsert() | |
| def upsert_rows(self, rows, upsert_keys: List[str], update_cols: Optional[List[str]] = None) -> List[CommitMessage]: | |
| context = self._context | |
| row_list = context._normalize_rows(rows) | |
| if not row_list: | |
| raise ValueError('rows must not be empty.') | |
| row_items = [(row, row_to_named_values(row, context.table.table_schema.fields)) for row in row_list] | |
| for _, values_by_name in row_items: | |
| context._validate_row_inputs(values_by_name, upsert_keys, update_cols) | |
| if update_cols is None or len(update_cols) == len(context.table.field_names): | |
| effective_update_cols = None | |
| else: | |
| effective_update_cols = update_cols | |
| columns = list(effective_update_cols) if effective_update_cols is not None else list(context.table.field_names) | |
| operation = context._create_native_upsert([values for _, values in row_items], upsert_keys, columns) | |
| if operation is None: | |
| operation = create_python_upsert_operation( | |
| context, row_items, upsert_keys, effective_update_cols, rows=True) | |
| return operation.upsert() | |
| def _create_native_upsert(self, data, upsert_keys, columns): | |
| context = self._context | |
| return prepare_native_write_operation('upsert', context.table, context.commit_user, data, upsert_keys, columns) | |
| def merge_into( | |
| target_table, source, *, on, when_matched=(), when_not_matched=(), | |
| commit_user, commit_identifier=BATCH_COMMIT_IDENTIFIER): | |
| """Select MERGE execution before consuming source input or staging files.""" | |
| if target_table.options.native_write_enabled(): | |
| from pypaimon.write.backend_dispatch import prepare_native_merge_into | |
| native = prepare_native_merge_into( | |
| target_table, source, on, when_matched, when_not_matched, | |
| commit_user, commit_identifier, | |
| ) | |
| if native is not None: | |
| return native.prepare_commit() | |
| return create_python_merge_into_operation( | |
| target_table, source, on=on, when_matched=when_matched, | |
| when_not_matched=when_not_matched, commit_user=commit_user, | |
| commit_identifier=commit_identifier, | |
| ).prepare_commit() |