This is an automated email from the ASF dual-hosted git repository. JingsongLi pushed a commit to branch codex/pypaimon-merge-into in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 45c9e4f3e7980f8878601602efe453d92778a6db Author: JingsongLi <[email protected]> AuthorDate: Wed Jul 1 18:04:59 2026 +0800 [python] Add merge_into for data evolution --- docs/docs/pypaimon/data-evolution.md | 337 ++++------ .../pypaimon/ray/data_evolution_merge_into.py | 54 +- .../pypaimon/table/data_evolution_merge_into.py | 690 +++++++++++++++++++++ .../pypaimon/tests/table_merge_into_test.py | 582 +++++++++++++++++ paimon-python/pypaimon/write/table_update.py | 57 +- 5 files changed, 1482 insertions(+), 238 deletions(-) diff --git a/docs/docs/pypaimon/data-evolution.md b/docs/docs/pypaimon/data-evolution.md index 1ec5a3a0a2..d1297b8051 100644 --- a/docs/docs/pypaimon/data-evolution.md +++ b/docs/docs/pypaimon/data-evolution.md @@ -3,9 +3,6 @@ title: "Data Evolution" sidebar_position: 5 --- -import Tabs from '@theme/Tabs'; -import TabItem from '@theme/TabItem'; - <!-- Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file @@ -36,46 +33,22 @@ To use partial updates / data evolution, enable both options when creating the t - **`row-tracking.enabled`**: `true` - **`data-evolution.enabled`**: `true` -## Batch vs Stream - -Data evolution supports both batch and stream modes. - -<Tabs groupId="pypaimon-data-evolution-mode"> +## Batch Update APIs -<TabItem value="batch" label="Batch"> +The examples below use batch mode: - Builder: `table.new_batch_write_builder()` - Write: `BatchTableWrite` - Update: `BatchTableUpdate` - Commit: `BatchTableCommit` -- `commit_identifier`: not required - Lifecycle: one-shot; each instance can commit only once - `prepare_commit()`: `write.prepare_commit()` - `update_by_arrow_with_row_id()`: `update.update_by_arrow_with_row_id(table)` - `update_by_predicate()`: `update.update_by_predicate(predicate, assignments)` - `upsert_by_arrow_with_key()`: `update.upsert_by_arrow_with_key(table, keys)` +- `merge_into()`: `update.merge_into(source, on=..., when_matched=..., when_not_matched=...)` - `commit()`: `commit.commit(messages)` -</TabItem> - -<TabItem value="stream" label="Stream"> - -- Builder: `table.new_stream_write_builder()` -- Write: `StreamTableWrite` -- Update: `StreamTableUpdate` -- Commit: `StreamTableCommit` -- `commit_identifier`: required; use a monotonically increasing integer -- Lifecycle: reusable; the same instance can commit many rounds -- `prepare_commit()`: `write.prepare_commit(commit_identifier)` -- `update_by_arrow_with_row_id()`: `update.update_by_arrow_with_row_id(table, commit_identifier)` -- `update_by_predicate()`: `update.update_by_predicate(predicate, assignments, commit_identifier)` -- `upsert_by_arrow_with_key()`: `update.upsert_by_arrow_with_key(table, keys, commit_identifier)` -- `commit()`: `commit.commit(messages, commit_identifier)` - -</TabItem> - -</Tabs> - ## Update Columns By Row ID You can use `update_by_arrow_with_row_id` to update columns in data evolution tables. @@ -87,10 +60,6 @@ to its corresponding `first_row_id`, then group rows with the same `first_row_id - **Update columns only**: include `_ROW_ID` plus the columns you want to update (partial schema is OK). -<Tabs groupId="pypaimon-data-evolution-mode"> - -<TabItem value="batch" label="Batch"> - ```python import pyarrow as pa from pypaimon import CatalogFactory, Schema @@ -140,64 +109,6 @@ table_commit.close() # 'f1': [-1001, 1002] ``` -</TabItem> - -<TabItem value="stream" label="Stream"> - -```python -import pyarrow as pa -from pypaimon import CatalogFactory, Schema - -catalog = CatalogFactory.create({'warehouse': '/tmp/warehouse'}) -catalog.create_database('default', False) - -simple_pa_schema = pa.schema([ - ('f0', pa.int8()), - ('f1', pa.int16()), -]) -schema = Schema.from_pyarrow_schema(simple_pa_schema, - options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}) -catalog.create_table('default.test_stream', schema, False) -table = catalog.get_table('default.test_stream') - -# write initial data -write_builder = table.new_batch_write_builder() -table_write = write_builder.new_write() -table_commit = write_builder.new_commit() -table_write.write_arrow(pa.Table.from_pydict({ - 'f0': [-1, 2], - 'f1': [-1001, 1002] -}, schema=simple_pa_schema)) -table_commit.commit(table_write.prepare_commit()) -table_write.close() -table_commit.close() - -# stream update: each round uses a new commit_identifier -stream_builder = table.new_stream_write_builder() -table_update = stream_builder.new_update().with_update_type(['f0']) -table_commit = stream_builder.new_commit() - -data1 = pa.Table.from_pydict({ - '_ROW_ID': [0], - 'f0': [5], -}, schema=pa.schema([('_ROW_ID', pa.int64()), ('f0', pa.int8())])) -cmts1 = table_update.update_by_arrow_with_row_id(data1, commit_identifier=1) -table_commit.commit(cmts1, commit_identifier=1) - -data2 = pa.Table.from_pydict({ - '_ROW_ID': [1], - 'f0': [6], -}, schema=pa.schema([('_ROW_ID', pa.int64()), ('f0', pa.int8())])) -cmts2 = table_update.update_by_arrow_with_row_id(data2, commit_identifier=2) -table_commit.commit(cmts2, commit_identifier=2) - -table_commit.close() -``` - -</TabItem> - -</Tabs> - ## Update Columns By Predicate You can use `update_by_predicate` for SQL-like `UPDATE ... SET ... WHERE ...` @@ -207,10 +118,6 @@ When global indexes are available, `update_by_predicate` discovers matching `_ROW_ID` values with `global-index.search-mode=full` on the configured point-in-time scan snapshot or, if none is configured, the latest snapshot. -<Tabs groupId="pypaimon-data-evolution-mode"> - -<TabItem value="batch" label="Batch"> - ```python import pyarrow as pa from pypaimon import CatalogFactory, Schema @@ -253,31 +160,6 @@ commit.commit(messages) commit.close() ``` -</TabItem> - -<TabItem value="stream" label="Stream"> - -Pass the `commit_identifier` to both `update_by_predicate` and `commit`: - -```python -stream_builder = table.new_stream_write_builder() -table_update = stream_builder.new_update() -table_commit = stream_builder.new_commit() - -predicate = table_update.new_predicate_builder().equal('id', 2) -messages = table_update.update_by_predicate( - predicate, - {'name': 'Bob_v2'}, - commit_identifier=1, -) -table_commit.commit(messages, commit_identifier=1) -table_commit.close() -``` - -</TabItem> - -</Tabs> - ## Filter by _ROW_ID Requires the same [Prerequisites](#prerequisites) (row-tracking and data-evolution enabled). On such tables you can filter by `_ROW_ID` to prune files at scan time. Supported: `equal('_ROW_ID', id)`, `is_in('_ROW_ID', [id1, ...])`, `between('_ROW_ID', low, high)`. @@ -304,10 +186,6 @@ If you want to **upsert** (update-or-insert) rows by one or more business key co **automatically stripped** from `upsert_keys` during matching (since each partition is processed independently), so you do **not** need to include them in `upsert_keys`. -<Tabs groupId="pypaimon-data-evolution-mode"> - -<TabItem value="batch" label="Batch"> - **Example: basic upsert** ```python @@ -420,13 +298,37 @@ table_commit.close() - Duplicate keys in the input data are automatically deduplicated — the **last occurrence** is kept. - The upsert is atomic per commit — all matched updates and new appends are included in the same commit. -</TabItem> +## Merge Into + +Use `merge_into` when your source data should update matched target rows and +optionally insert rows that do not match, similar to SQL `MERGE INTO`. +`merge_into` is exposed from `TableUpdate`, so it follows the same +commit-message lifecycle as other PyPaimon update APIs. The PyPaimon +implementation runs in a single process and materializes the rows it needs +locally. -<TabItem value="stream" label="Stream"> +Matched rows are updated by `_ROW_ID` internally. Only the columns touched by +the matched clauses are rewritten. `merge_into` derives the update columns from +the `WhenMatched` clauses; `with_update_type` is not needed. + +**Requirements** + +- The target table must have `data-evolution.enabled = true` and + `row-tracking.enabled = true`. +- `source` must be a `pyarrow.Table`, `pandas.DataFrame`, or another PyPaimon + table object. +- `on` can be a list of same-named key columns, or `{target_col: source_col}` + for renamed source keys. +- If multiple source rows match the same target `_ROW_ID`, `merge_into` raises + an error. Deduplicate the source before merging. ```python import pyarrow as pa from pypaimon import CatalogFactory, Schema +from pypaimon.table.data_evolution_merge_into import ( + WhenMatched, + WhenNotMatched, +) catalog = CatalogFactory.create({'warehouse': '/tmp/warehouse'}) catalog.create_database('default', False) @@ -440,8 +342,8 @@ schema = Schema.from_pyarrow_schema( pa_schema, options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}, ) -catalog.create_table('default.users_stream', schema, False) -table = catalog.get_table('default.users_stream') +catalog.create_table('default.users_merge', schema, False) +table = catalog.get_table('default.users_merge') # write initial data write_builder = table.new_batch_write_builder() @@ -455,33 +357,78 @@ commit.commit(write.prepare_commit()) write.close() commit.close() -# stream upsert: each round uses a new commit_identifier -stream_builder = table.new_stream_write_builder() -table_update = stream_builder.new_update() -table_commit = stream_builder.new_commit() - -upsert_data1 = pa.Table.from_pydict( - {'id': [1, 3], 'name': ['Alice_v2', 'Charlie'], 'age': [31, 28]}, +# merge: update id=2, insert id=3 +source = pa.Table.from_pydict( + {'id': [2, 3], 'name': ['Bob_v2', 'Charlie'], 'age': [26, 28]}, schema=pa_schema, ) -cmts1 = table_update.upsert_by_arrow_with_key(upsert_data1, upsert_keys=['id'], commit_identifier=1) -table_commit.commit(cmts1, commit_identifier=1) -upsert_data2 = pa.Table.from_pydict( - {'id': [2, 4], 'name': ['Bob_v2', 'David'], 'age': [26, 40]}, - schema=pa_schema, -) -cmts2 = table_update.upsert_by_arrow_with_key(upsert_data2, upsert_keys=['id'], commit_identifier=2) -table_commit.commit(cmts2, commit_identifier=2) +write_builder = table.new_batch_write_builder() +table_update = write_builder.new_update() +table_commit = write_builder.new_commit() +messages = table_update.merge_into( + source, + on=['id'], + when_matched=[WhenMatched(update='*')], + when_not_matched=[WhenNotMatched(insert='*')], +) +table_commit.commit(messages) table_commit.close() ``` -The `with_update_type` and partitioned table patterns shown in Batch Mode also work in Stream Mode — just add `commit_identifier` to the `upsert_by_arrow_with_key` and `commit` calls. +`WhenMatched` and `WhenNotMatched` clauses can use `'*'` to copy same-named +columns from source, or a mapping for explicit assignments: -</TabItem> +```python +from pypaimon.table.data_evolution_merge_into import ( + WhenMatched, + WhenNotMatched, + lit, + source_col, + target_col, +) + +messages = table_update.merge_into( + source, + on={'id': 'source_id'}, + when_matched=[ + WhenMatched(update={ + 'age': source_col('new_age'), + 'name': target_col('name'), + }), + ], + when_not_matched=[ + WhenNotMatched(insert={ + 'id': source_col('source_id'), + 'name': source_col('name'), + 'age': lit(0), + }), + ], +) +``` -</Tabs> +Conditions use SQL-style expressions with `s.` (source) and `t.` (target) +column prefixes. `WhenNotMatched` conditions may only reference source columns +(`s.*`). Install the SQL extra before using conditions: +`pip install pypaimon[sql]`. + +```python +messages = table_update.merge_into( + source, + on=['id'], + when_matched=[WhenMatched(update='*', condition='s.age > t.age')], + when_not_matched=[WhenNotMatched(insert='*', condition='s.age > 18')], +) +``` + +**Notes** + +- Multiple clauses are evaluated in order; the first matching condition wins. +- Matched clauses cannot update partition key columns, because cross-partition + row movement is not implemented. +- Blob columns are not written by `merge_into`: matched updates leave existing + blob files untouched, and not-matched inserts fill blob columns with `NULL`. ## Update Columns By Shards @@ -495,10 +442,6 @@ If you want to **compute a derived column** (or **update an existing column base This is useful for backfilling a newly added column, or recomputing a column from other columns. -<Tabs groupId="pypaimon-data-evolution-mode"> - -<TabItem value="batch" label="Batch"> - **Example: compute `d = c + b - a`** ```python @@ -580,70 +523,38 @@ commit.commit(commit_messages) commit.close() ``` -</TabItem> - -<TabItem value="stream" label="Stream"> - -```python -import pyarrow as pa -from pypaimon import CatalogFactory, Schema - -catalog = CatalogFactory.create({'warehouse': '/tmp/warehouse'}) -catalog.create_database('default', False) - -table_schema = pa.schema([ - ('a', pa.int32()), - ('b', pa.int32()), - ('c', pa.int32()), - ('d', pa.int32()), -]) - -schema = Schema.from_pyarrow_schema( - table_schema, - options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}, -) -catalog.create_table('default.t_stream', schema, False) -table = catalog.get_table('default.t_stream') - -# write initial data (a, b, c only) -write_builder = table.new_batch_write_builder() -write = write_builder.new_write().with_write_type(['a', 'b', 'c']) -commit = write_builder.new_commit() -write.write_arrow(pa.Table.from_pydict({'a': [1, 2], 'b': [10, 20], 'c': [100, 200]})) -commit.commit(write.prepare_commit()) -write.close() -commit.close() - -# stream shard update: each round uses a new commit_identifier -stream_builder = table.new_stream_write_builder() -table_update = stream_builder.new_update() -table_update.with_read_projection(['a', 'b', 'c']) -table_update.with_update_type(['d']) -table_commit = stream_builder.new_commit() - -upd = table_update.new_shard_updator(0, 1) -reader = upd.arrow_reader() - -for batch in iter(reader.read_next_batch, None): - a = batch.column('a').to_pylist() - b = batch.column('b').to_pylist() - c = batch.column('c').to_pylist() - d = [ci + bi - ai for ai, bi, ci in zip(a, b, c)] - - upd.update_by_arrow_batch( - pa.RecordBatch.from_pydict({'d': d}, schema=pa.schema([('d', pa.int32())])) - ) - -commit_messages = upd.prepare_commit() -table_commit.commit(commit_messages, commit_identifier=1) -``` - -</TabItem> - -</Tabs> - **Notes** - **Row order matters**: the batches you write must have the **same number of rows** as the batches you read, in the same order for that shard. - **Parallelism**: run multiple shards by calling `new_shard_updator(shard_idx, num_shards)` for each shard. + +## Stream Mode + +Data evolution also supports stream mode. The operation semantics are the same +as the batch APIs above; the main differences are the builder lifecycle and the +required `commit_identifier`. + +- Use `table.new_stream_write_builder()` instead of + `table.new_batch_write_builder()`. +- `StreamTableWrite`, `StreamTableUpdate`, and `StreamTableCommit` are reusable + across multiple rounds. +- Each round must use a monotonically increasing `commit_identifier`. +- Pass the same `commit_identifier` to the write prepare step or update method, + and to the corresponding commit call for that round. + +The API mapping is: + +| Batch API | Stream API | +| --- | --- | +| `write.prepare_commit()` | `write.prepare_commit(commit_identifier)` | +| `update.update_by_arrow_with_row_id(table)` | `update.update_by_arrow_with_row_id(table, commit_identifier)` | +| `update.update_by_predicate(predicate, assignments)` | `update.update_by_predicate(predicate, assignments, commit_identifier)` | +| `update.upsert_by_arrow_with_key(table, keys)` | `update.upsert_by_arrow_with_key(table, keys, commit_identifier)` | +| `update.merge_into(source, on=..., when_matched=..., when_not_matched=...)` | `update.merge_into(source, on=..., when_matched=..., when_not_matched=..., commit_identifier=...)` | +| `commit.commit(messages)` | `commit.commit(messages, commit_identifier)` | + +For shard updates, create the updater from `StreamTableUpdate` in the same way +as batch mode. `new_shard_updator(...)`, `arrow_reader()`, +`update_by_arrow_batch(...)`, and `prepare_commit()` stay the same; pass +`commit_identifier` when committing the returned messages. diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py b/paimon-python/pypaimon/ray/data_evolution_merge_into.py index ab20a56dc7..5b35da148e 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py @@ -282,39 +282,41 @@ def _build_datasets( if ctx.is_self_merge: if matched_specs and base_snapshot is not None: update_cols_union = _union_update_cols(matched_specs) - update_ds = build_self_merge_update_ds( + if update_cols_union: + update_ds = build_self_merge_update_ds( + target_identifier=target, + 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, + ray_remote_args=ray_remote_args, + ) + return update_ds, insert_ds, update_cols_union + + # Mirror Spark: matched/not-matched run as two independent joins + # (inner / left_anti). One unified left_outer join would force + # joined.materialize() to feed both branches, which can OOM on large merges. + if matched_specs and base_snapshot is not None: + update_cols_union = _union_update_cols(matched_specs) + if update_cols_union: + update_ds = build_matched_update_ds( target_identifier=target, + source_ds=source_ds, + target_on=ctx.target_on_cols, + source_on=ctx.source_on_cols, clauses=matched_specs, - target_field_names=ctx.full_target_field_names, + target_field_names=ctx.settable_field_names, target_pa_schema=ctx.update_pa_schema, update_cols=update_cols_union, catalog_options=ctx.catalog_options, + num_partitions=num_partitions, resolve_target_projection=_resolve_target_projection, snapshot_id=base_snapshot_id, ray_remote_args=ray_remote_args, ) - return update_ds, insert_ds, update_cols_union - - # Mirror Spark: matched/not-matched run as two independent joins - # (inner / left_anti). One unified left_outer join would force - # joined.materialize() to feed both branches, which can OOM on large merges. - if matched_specs and base_snapshot is not None: - update_cols_union = _union_update_cols(matched_specs) - update_ds = build_matched_update_ds( - target_identifier=target, - source_ds=source_ds, - target_on=ctx.target_on_cols, - source_on=ctx.source_on_cols, - clauses=matched_specs, - target_field_names=ctx.settable_field_names, - target_pa_schema=ctx.update_pa_schema, - update_cols=update_cols_union, - catalog_options=ctx.catalog_options, - num_partitions=num_partitions, - resolve_target_projection=_resolve_target_projection, - snapshot_id=base_snapshot_id, - ray_remote_args=ray_remote_args, - ) if not_matched_specs: # Insert writes the full target schema; SET spec only covers @@ -533,6 +535,8 @@ def _normalize_set_spec( f"SET spec references unknown target column " f"'{val.column}'" ) + if val.column == key: + continue result[key] = val elif isinstance(val, LiteralValue): result[key] = val @@ -549,6 +553,8 @@ def _normalize_set_spec( raise ValueError( f"SET spec references unknown target column '{ref}'" ) + if ref == key: + continue result[key] = TargetColumnRef(ref) else: result[key] = LiteralValue(val) diff --git a/paimon-python/pypaimon/table/data_evolution_merge_into.py b/paimon-python/pypaimon/table/data_evolution_merge_into.py new file mode 100644 index 0000000000..a775f501cc --- /dev/null +++ b/paimon-python/pypaimon/table/data_evolution_merge_into.py @@ -0,0 +1,690 @@ +################################################################################ +# 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. +################################################################################ + +"""Single-process MERGE INTO for Paimon data-evolution tables.""" + +from dataclasses import dataclass +from typing import Any, List, Optional, Sequence + +import pyarrow as pa +import pyarrow.compute as pc + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.ray.data_evolution_merge_into import ( + _blob_col_names, + _normalize_on, + _normalize_set_spec, + _resolve_target_projection, + _union_update_cols, + _validate_source_has_target_cols, +) +from pypaimon.ray.data_evolution_merge_join import _build_matched_transform +from pypaimon.ray.data_evolution_merge_transform import ( + OnSpec, + SourceColumnRef, + WhenMatched, + WhenNotMatched, + _NormalizedClause, + build_update_schema, + lit, + source_col, + target_col, + vectorized_insert_transform, +) +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.table_write import BatchTableWrite, StreamTableWrite +from pypaimon.write.table_update_by_row_id import TableUpdateByRowId + +__all__ = [ + "merge_into", + "WhenMatched", + "WhenNotMatched", + "source_col", + "target_col", + "lit", +] + + +@dataclass(frozen=True) +class _PrepareCtx: + target_on_cols: List[str] + source_on_cols: List[str] + settable_field_names: List[str] + full_target_field_names: List[str] + update_pa_schema: pa.Schema + full_pa_schema: pa.Schema + is_self_merge: bool = False + + +def merge_into( + target_table, + source: Any, + *, + on: OnSpec, + when_matched: Sequence[WhenMatched] = (), + when_not_matched: Sequence[WhenNotMatched] = (), + commit_user: str, + commit_identifier: int = BATCH_COMMIT_IDENTIFIER, +) -> List[CommitMessage]: + """Prepare MERGE INTO commit messages in the current Python process. + + The returned messages must be committed by the caller's matching + ``TableCommit``. + """ + base_snapshot = target_table.snapshot_manager().get_latest_snapshot() + source_table, matched_specs, not_matched_specs, ctx = _prepare( + target_table, + source, + list(when_matched), + list(when_not_matched), + on, + ) + + update_table, insert_table, update_cols_union = _build_tables( + target_table, + source_table, + matched_specs, + not_matched_specs, + ctx, + base_snapshot, + ) + + return _prepare_commit_messages( + target_table, + update_table, + insert_table, + update_cols_union, + base_snapshot, + commit_user, + commit_identifier, + ) + + +def _prepare( + target_table, + source, + when_matched, + when_not_matched, + on, +): + if not when_matched and not when_not_matched: + raise ValueError( + "At least one of when_matched or when_not_matched must be non-empty." + ) + for label, clauses in [("when_matched", when_matched), + ("when_not_matched", when_not_matched)]: + for i, clause in enumerate(clauses[:-1]): + if clause.condition is None: + raise ValueError( + "Only the last {} clause may omit its condition. " + "Clause at index {} has no condition, making subsequent " + "clauses unreachable.".format(label, i) + ) + + target_on_cols, source_on_cols = _normalize_on(on) + if not target_table.options.data_evolution_enabled(): + raise ValueError( + "merge_into requires 'data-evolution.enabled' = 'true' on target table." + ) + if not target_table.options.row_tracking_enabled(): + raise ValueError( + "merge_into requires 'row-tracking.enabled' = 'true' on target table." + ) + + blob_cols = _blob_col_names(target_table) + full_target_field_names = list(target_table.field_names) + settable_field_names = [ + c for c in full_target_field_names if c not in blob_cols + ] + on_map = dict(zip(target_on_cols, source_on_cols)) + + matched_specs = [ + _NormalizedClause( + spec=_normalize_set_spec( + c.update, + settable_field_names, + on_map, + ), + condition=c.condition, + ) + for c in when_matched + ] + if matched_specs and target_table.partition_keys: + partition_set = set(target_table.partition_keys) + for clause in matched_specs: + modified_partition_cols = partition_set & set(clause.spec.keys()) + if modified_partition_cols: + raise ValueError( + "merge_into does not support updating partition columns " + "{}; cross-partition row movement is not implemented." + .format(sorted(modified_partition_cols)) + ) + + has_condition = any( + c.condition is not None + for c in list(when_matched) + list(when_not_matched) + ) + if has_condition: + from pypaimon.ray.merge_condition import ( + _require_datafusion, + extract_target_columns, + ) + + _require_datafusion() + for c in when_not_matched: + if c.condition is not None: + t_refs = extract_target_columns(c.condition) + if t_refs: + raise ValueError( + "WhenNotMatched condition must not reference " + "target columns (t.*), but found: {}".format( + sorted(t_refs) + ) + ) + for c in list(when_matched) + list(when_not_matched): + if c.condition is not None: + blob_refs = extract_target_columns(c.condition) & blob_cols + if blob_refs: + raise ValueError( + "condition must not reference blob columns, " + "but found: {}".format(sorted(blob_refs)) + ) + + not_matched_specs = [] + for c in when_not_matched: + spec = _normalize_set_spec( + c.insert, + settable_field_names, + on_map, + allow_target_refs=False, + ) + for tk, sk in on_map.items(): + if tk in settable_field_names and tk not in spec: + spec[tk] = SourceColumnRef(sk) + not_matched_specs.append( + _NormalizedClause(spec=spec, condition=c.condition) + ) + + is_self_merge = _is_self_merge( + target_table, source, target_on_cols, source_on_cols + ) + if is_self_merge and not_matched_specs: + raise ValueError( + "Self-merge (source == target with ON _ROW_ID) does not " + "support WHEN NOT MATCHED clauses." + ) + + if is_self_merge: + source_table = None + source_col_names = set(full_target_field_names) | set(source_on_cols) + else: + source_table = _normalize_source(source) + _validate_source_on_cols(source_table, source_on_cols) + source_col_names = set(source_table.schema.names) + + _validate_source_has_target_cols( + source_col_names, matched_specs + not_matched_specs + ) + + if has_condition: + from pypaimon.ray.merge_condition import extract_columns + + target_names = set(full_target_field_names) + if is_self_merge: + target_names |= set(target_on_cols) + for c in list(when_matched) + list(when_not_matched): + if c.condition is not None: + for ref in extract_columns(c.condition): + prefix, col = ref.split(".", 1) + if prefix == "s" and col not in source_col_names: + raise ValueError( + "condition references unknown source column '{}'" + .format(col) + ) + if prefix == "t" and col not in target_names: + raise ValueError( + "condition references unknown target column '{}'" + .format(col) + ) + + from pypaimon.schema.data_types import PyarrowFieldParser + + full_pa_schema = PyarrowFieldParser.from_paimon_schema( + target_table.table_schema.fields + ) + update_pa_schema = pa.schema( + [full_pa_schema.field(c) for c in settable_field_names] + ) + ctx = _PrepareCtx( + target_on_cols=target_on_cols, + source_on_cols=source_on_cols, + settable_field_names=settable_field_names, + full_target_field_names=full_target_field_names, + update_pa_schema=update_pa_schema, + full_pa_schema=full_pa_schema, + is_self_merge=is_self_merge, + ) + return source_table, matched_specs, not_matched_specs, ctx + + +def _build_tables( + target_table, + source_table: Optional[pa.Table], + matched_specs: List[_NormalizedClause], + not_matched_specs: List[_NormalizedClause], + ctx: _PrepareCtx, + base_snapshot, +): + base_snapshot_id = base_snapshot.id if base_snapshot is not None else None + update_table = None + insert_table = None + update_cols_union: List[str] = [] + + if ctx.is_self_merge: + if matched_specs and base_snapshot is not None: + update_cols_union = _union_update_cols(matched_specs) + if update_cols_union: + update_table = _build_self_merge_update_table( + target_table, + matched_specs, + ctx, + update_cols_union, + base_snapshot_id, + ) + return update_table, insert_table, update_cols_union + + if matched_specs and base_snapshot is not None: + update_cols_union = _union_update_cols(matched_specs) + if update_cols_union: + update_table = _build_matched_update_table( + target_table, + source_table, + matched_specs, + ctx, + update_cols_union, + base_snapshot_id, + ) + + if not_matched_specs: + insert_table = _build_not_matched_insert_table( + target_table, + source_table, + not_matched_specs, + ctx, + base_snapshot_id, + target_empty=base_snapshot is None, + ) + + return update_table, insert_table, update_cols_union + + +def _build_self_merge_update_table( + target_table, + clauses: List[_NormalizedClause], + ctx: _PrepareCtx, + update_cols: Sequence[str], + snapshot_id: Optional[int], +) -> pa.Table: + row_id_name = SpecialFields.ROW_ID.name + needed_cols = set( + _resolve_target_projection( + clauses, + [row_id_name], + update_cols, + ctx.full_target_field_names, + ) + ) + for clause in clauses: + for value in clause.spec.values(): + if isinstance(value, SourceColumnRef): + needed_cols.add(value.column) + target_set = set(ctx.full_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 ctx.full_target_field_names if c in needed_cols + ] + + target = _read_table(target_table, projection=projection, snapshot_id=snapshot_id) + update_schema = build_update_schema( + ctx.update_pa_schema, update_cols, row_id_name + ) + if target.num_rows == 0: + return update_schema.empty_table() + + orig_names = list(target.schema.names) + target_renamed = _rename_with_prefix(target, "t.") + aliased = _add_self_merge_source_aliases( + target_renamed, orig_names, row_id_name + ) + transform = _build_matched_transform( + clauses, + on_map={row_id_name: row_id_name}, + on_pairs=[(row_id_name, row_id_name)], + update_cols=list(update_cols), + row_id_name=row_id_name, + update_schema=update_schema, + ) + return transform(aliased) + + +def _build_matched_update_table( + target_table, + source_table: pa.Table, + clauses: List[_NormalizedClause], + ctx: _PrepareCtx, + update_cols: Sequence[str], + snapshot_id: Optional[int], +) -> pa.Table: + row_id_name = SpecialFields.ROW_ID.name + needed_cols = _resolve_target_projection( + clauses, + ctx.target_on_cols, + update_cols, + ctx.settable_field_names, + ) + projection = [row_id_name] + [c for c in needed_cols if c != row_id_name] + target = _read_table(target_table, projection=projection, snapshot_id=snapshot_id) + update_schema = build_update_schema( + ctx.update_pa_schema, update_cols, row_id_name + ) + if target.num_rows == 0 or source_table.num_rows == 0: + return update_schema.empty_table() + + target_renamed = _rename_with_prefix(target, "t.") + source_renamed = _rename_with_prefix(source_table, "s.") + joined = target_renamed.join( + source_renamed, + keys=["t.{}".format(c) for c in ctx.target_on_cols], + right_keys=["s.{}".format(c) for c in ctx.source_on_cols], + join_type="inner", + ) + transform = _build_matched_transform( + clauses, + on_map=dict(zip(ctx.source_on_cols, ctx.target_on_cols)), + on_pairs=list(zip(ctx.source_on_cols, ctx.target_on_cols)), + update_cols=list(update_cols), + row_id_name=row_id_name, + update_schema=update_schema, + ) + return transform(joined) + + +def _build_not_matched_insert_table( + target_table, + source_table: pa.Table, + clauses: List[_NormalizedClause], + ctx: _PrepareCtx, + snapshot_id: Optional[int], + target_empty: bool = False, +) -> pa.Table: + source_renamed = _rename_with_prefix(source_table, "s.") + if source_renamed.num_rows == 0: + return ctx.full_pa_schema.empty_table() + + if target_empty: + unmatched = source_renamed + else: + target = _read_table( + target_table, + projection=list(ctx.target_on_cols), + snapshot_id=snapshot_id, + ) + if target.num_rows == 0: + unmatched = source_renamed + else: + target_renamed = _rename_with_prefix(target, "t.") + unmatched = source_renamed.join( + target_renamed, + keys=["s.{}".format(c) for c in ctx.source_on_cols], + right_keys=["t.{}".format(c) for c in ctx.target_on_cols], + join_type="left anti", + ) + + transform = _build_insert_transform( + clauses, ctx.full_target_field_names, ctx.full_pa_schema + ) + return transform(unmatched) + + +def _prepare_commit_messages( + table, + update_table: Optional[pa.Table], + insert_table: Optional[pa.Table], + update_cols_union: Sequence[str], + base_snapshot, + commit_user: str, + commit_identifier: int, +) -> List[CommitMessage]: + all_msgs: list = [] + if update_table is not None and update_table.num_rows > 0: + _validate_unique_row_ids(update_table) + update_snapshot_table = _copy_at_snapshot( + table, base_snapshot.id if base_snapshot is not None else None + ) + updater = TableUpdateByRowId( + update_snapshot_table, + commit_user, + commit_identifier, + ) + update_msgs = updater.update_columns( + update_table, list(update_cols_union) + ) + all_msgs.extend(update_msgs) + + if insert_table is not None and insert_table.num_rows > 0: + writer = _new_table_write(table, commit_user, commit_identifier) + try: + writer.write_arrow(insert_table) + if commit_identifier == BATCH_COMMIT_IDENTIFIER: + insert_msgs = writer.prepare_commit() + else: + insert_msgs = writer.prepare_commit(commit_identifier) + finally: + writer.close() + all_msgs.extend(insert_msgs) + + return all_msgs + + +def _new_table_write(table, commit_user: str, commit_identifier: int): + if commit_identifier == BATCH_COMMIT_IDENTIFIER: + return BatchTableWrite(table, commit_user) + return StreamTableWrite(table, commit_user) + + +def _build_insert_transform( + clauses: List[_NormalizedClause], + target_field_names: Sequence[str], + target_pa_schema: pa.Schema, +): + from pypaimon.ray.shuffle import _coerce_large_string_types + + 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 = 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 in prepared_clauses: + if remaining.num_rows == 0: + break + if rewritten is not None: + matched = _filter_batch( + remaining, rewritten, _pre_rewritten=True + ) + if matched.num_rows > 0: + parts.append( + vectorized_insert_transform( + matched, + spec, + target_field_names, + target_pa_schema, + ) + ) + if matched.num_rows < remaining.num_rows: + not_cond = "COALESCE(NOT ({}), TRUE)".format(rewritten) + remaining = _filter_batch( + remaining, not_cond, _pre_rewritten=True + ) + else: + remaining = remaining.slice(0, 0) + else: + parts.append( + vectorized_insert_transform( + remaining, + spec, + target_field_names, + target_pa_schema, + ) + ) + remaining = remaining.slice(0, 0) + if not parts: + return _coerce_large_string_types(target_pa_schema.empty_table()) + return _coerce_large_string_types(pa.concat_tables(parts)) + + return _transform + + +def _normalize_source(source: Any) -> pa.Table: + if isinstance(source, pa.Table): + return source + if _is_table_like(source): + snapshot = source.snapshot_manager().get_latest_snapshot() + snapshot_id = snapshot.id if snapshot is not None else None + return _read_table(source, snapshot_id=snapshot_id) + try: + import pandas as pd + except ImportError: + pd = None + if pd is not None and isinstance(source, pd.DataFrame): + return pa.Table.from_pandas(source, preserve_index=False) + raise TypeError( + "source must be a pyarrow.Table, a pandas.DataFrame, or a " + "Paimon table; got {}.".format(type(source).__name__) + ) + + +def _read_table(table, projection: Optional[List[str]] = None, + snapshot_id: Optional[int] = None) -> pa.Table: + read_table = _copy_at_snapshot(table, snapshot_id) + read_builder = read_table.new_read_builder() + if projection is not None: + read_builder.with_projection(projection) + scan = read_builder.new_scan() + return read_builder.new_read().to_arrow(scan.plan().splits()) + + +def _copy_at_snapshot(table, snapshot_id: Optional[int]): + if snapshot_id is None: + return table + return table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(snapshot_id)}) + + +def _validate_source_on_cols(source_table: pa.Table, on: Sequence[str]) -> None: + names = set(source_table.schema.names) + missing = [c for c in on if c not in names] + if missing: + raise ValueError( + "'on' columns {} missing from source schema {}.".format( + missing, list(names) + ) + ) + + +def _validate_unique_row_ids(update_table: pa.Table) -> None: + row_id_name = SpecialFields.ROW_ID.name + if ( + pc.count_distinct(update_table.column(row_id_name)).as_py() + != update_table.num_rows + ): + raise ValueError( + "MERGE matched multiple source rows to the same target _ROW_ID. " + "Deduplicate the source before merging." + ) + + +def _is_self_merge(target_table, source, target_on, source_on) -> bool: + row_id_name = SpecialFields.ROW_ID.name + return ( + _is_table_like(source) + and _same_table(target_table, source) + and target_on == [row_id_name] + and source_on == [row_id_name] + ) + + +def _is_table_like(obj) -> bool: + return hasattr(obj, "new_read_builder") and hasattr(obj, "snapshot_manager") + + +def _same_table(left, right) -> bool: + if left is right: + return True + return ( + getattr(left, "table_path", None) == getattr(right, "table_path", None) + and str(getattr(left, "identifier", "")) == str( + getattr(right, "identifier", "") + ) + and _current_branch(left) == _current_branch(right) + ) + + +def _current_branch(table): + current_branch = getattr(table, "current_branch", None) + return current_branch() if current_branch is not None else None + + +def _rename_with_prefix(table: pa.Table, prefix: str) -> pa.Table: + return table.rename_columns( + ["{}{}".format(prefix, name) for name in table.schema.names] + ) + + +def _add_self_merge_source_aliases( + target_renamed: pa.Table, orig_names: Sequence[str], row_id_name: str +) -> pa.Table: + columns = list(target_renamed.columns) + names = list(target_renamed.schema.names) + for orig in orig_names: + if orig == row_id_name: + continue + target_name = "t.{}".format(orig) + if target_name in names: + idx = names.index(target_name) + columns.append(columns[idx]) + names.append("s.{}".format(orig)) + return pa.table(columns, names=names) diff --git a/paimon-python/pypaimon/tests/table_merge_into_test.py b/paimon-python/pypaimon/tests/table_merge_into_test.py new file mode 100644 index 0000000000..f5ce03a28d --- /dev/null +++ b/paimon-python/pypaimon/tests/table_merge_into_test.py @@ -0,0 +1,582 @@ +################################################################################ +# 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. +################################################################################ + +import unittest + +import pyarrow as pa + +from pypaimon.table.data_evolution_merge_into import ( + WhenMatched, + WhenNotMatched, + lit, + source_col, + target_col, +) +from pypaimon.tests.data_evolution_test_helpers import ( + BatchModeMixin, + DataEvolutionTestBase, + StreamModeMixin, +) + +try: + import datafusion # noqa: F401 + + _HAS_DATAFUSION = True +except ImportError: + _HAS_DATAFUSION = False + +_SKIP_CONDITION = not _HAS_DATAFUSION +_SKIP_REASON = "datafusion not installed" + + +class TableMergeIntoTest(BatchModeMixin, DataEvolutionTestBase, unittest.TestCase): + + def _read_sorted(self, table): + return self._read_all(table).sort_by("id").to_pydict() + + def _read_projected_sorted(self, table, projection): + rb = table.new_read_builder().with_projection(projection) + return rb.new_read().to_arrow(rb.new_scan().plan().splits()).sort_by( + "id" + ).to_pydict() + + def _row_ids_by_id(self, table): + rows = self._read_projected_sorted(table, ["id", "_ROW_ID"]) + return dict(zip(rows["id"], rows["_ROW_ID"])) + + def _merge_and_commit(self, table, source, **kwargs): + wb = table.new_batch_write_builder() + table_update = wb.new_update() + msgs = table_update.merge_into(source, **kwargs) + commit = wb.new_commit() + commit.commit(msgs) + commit.close() + return msgs + + def test_table_merge_into_updates_and_inserts(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alice", "Bob"], + "age": pa.array([25, 30], type=pa.int32()), + "city": ["NYC", "LA"], + }, schema=self.pa_schema)) + + source = pa.Table.from_pydict({ + "id": pa.array([2, 3], type=pa.int32()), + "name": ["Bobby", "Cindy"], + "age": pa.array([31, 22], type=pa.int32()), + "city": ["LA2", "SF"], + }, schema=self.pa_schema) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_matched=[WhenMatched(update="*")], + when_not_matched=[WhenNotMatched(insert="*")], + ) + + self.assertTrue(msgs) + self.assertEqual( + { + "id": [1, 2, 3], + "name": ["Alice", "Bobby", "Cindy"], + "age": [25, 31, 22], + "city": ["NYC", "LA2", "SF"], + }, + self._read_sorted(target), + ) + + def test_table_merge_into_insert_only_on_empty_table(self): + target = self._create_table() + source = pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alice", "Bob"], + "age": pa.array([25, 30], type=pa.int32()), + "city": ["NYC", "LA"], + }, schema=self.pa_schema) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_not_matched=[WhenNotMatched(insert="*")], + ) + + self.assertTrue(msgs) + self.assertEqual( + { + "id": [1, 2], + "name": ["Alice", "Bob"], + "age": [25, 30], + "city": ["NYC", "LA"], + }, + self._read_sorted(target), + ) + + def test_table_merge_into_with_renamed_on_key_fills_insert_key(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + source = pa.Table.from_pydict({ + "source_id": pa.array([1, 2], type=pa.int32()), + "new_name": ["Alicia", "Bob"], + "new_age": pa.array([26, 30], type=pa.int32()), + "new_city": ["SF", "LA"], + }) + + self._merge_and_commit( + target, + source, + on={"id": "source_id"}, + when_matched=[ + WhenMatched(update={ + "name": source_col("new_name"), + "age": source_col("new_age"), + "city": source_col("new_city"), + }), + ], + when_not_matched=[ + WhenNotMatched(insert={ + "name": source_col("new_name"), + "age": source_col("new_age"), + "city": source_col("new_city"), + }), + ], + ) + + self.assertEqual( + { + "id": [1, 2], + "name": ["Alicia", "Bob"], + "age": [26, 30], + "city": ["SF", "LA"], + }, + self._read_sorted(target), + ) + + def test_table_merge_into_ignores_target_self_assignment(self): + for assignment in [target_col("name"), "t.name"]: + with self.subTest(assignment=assignment): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + row_ids_before = self._row_ids_by_id(target) + source = pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alicia"], + "age": pa.array([26], type=pa.int32()), + "city": ["SF"], + }, schema=self.pa_schema) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_matched=[WhenMatched(update={"name": assignment})], + ) + + self.assertEqual([], msgs) + self.assertEqual( + { + "id": [1], + "name": ["Alice"], + "age": [25], + "city": ["NYC"], + }, + self._read_sorted(target), + ) + self.assertEqual(row_ids_before, self._row_ids_by_id(target)) + + def test_table_merge_into_source_assignment_same_name_is_modified(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + source = pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alicia"], + "age": pa.array([26], type=pa.int32()), + "city": ["SF"], + }, schema=self.pa_schema) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_matched=[WhenMatched(update={"name": source_col("name")})], + ) + + self.assertTrue(msgs) + self.assertEqual(["Alicia"], self._read_sorted(target)["name"]) + self.assertEqual([25], self._read_sorted(target)["age"]) + + def test_table_merge_into_rejects_nested_assignment_key(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + source = pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alicia"], + "age": pa.array([26], type=pa.int32()), + "city": ["SF"], + }, schema=self.pa_schema) + + with self.assertRaisesRegex(ValueError, "unknown target column"): + target.new_batch_write_builder().new_update().merge_into( + source, + on=["id"], + when_matched=[WhenMatched(update={"name.first": lit("A")})], + ) + + def test_table_merge_into_rejects_partition_column_update(self): + target = self._create_table(partition_keys=["city"]) + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + source = pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([26], type=pa.int32()), + "city": ["SF"], + }, schema=self.pa_schema) + + with self.assertRaisesRegex(ValueError, "partition columns"): + target.new_batch_write_builder().new_update().merge_into( + source, + on=["id"], + when_matched=[WhenMatched(update={"city": source_col("city")})], + ) + + def test_table_merge_into_preserves_row_ids_for_matched_rows(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alice", "Bob"], + "age": pa.array([25, 30], type=pa.int32()), + "city": ["NYC", "LA"], + }, schema=self.pa_schema)) + row_ids_before = self._row_ids_by_id(target) + + self._merge_and_commit( + target, + pa.Table.from_pydict({ + "id": pa.array([2, 3], type=pa.int32()), + "name": ["Bobby", "Cindy"], + "age": pa.array([31, 22], type=pa.int32()), + "city": ["LA2", "SF"], + }, schema=self.pa_schema), + on=["id"], + when_matched=[WhenMatched(update="*")], + when_not_matched=[WhenNotMatched(insert="*")], + ) + + row_ids_after = self._row_ids_by_id(target) + self.assertEqual(row_ids_before[1], row_ids_after[1]) + self.assertEqual(row_ids_before[2], row_ids_after[2]) + self.assertGreater(row_ids_after[3], max(row_ids_before.values())) + + def test_table_merge_into_second_round_updates_first_round_insert(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + + self._merge_and_commit( + target, + pa.Table.from_pydict({ + "id": pa.array([2], type=pa.int32()), + "name": ["Bob"], + "age": pa.array([30], type=pa.int32()), + "city": ["LA"], + }, schema=self.pa_schema), + on=["id"], + when_not_matched=[WhenNotMatched(insert="*")], + ) + row_ids_after_insert = self._row_ids_by_id(target) + + self._merge_and_commit( + target, + pa.Table.from_pydict({ + "id": pa.array([2], type=pa.int32()), + "name": ["Bobby"], + "age": pa.array([31], type=pa.int32()), + "city": ["LA2"], + }, schema=self.pa_schema), + on=["id"], + when_matched=[WhenMatched(update="*")], + ) + + self.assertEqual( + { + "id": [1, 2], + "name": ["Alice", "Bobby"], + "age": [25, 31], + "city": ["NYC", "LA2"], + }, + self._read_sorted(target), + ) + self.assertEqual(row_ids_after_insert, self._row_ids_by_id(target)) + + def test_table_merge_into_accepts_paimon_source_table(self): + target = self._create_table() + source = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + self._write_arrow(source, pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alicia", "Bob"], + "age": pa.array([26, 30], type=pa.int32()), + "city": ["SF", "LA"], + }, schema=self.pa_schema)) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_matched=[WhenMatched(update="*")], + when_not_matched=[WhenNotMatched(insert="*")], + ) + + self.assertTrue(msgs) + self.assertEqual( + { + "id": [1, 2], + "name": ["Alicia", "Bob"], + "age": [26, 30], + "city": ["SF", "LA"], + }, + self._read_sorted(target), + ) + + @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON) + def test_table_merge_into_condition_clauses(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["old_1", "old_2"], + "age": pa.array([10, 40], type=pa.int32()), + "city": ["NYC", "LA"], + }, schema=self.pa_schema)) + + source = pa.Table.from_pydict({ + "id": pa.array([1, 2, 3], type=pa.int32()), + "name": ["new_1", "new_2", "new_3"], + "age": pa.array([15, 35, 50], type=pa.int32()), + "city": ["SF", "SEA", "DAL"], + }, schema=self.pa_schema) + + msgs = self._merge_and_commit( + target, + source, + on=["id"], + when_matched=[ + WhenMatched(update={"age": lit(99)}, condition="s.age > t.age"), + WhenMatched(update={"name": lit("kept")}), + ], + when_not_matched=[ + WhenNotMatched(insert="*", condition="s.age > 45"), + ], + ) + + self.assertTrue(msgs) + self.assertEqual( + { + "id": [1, 2, 3], + "name": ["old_1", "kept", "new_3"], + "age": [99, 40, 50], + "city": ["NYC", "LA", "DAL"], + }, + self._read_sorted(target), + ) + + def test_table_merge_into_duplicate_source_match_rejected(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + + source = pa.Table.from_pydict({ + "id": pa.array([1, 1], type=pa.int32()), + "name": ["A1", "A2"], + "age": pa.array([26, 27], type=pa.int32()), + "city": ["SF", "LA"], + }, schema=self.pa_schema) + + with self.assertRaisesRegex(ValueError, "multiple source rows"): + target.new_batch_write_builder().new_update().merge_into( + source, + on=["id"], + when_matched=[WhenMatched(update={"age": "s.age"})], + ) + + def test_table_merge_into_self_merge_by_row_id(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alice", "Bob"], + "age": pa.array([25, 30], type=pa.int32()), + "city": ["NYC", "LA"], + }, schema=self.pa_schema)) + + msgs = self._merge_and_commit( + target, + target, + on=["_ROW_ID"], + when_matched=[WhenMatched(update={"name": lit("updated")})], + ) + + self.assertTrue(msgs) + self.assertEqual(["updated", "updated"], self._read_sorted(target)["name"]) + + +class StreamTableMergeIntoTest(StreamModeMixin, DataEvolutionTestBase, unittest.TestCase): + + def _read_sorted(self, table): + return self._read_all(table).sort_by("id").to_pydict() + + def test_stream_table_update_merge_into(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + + source = pa.Table.from_pydict({ + "id": pa.array([1, 2], type=pa.int32()), + "name": ["Alicia", "Bob"], + "age": pa.array([26, 30], type=pa.int32()), + "city": ["SF", "LA"], + }, schema=self.pa_schema) + + base_snapshot_id = self._latest_snapshot_id(target) + wb = target.new_stream_write_builder() + cid = self._next_commit_id() + msgs = wb.new_update().merge_into( + source, + on=["id"], + when_matched=[WhenMatched(update="*")], + when_not_matched=[WhenNotMatched(insert="*")], + commit_identifier=cid, + ) + commit = wb.new_commit() + commit.commit(msgs, cid) + commit.close() + + self._assert_stream_builder_snapshots( + target, wb, base_snapshot_id, [cid] + ) + self.assertEqual( + { + "id": [1, 2], + "name": ["Alicia", "Bob"], + "age": [26, 30], + "city": ["SF", "LA"], + }, + self._read_sorted(target), + ) + + def test_stream_table_update_merge_into_multiple_rounds(self): + target = self._create_table() + self._write_arrow(target, pa.Table.from_pydict({ + "id": pa.array([1], type=pa.int32()), + "name": ["Alice"], + "age": pa.array([25], type=pa.int32()), + "city": ["NYC"], + }, schema=self.pa_schema)) + + base_snapshot_id = self._latest_snapshot_id(target) + wb = target.new_stream_write_builder() + commit = wb.new_commit() + + cid1 = self._next_commit_id() + msgs1 = wb.new_update().merge_into( + pa.Table.from_pydict({ + "id": pa.array([2], type=pa.int32()), + "name": ["Bob"], + "age": pa.array([30], type=pa.int32()), + "city": ["LA"], + }, schema=self.pa_schema), + on=["id"], + when_not_matched=[WhenNotMatched(insert="*")], + commit_identifier=cid1, + ) + commit.commit(msgs1, cid1) + + cid2 = self._next_commit_id() + msgs2 = wb.new_update().merge_into( + pa.Table.from_pydict({ + "id": pa.array([2], type=pa.int32()), + "name": ["Bobby"], + "age": pa.array([31], type=pa.int32()), + "city": ["LA2"], + }, schema=self.pa_schema), + on=["id"], + when_matched=[WhenMatched(update="*")], + commit_identifier=cid2, + ) + commit.commit(msgs2, cid2) + commit.close() + + self._assert_stream_builder_snapshots( + target, wb, base_snapshot_id, [cid1, cid2] + ) + self.assertEqual( + { + "id": [1, 2], + "name": ["Alice", "Bobby"], + "age": [25, 31], + "city": ["NYC", "LA2"], + }, + self._read_sorted(target), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/pypaimon/write/table_update.py b/paimon-python/pypaimon/write/table_update.py index fe1bbee83a..861396e0ba 100644 --- a/paimon-python/pypaimon/write/table_update.py +++ b/paimon-python/pypaimon/write/table_update.py @@ -16,7 +16,7 @@ # under the License. from collections import defaultdict -from typing import Any, List, Mapping, Optional, Tuple +from typing import Any, List, Mapping, Optional, Sequence, Tuple import pyarrow import pyarrow as pa @@ -185,6 +185,26 @@ class TableUpdate: self.table, self.commit_user, commit_identifier ).upsert(table, upsert_keys, self.update_cols) + def _merge_into( + self, + source: Any, + on, + when_matched: Sequence, + when_not_matched: Sequence, + commit_identifier: int, + ) -> List[CommitMessage]: + from pypaimon.table.data_evolution_merge_into import merge_into + + return merge_into( + self.table, + source, + on=on, + when_matched=when_matched, + when_not_matched=when_not_matched, + commit_user=self.commit_user, + commit_identifier=commit_identifier, + ) + def _update_by_predicate( self, predicate: Optional[Predicate], @@ -342,6 +362,23 @@ class BatchTableUpdate(TableUpdate): predicate, assignments, BATCH_COMMIT_IDENTIFIER ) + def merge_into( + self, + source: Any, + *, + on, + when_matched: Sequence = (), + when_not_matched: Sequence = (), + ) -> List[CommitMessage]: + """Prepare batch MERGE INTO commit messages.""" + return self._merge_into( + source, + on, + when_matched, + when_not_matched, + BATCH_COMMIT_IDENTIFIER, + ) + class StreamTableUpdate(TableUpdate): """Stream-mode table update; the same instance may drive many rounds, @@ -378,6 +415,24 @@ class StreamTableUpdate(TableUpdate): predicate, assignments, commit_identifier ) + def merge_into( + self, + source: Any, + *, + on, + when_matched: Sequence = (), + when_not_matched: Sequence = (), + commit_identifier: int, + ) -> List[CommitMessage]: + """Prepare stream MERGE INTO commit messages.""" + return self._merge_into( + source, + on, + when_matched, + when_not_matched, + commit_identifier, + ) + class ShardTableUpdator:
