You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by lz...@apache.org on 2026/03/18 07:38:03 UTC
(paimon) branch master updated: [python] Introduce TableUpsertByKey for data-evolution table (#7462)
This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new f0e6e9ff2c [python] Introduce TableUpsertByKey for data-evolution table (#7462)
f0e6e9ff2c is described below
commit f0e6e9ff2cf19b698ffea75acb50d0bd29ba220a
Author: Jingsong Lee <ji...@gmail.com>
AuthorDate: Wed Mar 18 15:37:57 2026 +0800
[python] Introduce TableUpsertByKey for data-evolution table (#7462)
---
docs/content/pypaimon/data-evolution.md | 128 ++++
.../pypaimon/tests/table_upsert_by_key_test.py | 739 +++++++++++++++++++++
paimon-python/pypaimon/write/table_update.py | 19 +
.../pypaimon/write/table_upsert_by_key.py | 378 +++++++++++
4 files changed, 1264 insertions(+)
diff --git a/docs/content/pypaimon/data-evolution.md b/docs/content/pypaimon/data-evolution.md
index 91714e52c7..b5daafae57 100644
--- a/docs/content/pypaimon/data-evolution.md
+++ b/docs/content/pypaimon/data-evolution.md
@@ -106,6 +106,134 @@ rb = table.new_read_builder().with_filter(pb.equal('_ROW_ID', 0))
result = rb.new_read().to_arrow(rb.new_scan().plan().splits())
```
+## Upsert By Key
+
+If you want to **upsert** (update-or-insert) rows by one or more business key columns — without manually providing
+`_ROW_ID` — use `upsert_by_arrow_with_key`. For each input row:
+
+- **Key matches** an existing row → update that row in place.
+- **No match** → append as a new row.
+
+**Requirements**
+
+- The table must have `data-evolution.enabled = true` and `row-tracking.enabled = true`.
+- All `upsert_keys` must exist in both the table schema and the input data.
+- For **partitioned tables**, the input data must contain all partition key columns. Partition keys are
+ **automatically stripped** from `upsert_keys` during matching (since each partition is processed independently),
+ so you do **not** need to include them in `upsert_keys`.
+
+**Example: basic upsert**
+
+```python
+import pyarrow as pa
+from pypaimon import CatalogFactory, Schema
+
+catalog = CatalogFactory.create({'warehouse': '/tmp/warehouse'})
+catalog.create_database('default', False)
+
+pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('age', pa.int32()),
+])
+schema = Schema.from_pyarrow_schema(
+ pa_schema,
+ options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'},
+)
+catalog.create_table('default.users', schema, False)
+table = catalog.get_table('default.users')
+
+# write initial data
+write_builder = table.new_batch_write_builder()
+write = write_builder.new_write()
+commit = write_builder.new_commit()
+write.write_arrow(pa.Table.from_pydict(
+ {'id': [1, 2], 'name': ['Alice', 'Bob'], 'age': [30, 25]},
+ schema=pa_schema,
+))
+commit.commit(write.prepare_commit())
+write.close()
+commit.close()
+
+# upsert: update id=1, insert id=3
+write_builder = table.new_batch_write_builder()
+table_update = write_builder.new_update()
+table_commit = write_builder.new_commit()
+
+upsert_data = pa.Table.from_pydict(
+ {'id': [1, 3], 'name': ['Alice_v2', 'Charlie'], 'age': [31, 28]},
+ schema=pa_schema,
+)
+cmts = table_update.upsert_by_arrow_with_key(upsert_data, upsert_keys=['id'])
+table_commit.commit(cmts)
+table_commit.close()
+
+# content should be:
+# id=1: name='Alice_v2', age=31 (updated)
+# id=2: name='Bob', age=25 (unchanged)
+# id=3: name='Charlie', age=28 (new)
+```
+
+**Example: partial-column upsert with `update_cols`**
+
+Combine `with_update_type` with `upsert_by_arrow_with_key` to update only specific columns for
+matched rows while still appending full rows for new keys:
+
+```python
+write_builder = table.new_batch_write_builder()
+table_update = write_builder.new_update().with_update_type(['age'])
+table_commit = write_builder.new_commit()
+
+upsert_data = pa.Table.from_pydict(
+ {'id': [1, 4], 'name': ['ignored', 'David'], 'age': [99, 22]},
+ schema=pa_schema,
+)
+cmts = table_update.upsert_by_arrow_with_key(upsert_data, upsert_keys=['id'])
+table_commit.commit(cmts)
+table_commit.close()
+
+# id=1: only 'age' is updated to 99; 'name' remains 'Alice_v2'
+# id=4: appended as a full new row
+```
+
+**Example: partitioned table with composite key**
+
+```python
+partitioned_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('region', pa.string()),
+])
+schema = Schema.from_pyarrow_schema(
+ partitioned_schema,
+ partition_keys=['region'],
+ options={'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'},
+)
+catalog.create_table('default.users_partitioned', schema, False)
+table = catalog.get_table('default.users_partitioned')
+
+# ... write initial data ...
+
+write_builder = table.new_batch_write_builder()
+table_update = write_builder.new_update()
+table_commit = write_builder.new_commit()
+
+upsert_data = pa.Table.from_pydict(
+ {'id': [1, 3], 'name': ['Alice_v2', 'Charlie'], 'region': ['US', 'EU']},
+ schema=partitioned_schema,
+)
+# upsert_keys=['id'] only; partition key 'region' is auto-stripped
+cmts = table_update.upsert_by_arrow_with_key(upsert_data, upsert_keys=['id'])
+table_commit.commit(cmts)
+table_commit.close()
+```
+
+**Notes**
+
+- Execution is driven **partition-by-partition**: only one partition's key set is loaded into memory at a time.
+- Duplicate keys in the input data will raise an error.
+- The upsert is atomic per commit — all matched updates and new appends are included in the same commit.
+
## Update Columns By Shards
If you want to **compute a derived column** (or **update an existing column based on other columns**) without providing
diff --git a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
new file mode 100644
index 0000000000..37511afd54
--- /dev/null
+++ b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
@@ -0,0 +1,739 @@
+"""
+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 os
+import shutil
+import tempfile
+import unittest
+import uuid
+
+import pyarrow as pa
+
+from pypaimon import CatalogFactory, Schema
+
+
+class TableUpsertByKeyTest(unittest.TestCase):
+ """Tests for TableUpsertByKey (upsert_by_arrow_with_key)."""
+
+ @classmethod
+ def setUpClass(cls):
+ cls.tempdir = tempfile.mkdtemp()
+ cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
+ cls.catalog = CatalogFactory.create({
+ 'warehouse': cls.warehouse
+ })
+ cls.catalog.create_database('default', True)
+
+ # Schema for non-partitioned table
+ cls.pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('age', pa.int32()),
+ ('city', pa.string()),
+ ])
+
+ # Schema for partitioned table
+ cls.partitioned_pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('age', pa.int32()),
+ ('region', pa.string()),
+ ])
+
+ cls.table_options = {
+ 'row-tracking.enabled': 'true',
+ 'data-evolution.enabled': 'true',
+ }
+
+ @classmethod
+ def tearDownClass(cls):
+ shutil.rmtree(cls.tempdir, ignore_errors=True)
+
+ # ------------------------------------------------------------------
+ # helpers
+ # ------------------------------------------------------------------
+
+ def _unique_name(self, prefix='upsert'):
+ return f'default.{prefix}_{uuid.uuid4().hex[:8]}'
+
+ def _create_table(self, pa_schema=None, partition_keys=None):
+ """Create a table and return the table object."""
+ pa_schema = pa_schema or self.pa_schema
+ name = self._unique_name()
+ if partition_keys:
+ schema = Schema.from_pyarrow_schema(
+ pa_schema, partition_keys=partition_keys, options=self.table_options
+ )
+ else:
+ schema = Schema.from_pyarrow_schema(pa_schema, options=self.table_options)
+ self.catalog.create_table(name, schema, False)
+ return self.catalog.get_table(name)
+
+ def _write(self, table, data):
+ """Write an Arrow table and commit."""
+ wb = table.new_batch_write_builder()
+ tw = wb.new_write()
+ tc = wb.new_commit()
+ tw.write_arrow(data)
+ tc.commit(tw.prepare_commit())
+ tw.close()
+ tc.close()
+
+ def _read_all(self, table):
+ """Read the full table and return as Arrow table."""
+ rb = table.new_read_builder()
+ tr = rb.new_read()
+ splits = rb.new_scan().plan().splits()
+ return tr.to_arrow(splits)
+
+ def _upsert(self, table, data, upsert_keys, update_cols=None):
+ """Run upsert_by_arrow_with_key and commit."""
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ if update_cols:
+ tu.with_update_type(update_cols)
+ msgs = tu.upsert_by_arrow_with_key(data, upsert_keys)
+ tc = wb.new_commit()
+ tc.commit(msgs)
+ tc.close()
+
+ # ==================================================================
+ # Basic upsert tests (non-partitioned)
+ # ==================================================================
+
+ def test_all_new_rows(self):
+ """Upsert into an empty table – every row is new."""
+ table = self._create_table()
+ data = pa.Table.from_pydict({
+ 'id': [1, 2, 3],
+ 'name': ['Alice', 'Bob', 'Carol'],
+ 'age': [25, 30, 35],
+ 'city': ['NYC', 'LA', 'Chicago'],
+ }, schema=self.pa_schema)
+
+ self._upsert(table, data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 3)
+ self.assertEqual(sorted(result['id'].to_pylist()), [1, 2, 3])
+
+ def test_all_matched_rows(self):
+ """Upsert where every row matches an existing key – pure update."""
+ table = self._create_table()
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2, 3],
+ 'name': ['Alice', 'Bob', 'Carol'],
+ 'age': [25, 30, 35],
+ 'city': ['NYC', 'LA', 'Chicago'],
+ }, schema=self.pa_schema)
+ self._write(table, initial)
+
+ update = pa.Table.from_pydict({
+ 'id': [1, 2, 3],
+ 'name': ['Alice2', 'Bob2', 'Carol2'],
+ 'age': [26, 31, 36],
+ 'city': ['NYC2', 'LA2', 'Chicago2'],
+ }, schema=self.pa_schema)
+ self._upsert(table, update, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 3)
+ names = sorted(zip(result['id'].to_pylist(), result['name'].to_pylist()))
+ self.assertEqual(names, [(1, 'Alice2'), (2, 'Bob2'), (3, 'Carol2')])
+
+ def test_mixed_update_and_append(self):
+ """Upsert where some rows match and some are new."""
+ table = self._create_table()
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema)
+ self._write(table, initial)
+
+ upsert_data = pa.Table.from_pydict({
+ 'id': [2, 3],
+ 'name': ['Bob_new', 'Carol'],
+ 'age': [31, 35],
+ 'city': ['LA_new', 'Chicago'],
+ }, schema=self.pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 3)
+ rows = sorted(
+ zip(result['id'].to_pylist(), result['name'].to_pylist()),
+ key=lambda x: x[0],
+ )
+ self.assertEqual(rows, [(1, 'Alice'), (2, 'Bob_new'), (3, 'Carol')])
+
+ def test_upsert_preserves_untouched_columns(self):
+ """Columns not listed in update_cols remain unchanged after upsert."""
+ table = self._create_table()
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema)
+ self._write(table, initial)
+
+ # Only update 'age' for matched row (id=1)
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1],
+ 'name': ['_ignored_'],
+ 'age': [99],
+ 'city': ['_ignored_'],
+ }, schema=self.pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'], update_cols=['age'])
+
+ result = self._read_all(table)
+ row = {k: result[k].to_pylist() for k in result.column_names}
+ idx = row['id'].index(1)
+ self.assertEqual(row['age'][idx], 99)
+ self.assertEqual(row['name'][idx], 'Alice') # unchanged
+ self.assertEqual(row['city'][idx], 'NYC') # unchanged
+
+ # ==================================================================
+ # Composite key tests
+ # ==================================================================
+
+ def test_composite_key_upsert(self):
+ """Upsert using two columns as composite key."""
+ table = self._create_table()
+ initial = pa.Table.from_pydict({
+ 'id': [1, 1, 2],
+ 'name': ['Alice', 'Alice', 'Bob'],
+ 'age': [25, 30, 35],
+ 'city': ['NYC', 'LA', 'Chicago'],
+ }, schema=self.pa_schema)
+ self._write(table, initial)
+
+ # (id=1, name='Alice') matches two existing rows with different cities;
+ # since keys contain (id, name) the match is unique per key column combo.
+ # But our data has (1, 'Alice') twice in initial → validation should still
+ # pass because we check input data, not table data.
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Carol'],
+ 'age': [99, 40],
+ 'city': ['Updated', 'Dallas'],
+ }, schema=self.pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id', 'name'])
+
+ result = self._read_all(table)
+ # 3 original + 1 new (2, Carol) = 4 rows; (1, Alice) matched the first occurrence
+ self.assertEqual(result.num_rows, 4)
+ ids = sorted(zip(
+ result['id'].to_pylist(),
+ result['name'].to_pylist(),
+ result['city'].to_pylist(),
+ ))
+ # (1,'Alice','LA') row should remain; (1,'Alice','NYC')-keyed row updated; (2,'Bob') unchanged; (2,'Carol') new
+ self.assertIn((2, 'Carol', 'Dallas'), ids)
+
+ # ==================================================================
+ # Sequential upsert tests
+ # ==================================================================
+
+ def test_sequential_upserts(self):
+ """Two sequential upserts: second upsert sees results of the first."""
+ table = self._create_table()
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema))
+
+ # First upsert: update id=1, add id=3
+ self._upsert(table, pa.Table.from_pydict({
+ 'id': [1, 3],
+ 'name': ['Alice_v2', 'Carol'],
+ 'age': [26, 35],
+ 'city': ['NYC', 'Chicago'],
+ }, schema=self.pa_schema), upsert_keys=['id'])
+
+ # Second upsert: update id=3 (just inserted), add id=4
+ self._upsert(table, pa.Table.from_pydict({
+ 'id': [3, 4],
+ 'name': ['Carol_v2', 'Dave'],
+ 'age': [36, 40],
+ 'city': ['Houston', 'Phoenix'],
+ }, schema=self.pa_schema), upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 4)
+ rows = {r: v for r, v in zip(result['id'].to_pylist(), result['name'].to_pylist())}
+ self.assertEqual(rows[1], 'Alice_v2')
+ self.assertEqual(rows[3], 'Carol_v2')
+ self.assertEqual(rows[4], 'Dave')
+
+ # ==================================================================
+ # Partitioned table tests
+ # ==================================================================
+
+ def test_partitioned_table_upsert(self):
+ """Upsert on a partitioned table with data in multiple partitions."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ # Write initial data across two partitions
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2, 3, 4],
+ 'name': ['Alice', 'Bob', 'Carol', 'Dave'],
+ 'age': [25, 30, 35, 40],
+ 'region': ['US', 'US', 'EU', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._write(table, initial)
+
+ # Upsert: update one row per partition, add one new
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 3, 5],
+ 'name': ['Alice_v2', 'Carol_v2', 'Eve'],
+ 'age': [26, 36, 45],
+ 'region': ['US', 'EU', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ # upsert_keys=['id'] only – partition key 'region' is auto-stripped
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 5)
+ rows = sorted(zip(
+ result['id'].to_pylist(),
+ result['name'].to_pylist(),
+ result['region'].to_pylist(),
+ ))
+ self.assertIn((1, 'Alice_v2', 'US'), rows)
+ self.assertIn((2, 'Bob', 'US'), rows)
+ self.assertIn((3, 'Carol_v2', 'EU'), rows)
+ self.assertIn((4, 'Dave', 'EU'), rows)
+ self.assertIn((5, 'Eve', 'EU'), rows)
+
+ def test_partitioned_upsert_single_partition(self):
+ """Upsert targeting only one partition does not affect the other."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2, 3],
+ 'name': ['Alice', 'Bob', 'Carol'],
+ 'age': [25, 30, 35],
+ 'region': ['US', 'US', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._write(table, initial)
+
+ # Only touch US partition
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 4],
+ 'name': ['Alice_v2', 'Dave'],
+ 'age': [26, 40],
+ 'region': ['US', 'US'],
+ }, schema=self.partitioned_pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 4)
+ eu_rows = [
+ (i, n) for i, n, r in zip(
+ result['id'].to_pylist(),
+ result['name'].to_pylist(),
+ result['region'].to_pylist(),
+ ) if r == 'EU'
+ ]
+ # EU partition unchanged
+ self.assertEqual(eu_rows, [(3, 'Carol')])
+
+ def test_partitioned_all_new_rows(self):
+ """Upsert into an empty partitioned table – all new."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ data = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'region': ['US', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._upsert(table, data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 2)
+
+ # ==================================================================
+ # Multiple data files
+ # ==================================================================
+
+ def test_upsert_across_multiple_data_files(self):
+ """Upsert matches rows that live in different data files (commits)."""
+ table = self._create_table()
+ # First commit
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema))
+ # Second commit
+ self._write(table, pa.Table.from_pydict({
+ 'id': [3, 4],
+ 'name': ['Carol', 'Dave'],
+ 'age': [35, 40],
+ 'city': ['Chicago', 'Houston'],
+ }, schema=self.pa_schema))
+
+ # Upsert hits both files: update id=2 (file1) and id=3 (file2), add id=5
+ upsert_data = pa.Table.from_pydict({
+ 'id': [2, 3, 5],
+ 'name': ['Bob_v2', 'Carol_v2', 'Eve'],
+ 'age': [31, 36, 45],
+ 'city': ['LA2', 'Chicago2', 'Phoenix'],
+ }, schema=self.pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 5)
+ rows = {r: v for r, v in zip(result['id'].to_pylist(), result['name'].to_pylist())}
+ self.assertEqual(rows[1], 'Alice')
+ self.assertEqual(rows[2], 'Bob_v2')
+ self.assertEqual(rows[3], 'Carol_v2')
+ self.assertEqual(rows[4], 'Dave')
+ self.assertEqual(rows[5], 'Eve')
+
+ # ==================================================================
+ # Validation / error tests
+ # ==================================================================
+
+ def test_empty_upsert_keys_raises(self):
+ """Empty upsert_keys list should raise ValueError."""
+ table = self._create_table()
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1], 'name': ['A'], 'age': [1], 'city': ['X'],
+ }, schema=self.pa_schema))
+
+ data = pa.Table.from_pydict({
+ 'id': [1], 'name': ['B'], 'age': [2], 'city': ['Y'],
+ }, schema=self.pa_schema)
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=[])
+ self.assertIn('must not be empty', str(ctx.exception))
+
+ def test_upsert_key_not_in_schema_raises(self):
+ """upsert_key not present in table schema should raise ValueError."""
+ table = self._create_table()
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1], 'name': ['A'], 'age': [1], 'city': ['X'],
+ }, schema=self.pa_schema))
+
+ data = pa.Table.from_pydict({
+ 'id': [1], 'name': ['B'], 'age': [2], 'city': ['Y'],
+ }, schema=self.pa_schema)
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['nonexistent'])
+ self.assertIn('not in table schema', str(ctx.exception))
+
+ def test_upsert_key_not_in_data_raises(self):
+ """upsert_key not present in input data should raise ValueError."""
+ table = self._create_table()
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1], 'name': ['A'], 'age': [1], 'city': ['X'],
+ }, schema=self.pa_schema))
+
+ # Input data missing 'city' column but upsert_keys=['city']
+ data = pa.Table.from_pydict({
+ 'id': [1], 'name': ['B'], 'age': [2],
+ })
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['city'])
+ self.assertIn('not in input data', str(ctx.exception))
+
+ def test_empty_data_raises(self):
+ """Empty input data should raise ValueError."""
+ table = self._create_table()
+ data = pa.Table.from_pydict({
+ 'id': pa.array([], type=pa.int32()),
+ 'name': pa.array([], type=pa.string()),
+ 'age': pa.array([], type=pa.int32()),
+ 'city': pa.array([], type=pa.string()),
+ })
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['id'])
+ self.assertIn('empty', str(ctx.exception))
+
+ def test_duplicate_keys_in_input_raises(self):
+ """Duplicate composite keys in input data should raise ValueError."""
+ table = self._create_table()
+ data = pa.Table.from_pydict({
+ 'id': [1, 1],
+ 'name': ['A', 'B'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema)
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['id'])
+ self.assertIn('duplicate', str(ctx.exception).lower())
+
+ def test_partitioned_table_missing_partition_col_in_data_raises(self):
+ """Input data missing partition column should raise ValueError."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ # Input data does NOT include the 'region' partition column
+ data = pa.Table.from_pydict({
+ 'id': [1], 'name': ['A'], 'age': [25],
+ })
+
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['id'])
+ self.assertIn('partition key', str(ctx.exception).lower())
+
+ def test_same_key_in_different_partitions(self):
+ """Same upsert key value in different partitions is valid."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ initial = pa.Table.from_pydict({
+ 'id': [1, 1],
+ 'name': ['Alice_US', 'Alice_EU'],
+ 'age': [25, 30],
+ 'region': ['US', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._write(table, initial)
+
+ # Upsert id=1 in both partitions (same key, different partitions)
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 1],
+ 'name': ['Alice_US_v2', 'Alice_EU_v2'],
+ 'age': [26, 31],
+ 'region': ['US', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, 2)
+ rows = sorted(zip(
+ result['id'].to_pylist(),
+ result['name'].to_pylist(),
+ result['region'].to_pylist(),
+ ), key=lambda x: x[2])
+ self.assertEqual(rows[0], (1, 'Alice_EU_v2', 'EU'))
+ self.assertEqual(rows[1], (1, 'Alice_US_v2', 'US'))
+
+ def test_invalid_update_cols_raises(self):
+ """update_cols referencing non-existent column should raise."""
+ table = self._create_table()
+ data = pa.Table.from_pydict({
+ 'id': [1], 'name': ['A'], 'age': [25], 'city': ['NYC'],
+ }, schema=self.pa_schema)
+
+ with self.assertRaises(ValueError):
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update().with_update_type(['nonexistent_col'])
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['id'])
+
+ # ==================================================================
+ # update_cols (partial column update) tests
+ # ==================================================================
+
+ def test_update_cols_partial_update(self):
+ """update_cols limits which columns are updated for matched rows."""
+ table = self._create_table()
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'city': ['NYC', 'LA'],
+ }, schema=self.pa_schema)
+ self._write(table, initial)
+
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['_X_', '_Y_'],
+ 'age': [99, 88],
+ 'city': ['_X_', '_Y_'],
+ }, schema=self.pa_schema)
+ # Only update 'age'
+ self._upsert(table, upsert_data, upsert_keys=['id'], update_cols=['age'])
+
+ result = self._read_all(table)
+ rows = sorted(zip(
+ result['id'].to_pylist(),
+ result['name'].to_pylist(),
+ result['age'].to_pylist(),
+ result['city'].to_pylist(),
+ ))
+ # name and city should be unchanged; age updated
+ self.assertEqual(rows[0], (1, 'Alice', 99, 'NYC'))
+ self.assertEqual(rows[1], (2, 'Bob', 88, 'LA'))
+
+ # ==================================================================
+ # Large data test
+ # ==================================================================
+
+ def test_large_table_upsert(self):
+ """Upsert on a table with many rows spread across multiple commits."""
+ table = self._create_table()
+ n = 200
+
+ # Write 200 rows in 2 commits
+ self._write(table, pa.Table.from_pydict({
+ 'id': list(range(n // 2)),
+ 'name': [f'Name_{i}' for i in range(n // 2)],
+ 'age': [20 + i for i in range(n // 2)],
+ 'city': [f'City_{i}' for i in range(n // 2)],
+ }, schema=self.pa_schema))
+ self._write(table, pa.Table.from_pydict({
+ 'id': list(range(n // 2, n)),
+ 'name': [f'Name_{i}' for i in range(n // 2, n)],
+ 'age': [20 + i for i in range(n // 2, n)],
+ 'city': [f'City_{i}' for i in range(n // 2, n)],
+ }, schema=self.pa_schema))
+
+ # Upsert: update every 10th row, add 10 new rows
+ update_ids = list(range(0, n, 10))
+ new_ids = list(range(n, n + 10))
+ all_ids = update_ids + new_ids
+
+ upsert_data = pa.Table.from_pydict({
+ 'id': all_ids,
+ 'name': [f'Upserted_{i}' for i in all_ids],
+ 'age': [1000 + i for i in all_ids],
+ 'city': [f'UCity_{i}' for i in all_ids],
+ }, schema=self.pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'])
+
+ result = self._read_all(table)
+ self.assertEqual(result.num_rows, n + 10)
+
+ result_map = {
+ rid: rname
+ for rid, rname in zip(result['id'].to_pylist(), result['name'].to_pylist())
+ }
+ # Verify updated rows
+ for uid in update_ids:
+ self.assertEqual(result_map[uid], f'Upserted_{uid}')
+ # Verify new rows
+ for nid in new_ids:
+ self.assertEqual(result_map[nid], f'Upserted_{nid}')
+ # Verify untouched rows
+ for i in range(n):
+ if i not in update_ids:
+ self.assertEqual(result_map[i], f'Name_{i}')
+
+ def test_non_data_evolution_table_raises(self):
+ """Upsert on a table without data-evolution enabled should raise."""
+ pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ])
+ name = self._unique_name()
+ # No data-evolution option
+ schema = Schema.from_pyarrow_schema(pa_schema)
+ self.catalog.create_table(name, schema, False)
+ table = self.catalog.get_table(name)
+
+ data = pa.Table.from_pydict({'id': [1], 'name': ['A']}, schema=pa_schema)
+ with self.assertRaises(ValueError) as ctx:
+ wb = table.new_batch_write_builder()
+ tu = wb.new_update()
+ tu.upsert_by_arrow_with_key(data, upsert_keys=['id'])
+ self.assertIn('data-evolution.enabled', str(ctx.exception))
+
+ def test_partitioned_update_cols_with_new_rows(self):
+ """When update_cols is set, matched rows update only those columns,
+ but new rows are appended with ALL columns via the standard write path."""
+ table = self._create_table(
+ pa_schema=self.partitioned_pa_schema,
+ partition_keys=['region'],
+ )
+
+ initial = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['Alice', 'Bob'],
+ 'age': [25, 30],
+ 'region': ['US', 'EU'],
+ }, schema=self.partitioned_pa_schema)
+ self._write(table, initial)
+
+ # Upsert with update_cols=['age'] -- the partition key 'region' is NOT
+ # in update_cols. id=1 matches (update), id=3 is new (append).
+ upsert_data = pa.Table.from_pydict({
+ 'id': [1, 3],
+ 'name': ['Alice_v2', 'Carol'],
+ 'age': [99, 50],
+ 'region': ['US', 'US'],
+ }, schema=self.partitioned_pa_schema)
+ self._upsert(table, upsert_data, upsert_keys=['id'],
+ update_cols=['age'])
+
+ result = self._read_all(table)
+ # 2 original + 1 new = 3
+ self.assertEqual(result.num_rows, 3)
+
+ ids = result['id'].to_pylist()
+ names = result['name'].to_pylist()
+ ages = result['age'].to_pylist()
+ regions = result['region'].to_pylist()
+
+ # id=1 was updated: only age changed
+ idx1 = ids.index(1)
+ self.assertEqual(ages[idx1], 99)
+ self.assertEqual(names[idx1], 'Alice') # unchanged
+ self.assertEqual(regions[idx1], 'US')
+
+ # id=2 was untouched
+ idx2 = ids.index(2)
+ self.assertEqual(ages[idx2], 30)
+ self.assertEqual(regions[idx2], 'EU')
+
+ # id=3 is new — all columns written (not partial)
+ idx3 = ids.index(3)
+ self.assertEqual(ages[idx3], 50)
+ self.assertEqual(names[idx3], 'Carol')
+ self.assertEqual(regions[idx3], 'US')
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/paimon-python/pypaimon/write/table_update.py b/paimon-python/pypaimon/write/table_update.py
index 9ad86aa95d..b26194eb55 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -27,6 +27,7 @@ from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.read.split import DataSplit
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+from pypaimon.write.table_upsert_by_key import TableUpsertByKey
from pypaimon.write.writer.data_writer import DataWriter
from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter
@@ -127,6 +128,24 @@ class TableUpdate:
update_by_row_id.update_columns(table, self.update_cols)
return update_by_row_id.commit_messages
+ def upsert_by_arrow_with_key(self, table: pa.Table, upsert_keys: List[str]) -> List[CommitMessage]:
+ """Upsert rows into an append-only table by one or more key columns.
+
+ For each row in the input Arrow table:
+ - If a row with the same composite upsert_keys value already exists -> update that row
+ in-place.
+ - If no matching row exists -> append as a new row.
+
+ Args:
+ table: Input Arrow table containing rows to upsert.
+ upsert_keys: One or more column names used together as a composite match key.
+
+ Returns:
+ List of CommitMessages to be committed.
+ """
+ upsert = TableUpsertByKey(self.table, self.commit_user)
+ return upsert.upsert(table, upsert_keys, self.update_cols)
+
class ShardTableUpdator:
diff --git a/paimon-python/pypaimon/write/table_upsert_by_key.py b/paimon-python/pypaimon/write/table_upsert_by_key.py
new file mode 100644
index 0000000000..3c895e4904
--- /dev/null
+++ b/paimon-python/pypaimon/write/table_upsert_by_key.py
@@ -0,0 +1,378 @@
+################################################################################
+# 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 typing import Any, Dict, List, Optional, Tuple
+
+import pyarrow as pa
+import logging
+
+from pypaimon.read.table_read import TableRead
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.write.commit_message import CommitMessage
+from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+from pypaimon.write.table_write import BatchTableWrite
+
+# Composite key is represented as a tuple of values
+_KeyTuple = Tuple[Any, ...]
+
+logger = logging.getLogger(__name__)
+
+
+class TableUpsertByKey:
+ """
+ Table upsert by one or more user-specified key columns for append-only tables.
+
+ For each row in the input Arrow table:
+ - If a row with the same upsert_keys composite value already exists → update that row
+ (in-place rewrite).
+ - If no matching row exists → append as a new row.
+
+ All upsert_keys must be columns present in both the input data and the table schema.
+ """
+
+ def __init__(self, table, commit_user: str):
+ from pypaimon.table.file_store_table import FileStoreTable
+
+ self.table: FileStoreTable = table
+ self.commit_user = commit_user
+
+ 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.
+ """
+ self._validate_inputs(data, upsert_keys, update_cols)
+
+ # Determine which columns to update
+ if update_cols is None or len(update_cols) == len(self.table.field_names):
+ effective_update_cols = None # means all columns
+ else:
+ effective_update_cols = update_cols
+
+ all_commit_messages: List[CommitMessage] = []
+
+ # Process each partition independently
+ for partition_spec, partition_data in self._group_by_partition(data):
+ msgs = self._upsert_partition(
+ partition_data, upsert_keys, partition_spec, effective_update_cols
+ )
+ all_commit_messages.extend(msgs)
+
+ return all_commit_messages
+
+ # ------------------------------------------------------------------
+ # Partition grouping
+ # ------------------------------------------------------------------
+
+ def _group_by_partition(
+ self, data: pa.Table
+ ) -> List[Tuple[Dict[str, Any], pa.Table]]:
+ """
+ Split *data* into ``(partition_spec, partition_rows)`` pairs.
+
+ For non-partitioned tables a single pair ``({}, data)`` is returned.
+ """
+ partition_keys = self.table.partition_keys
+ if not partition_keys:
+ return [({}, data)]
+
+ # Materialise partition columns once
+ part_columns = [data[k].to_pylist() for k in partition_keys]
+
+ # Discover unique partitions and collect row indices
+ seen_order: List[Tuple[Any, ...]] = [] # preserves insertion order
+ partition_to_indices: Dict[Tuple[Any, ...], List[int]] = {}
+ for i in range(data.num_rows):
+ part_tuple = tuple(col[i] for col in part_columns)
+ if part_tuple not in partition_to_indices:
+ seen_order.append(part_tuple)
+ partition_to_indices[part_tuple] = []
+ partition_to_indices[part_tuple].append(i)
+
+ result: List[Tuple[Dict[str, Any], pa.Table]] = []
+ for part_tuple in seen_order:
+ spec = dict(zip(partition_keys, part_tuple))
+ indices = pa.array(partition_to_indices[part_tuple], type=pa.int64())
+ result.append((spec, data.take(indices)))
+ return result
+
+ # ------------------------------------------------------------------
+ # Per-partition upsert
+ # ------------------------------------------------------------------
+
+ def _upsert_partition(
+ self,
+ partition_data: pa.Table,
+ upsert_keys: List[str],
+ partition_spec: Dict[str, Any],
+ update_cols: Optional[List[str]],
+ ) -> List[CommitMessage]:
+ """
+ Full upsert cycle scoped to a single partition.
+
+ Partition key columns are stripped from *upsert_keys* before matching
+ because all rows within this partition share the same partition values.
+
+ The scan reads partition data in batches and filters against the input
+ key set on-the-fly, so only matching key → _ROW_ID pairs are kept in
+ memory (instead of the entire partition's key set).
+ """
+ # Strip partition columns – they are constant inside one partition
+ partition_key_set = set(self.table.partition_keys)
+ match_keys = [k for k in upsert_keys if k not in partition_key_set]
+
+ # 1. Build input key tuples and a lookup set
+ key_columns = [partition_data[k].to_pylist() for k in match_keys]
+ input_key_tuples = [
+ tuple(col[i] for col in key_columns)
+ for i in range(partition_data.num_rows)
+ ]
+
+ # 2. Per-partition duplicate check
+ input_key_set = set(input_key_tuples)
+ if len(input_key_tuples) != len(input_key_set):
+ raise ValueError(
+ f"Input data contains duplicate values in upsert_keys columns "
+ f"{match_keys} within partition {partition_spec}."
+ )
+
+ # 3. Scan partition in batches, build key → _ROW_ID only for
+ # keys present in the input (avoids full-partition materialisation).
+ key_to_row_id = self._build_key_to_row_id_map(
+ match_keys, partition_spec, input_key_set
+ )
+
+ # 4. Split into matched (update) vs unmatched (append)
+ matched_indices: List[int] = []
+ new_indices: List[int] = []
+ for i, key_tuple in enumerate(input_key_tuples):
+ if key_tuple in key_to_row_id:
+ matched_indices.append(i)
+ else:
+ new_indices.append(i)
+
+ commit_messages: List[CommitMessage] = []
+
+ logger.info(
+ f"Upserting partition {partition_spec}: "
+ f"{len(matched_indices)} matched, {len(new_indices)} new"
+ )
+
+ # 5. In-place updates
+ if matched_indices:
+ commit_messages.extend(
+ self._do_updates(
+ partition_data, matched_indices,
+ input_key_tuples, key_to_row_id, update_cols
+ )
+ )
+
+ # 6. Appends
+ if new_indices:
+ commit_messages.extend(
+ self._do_appends(partition_data, new_indices)
+ )
+
+ return commit_messages
+
+ def _validate_inputs(self, data: pa.Table, upsert_keys: List[str],
+ update_cols: Optional[List[str]]):
+ """Validate inputs before processing."""
+ if not self.table.options.data_evolution_enabled():
+ raise ValueError(
+ "upsert_by_arrow_with_key requires 'data-evolution.enabled' = 'true'."
+ )
+
+ if not self.table.options.row_tracking_enabled():
+ raise ValueError(
+ "upsert_by_arrow_with_key requires 'row-tracking.enabled' = 'true'."
+ )
+
+ if not upsert_keys:
+ raise ValueError("upsert_keys must not be empty.")
+
+ for key in upsert_keys:
+ if key not in self.table.field_names:
+ raise ValueError(
+ f"upsert_key '{key}' is not in table schema fields: {self.table.field_names}"
+ )
+ if key not in data.column_names:
+ raise ValueError(
+ f"upsert_key '{key}' is not in input data columns: {data.column_names}"
+ )
+
+ # For partitioned tables, input data must contain partition columns
+ partition_keys = self.table.partition_keys
+ if partition_keys:
+ missing_in_data = [pk for pk in partition_keys if pk not in data.column_names]
+ if missing_in_data:
+ raise ValueError(
+ f"For partitioned tables, input data must contain all partition key "
+ f"columns. Missing: {missing_in_data}"
+ )
+
+ if update_cols is not None:
+ for col in update_cols:
+ if col not in self.table.field_names:
+ raise ValueError(
+ f"Column '{col}' in update_cols is not in table schema fields: "
+ f"{self.table.field_names}"
+ )
+
+ if data.num_rows == 0:
+ raise ValueError("Input data is empty.")
+
+ # NOTE: duplicate-key checking is deferred to _upsert_partition so
+ # that partition columns can be stripped first. The same non-partition
+ # key may legally appear in different partitions.
+
+ def _build_key_to_row_id_map(
+ self,
+ match_keys: List[str],
+ partition_spec: Optional[Dict[str, Any]],
+ input_key_set: set,
+ ) -> Dict[_KeyTuple, int]:
+ """
+ Scan the partition in batches and collect key → _ROW_ID only for
+ rows whose composite key is in *input_key_set*.
+
+ ``is_in`` predicates on each key column are pushed to the scan so
+ that files whose stats do not overlap the input values are pruned
+ entirely.
+
+ Args:
+ match_keys: Column names used as the composite match key
+ (partition columns already stripped).
+ partition_spec: Dict of partition_key → value to restrict the scan.
+ Pass an empty dict (or None) for non-partitioned tables.
+ input_key_set: Set of composite key tuples from the input data.
+ """
+ read_builder = self.table.new_read_builder()
+
+ if partition_spec:
+ predicate_builder = read_builder.new_predicate_builder()
+ sub_predicates = [
+ predicate_builder.equal(k, v)
+ for k, v in partition_spec.items()
+ ]
+ partition_predicate = predicate_builder.and_predicates(sub_predicates)
+ read_builder = read_builder.with_filter(partition_predicate)
+
+ scan = read_builder.new_scan()
+ splits = scan.plan().splits()
+ if not splits:
+ return {}
+
+ # Read only the key columns + _ROW_ID
+ key_fields = [self.table.field_dict[k] for k in match_keys]
+ read_type = key_fields + [SpecialFields.ROW_ID]
+
+ table_read = TableRead(
+ table=self.table, predicate=None, read_type=read_type
+ )
+
+ # Stream batches and filter against input_key_set on-the-fly
+ key_to_row_id: Dict[_KeyTuple, int] = {}
+ row_id_col = SpecialFields.ROW_ID.name
+ for batch in table_read.to_arrow_batch_reader(splits):
+ batch_key_cols = [batch.column(k).to_pylist() for k in match_keys]
+ batch_row_ids = batch.column(row_id_col).to_pylist()
+ for j, row_id in enumerate(batch_row_ids):
+ key_tuple = tuple(col[j] for col in batch_key_cols)
+ if key_tuple in input_key_set:
+ key_to_row_id[key_tuple] = row_id
+
+ return key_to_row_id
+
+ def _do_updates(
+ self,
+ data: pa.Table,
+ matched_indices: List[int],
+ input_key_tuples: List[_KeyTuple],
+ key_to_row_id: Dict[_KeyTuple, int],
+ update_cols: Optional[List[str]]
+ ) -> List[CommitMessage]:
+ """
+ Update rows that have matching upsert keys by rewriting them in-place.
+ """
+ matched_data = data.take(matched_indices)
+
+ # Build _ROW_ID values for matched rows
+ row_id_values = [key_to_row_id[input_key_tuples[i]] for i in matched_indices]
+ row_id_array = pa.array(row_id_values, type=pa.int64())
+
+ # Add _ROW_ID column
+ update_data = matched_data.append_column(SpecialFields.ROW_ID.name, row_id_array)
+
+ # Determine which columns to update
+ if update_cols is None:
+ cols_to_update = list(self.table.field_names)
+ else:
+ cols_to_update = list(update_cols)
+
+ # Use TableUpdateByRowId to do the actual in-place update
+ updater = TableUpdateByRowId(self.table, self.commit_user)
+ return updater.update_columns(update_data, cols_to_update)
+
+ def _do_appends(
+ self,
+ data: pa.Table,
+ new_indices: List[int],
+ ) -> List[CommitMessage]:
+ """
+ Append rows that have no matching upsert key.
+
+ New rows are written with all columns via the standard
+ BatchTableWrite API, which handles partition/bucket routing
+ automatically. ``update_cols`` only restricts which columns
+ are rewritten for *matched* rows.
+ """
+ new_data = data.take(new_indices)
+
+ # Reorder columns to match table schema order
+ all_ordered_cols = [c for c in self.table.field_names if c in new_data.column_names]
+ new_data = new_data.select(all_ordered_cols)
+
+ table_write = BatchTableWrite(self.table, self.commit_user)
+ try:
+ table_write.with_write_type(all_ordered_cols)
+ table_write.write_arrow(new_data)
+ return table_write.prepare_commit()
+ finally:
+ table_write.close()