| """ |
| 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_keeps_last(self): |
| """Duplicate keys in input data should keep the last occurrence.""" |
| 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)) |
| |
| # id=1 appears twice; the second row (name='A_last') should win |
| data = pa.Table.from_pydict({ |
| 'id': [1, 1], |
| 'name': ['A_first', 'A_last'], |
| 'age': [90, 91], |
| 'city': ['X', 'Y'], |
| }, schema=self.pa_schema) |
| self._upsert(table, data, upsert_keys=['id']) |
| |
| result = self._read_all(table) |
| rows = {r: (n, a, c) for r, n, a, c in zip( |
| result['id'].to_pylist(), |
| result['name'].to_pylist(), |
| result['age'].to_pylist(), |
| result['city'].to_pylist(), |
| )} |
| # id=1 updated with last duplicate row |
| self.assertEqual(rows[1], ('A_last', 91, 'Y')) |
| # id=2 unchanged |
| self.assertEqual(rows[2], ('Bob', 30, 'LA')) |
| |
| def test_duplicate_keys_all_new_keeps_last(self): |
| """Duplicate keys in input on empty table keeps the last occurrence.""" |
| table = self._create_table() |
| |
| # id=1 appears three times; last row should win |
| data = pa.Table.from_pydict({ |
| 'id': [1, 1, 1, 2], |
| 'name': ['A1', 'A2', 'A3', 'B'], |
| 'age': [10, 20, 30, 40], |
| 'city': ['X1', 'X2', 'X3', 'Y'], |
| }, schema=self.pa_schema) |
| self._upsert(table, data, upsert_keys=['id']) |
| |
| result = self._read_all(table) |
| self.assertEqual(result.num_rows, 2) |
| rows = {r: (n, a, c) for r, n, a, c in zip( |
| result['id'].to_pylist(), |
| result['name'].to_pylist(), |
| result['age'].to_pylist(), |
| result['city'].to_pylist(), |
| )} |
| self.assertEqual(rows[1], ('A3', 30, 'X3')) |
| self.assertEqual(rows[2], ('B', 40, 'Y')) |
| |
| def test_duplicate_keys_partitioned_keeps_last(self): |
| """Duplicate keys in a partitioned table keep the last per partition.""" |
| table = self._create_table( |
| pa_schema=self.partitioned_pa_schema, |
| partition_keys=['region'], |
| ) |
| self._write(table, pa.Table.from_pydict({ |
| 'id': [1, 2], |
| 'name': ['Alice', 'Bob'], |
| 'age': [25, 30], |
| 'region': ['US', 'EU'], |
| }, schema=self.partitioned_pa_schema)) |
| |
| # id=1 duplicated in US partition; id=2 duplicated in EU partition |
| data = pa.Table.from_pydict({ |
| 'id': [1, 1, 2, 2], |
| 'name': ['A_first', 'A_last', 'B_first', 'B_last'], |
| 'age': [50, 51, 60, 61], |
| 'region': ['US', 'US', 'EU', 'EU'], |
| }, schema=self.partitioned_pa_schema) |
| self._upsert(table, data, upsert_keys=['id']) |
| |
| result = self._read_all(table) |
| self.assertEqual(result.num_rows, 2) |
| rows = {(r, reg): (n, a) for r, n, a, reg in zip( |
| result['id'].to_pylist(), |
| result['name'].to_pylist(), |
| result['age'].to_pylist(), |
| result['region'].to_pylist(), |
| )} |
| self.assertEqual(rows[(1, 'US')], ('A_last', 51)) |
| self.assertEqual(rows[(2, 'EU')], ('B_last', 61)) |
| |
| 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() |
