| # 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 random | |
| import string | |
| import threading | |
| import unittest | |
| from unittest import mock | |
| import pyarrow as pa | |
| import pytest | |
| from pypaimon.read.read_builder import ReadBuilder | |
| from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER | |
| from pypaimon.tests.data_evolution_test_helpers import ( | |
| BatchModeMixin, | |
| DataEvolutionTestBase, | |
| StreamModeMixin, | |
| ) | |
| from pypaimon.write.table_update import BatchTableUpdate | |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId | |
| from pypaimon.write.write_builder import BatchWriteBuilder, StreamWriteBuilder | |
| # ====================================================================== | |
| # Shared base for batch & stream table-update tests | |
| # ====================================================================== | |
| @pytest.mark.parametrize("stream, commit_identifier", [ | |
| (False, BATCH_COMMIT_IDENTIFIER), | |
| (True, 123), | |
| ]) | |
| def test_write_builder_creates_row_id_updater_with_commit_identity( | |
| stream, commit_identifier): | |
| table = mock.Mock() | |
| table.options.commit_user_prefix.return_value = None | |
| files_info = mock.Mock( | |
| snapshot_id=42, | |
| first_row_ids=[0], | |
| first_row_id_index={}, | |
| valid_row_id_ranges=[], | |
| ) | |
| builder = StreamWriteBuilder(table) if stream else BatchWriteBuilder(table) | |
| update = builder.new_update() | |
| if stream: | |
| updater = update.new_update_by_row_id( | |
| commit_identifier, _precomputed_files_info=files_info) | |
| else: | |
| updater = update.new_update_by_row_id( | |
| _precomputed_files_info=files_info) | |
| assert updater.table is table | |
| assert updater.commit_user == update.commit_user == builder.commit_user | |
| assert updater.commit_identifier == commit_identifier | |
| assert updater.snapshot_id == 42 | |
| table.new_read_builder.assert_not_called() | |
| def test_batch_row_id_update_batches_reuse_file_index(): | |
| table = mock.MagicMock() | |
| table.field_names = ["value"] | |
| batches = [ | |
| pa.table({"_ROW_ID": [0], "value": [10]}), | |
| pa.table({"_ROW_ID": [1], "value": [20]}), | |
| ] | |
| with mock.patch.object( | |
| BatchTableUpdate, "_new_row_id_updater") as factory: | |
| updater = factory.return_value | |
| updater.commit_messages = [] | |
| def update_columns(batch, columns): | |
| updater.commit_messages.append((batch, columns)) | |
| return updater.commit_messages | |
| updater.update_columns.side_effect = update_columns | |
| messages = ( | |
| BatchTableUpdate(table, "user") | |
| .with_update_type(["value"]) | |
| .update_by_arrow_batches_with_row_id(iter(batches)) | |
| ) | |
| factory.assert_called_once_with(BATCH_COMMIT_IDENTIFIER) | |
| assert updater.update_columns.call_count == 2 | |
| assert messages == [ | |
| (batches[0], ["value"]), | |
| (batches[1], ["value"]), | |
| ] | |
| def test_batch_row_id_update_detects_overlap_before_write(): | |
| updater = TableUpdateByRowId.__new__(TableUpdateByRowId) | |
| updater.table = mock.Mock(field_names=['age']) | |
| updater.commit_messages = [] | |
| updater._updated_first_row_ids_by_column = {} | |
| def route(data): | |
| return data.append_column( | |
| TableUpdateByRowId.FIRST_ROW_ID_COLUMN, | |
| pa.array([0], type=pa.int64()), | |
| ) | |
| with mock.patch.object( | |
| updater, '_calculate_first_row_id', side_effect=route): | |
| with mock.patch.object(updater, '_write_by_first_row_id') as write: | |
| updater.update_columns( | |
| pa.table({'_ROW_ID': [0], 'age': [26]}), ['age']) | |
| with pytest.raises(ValueError, match='overlapping first_row_ids'): | |
| updater.update_columns( | |
| pa.table({'_ROW_ID': [1], 'age': [31]}), ['age']) | |
| assert write.call_count == 1 | |
| class _TableUpdateTestBase(DataEvolutionTestBase): | |
| """Shared tests for ``TableUpdate.update_by_arrow_with_row_id``. | |
| Concrete subclasses must inherit from :class:`unittest.TestCase` AND one | |
| of :class:`_BatchModeMixin` / :class:`_StreamModeMixin`, which add the | |
| operation-specific primitive ``_apply_update(tu, data, cid)`` on top of | |
| the framework primitives provided by | |
| :class:`DataEvolutionTestBase`/:class:`BatchModeMixin`/:class:`StreamModeMixin`. | |
| """ | |
| # ------------------------------------------------------------------ | |
| # Operation-specific primitive (overridden by mixins) | |
| # ------------------------------------------------------------------ | |
| def _apply_update(self, table_update, data, cid): | |
| raise NotImplementedError | |
| def _apply_update_by_predicate( | |
| self, table_update, predicate, assignments, cid, | |
| read_columns=None): | |
| raise NotImplementedError | |
| def _apply_delete_by_predicate(self, table_update, predicate, cid): | |
| raise NotImplementedError | |
| def _apply_delete_by_row_id(self, table_update, row_ids, cid): | |
| raise NotImplementedError | |
| # ------------------------------------------------------------------ | |
| # Helpers built on the primitives | |
| # ------------------------------------------------------------------ | |
| def _create_seeded_table(self, partition_keys=None, options=None): | |
| """Create the canonical 5-row / 2-file table used by most tests. | |
| Layout (row_id → row): | |
| 0: (1, Alice, 25, NYC) | |
| 1: (2, Bob, 30, LA) | |
| 2: (3, Charlie, 35, Chicago) | |
| 3: (4, David, 40, Houston) | |
| 4: (5, Eve, 45, Phoenix) | |
| """ | |
| table = self._create_table( | |
| partition_keys=partition_keys, | |
| options=options, | |
| ) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2], | |
| 'name': ['Alice', 'Bob'], | |
| 'age': [25, 30], | |
| 'city': ['NYC', 'LA'], | |
| }, schema=self.pa_schema)) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [3, 4, 5], | |
| 'name': ['Charlie', 'David', 'Eve'], | |
| 'age': [35, 40, 45], | |
| 'city': ['Chicago', 'Houston', 'Phoenix'], | |
| }, schema=self.pa_schema)) | |
| return table | |
| def _create_global_indexed_table_for_predicate_update(self, extra_options=None): | |
| options = dict(self.table_options) | |
| options.update({ | |
| 'global-index.enabled': 'true', | |
| 'bucket': '-1', | |
| 'file.format': 'parquet', | |
| }) | |
| options.update(extra_options or {}) | |
| table = self._create_table(options=options) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2], | |
| 'name': ['old', 'indexed'], | |
| 'age': [10, 15], | |
| 'city': ['NYC', 'LA'], | |
| }, schema=self.pa_schema)) | |
| self.assertEqual( | |
| 1, | |
| table.create_global_index( | |
| 'name', | |
| options={'sorted-index.records-per-range': '1000'}, | |
| ), | |
| ) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [3, 4], | |
| 'name': ['new', 'other'], | |
| 'age': [20, 30], | |
| 'city': ['LA', 'SF'], | |
| }, schema=self.pa_schema)) | |
| return table | |
| def _do_update(self, table, data, columns): | |
| """End-to-end ``update_by_arrow_with_row_id`` + commit. Returns the | |
| commit messages so callers can inspect produced files.""" | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(columns) | |
| cid = self._next_commit_id() | |
| msgs = self._apply_update(tu, data, cid) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| return msgs | |
| def _do_update_by_predicate( | |
| self, table, predicate, assignments, read_columns=None): | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| cid = self._next_commit_id() | |
| msgs = self._apply_update_by_predicate( | |
| tu, | |
| predicate, | |
| assignments, | |
| cid, | |
| read_columns, | |
| ) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| return msgs | |
| @staticmethod | |
| def _list_table_files(table): | |
| return { | |
| os.path.relpath(os.path.join(root, name), table.table_path) | |
| for root, _dirs, files in os.walk(table.table_path) | |
| for name in files | |
| } | |
| def _do_delete_by_predicate(self, table, predicate): | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| cid = self._next_commit_id() | |
| msgs = self._apply_delete_by_predicate(tu, predicate, cid) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| return msgs | |
| def _do_delete_by_row_id(self, table, row_ids): | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| cid = self._next_commit_id() | |
| msgs = self._apply_delete_by_row_id(tu, row_ids, cid) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| return msgs | |
| def _create_seeded_deletion_vector_table(self, partition_keys=None): | |
| options = dict(self.table_options) | |
| options['deletion-vectors.enabled'] = 'true' | |
| table = self._create_seeded_table( | |
| partition_keys=partition_keys, | |
| options=options, | |
| ) | |
| return table | |
| # ================================================================== | |
| # Shared tests (run under both batch and stream modes) | |
| # ================================================================== | |
| def test_update_existing_column(self): | |
| """Update a single column across both files (row_ids unordered).""" | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 0, 2, 3, 4], | |
| 'age': [31, 26, 36, 39, 42], | |
| }), ['age']) | |
| self.assertEqual( | |
| [26, 31, 36, 39, 42], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| @pytest.mark.python_plan | |
| def test_update_by_predicate(self): | |
| table = self._create_seeded_table() | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| pb = tu.new_predicate_builder() | |
| predicate = pb.greater_or_equal('age', 35) | |
| cid = self._next_commit_id() | |
| msgs = self._apply_update_by_predicate( | |
| tu, | |
| predicate, | |
| {'age': 99, 'city': 'Updated'}, | |
| cid, | |
| ) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| # Table scans are unordered; compare in stable business-key order. | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual( | |
| [25, 30, 99, 99, 99], | |
| result['age'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'Updated', 'Updated', 'Updated'], | |
| result['city'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| result['name'].to_pylist(), | |
| ) | |
| @pytest.mark.python_write | |
| def test_update_by_predicate_processes_one_file_group_at_a_time(self): | |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId | |
| table = self._create_seeded_table() | |
| self._do_update( | |
| table, | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}), | |
| ['age'], | |
| ) | |
| splits = table.new_read_builder().new_scan().plan_for_write().splits() | |
| self.assertEqual(1, len(splits)) | |
| group_sizes = [] | |
| original = TableUpdateByRowId.update_columns | |
| def capture(updater, data, columns): | |
| group_sizes.append(data.num_rows) | |
| return original(updater, data, columns) | |
| with mock.patch.object(TableUpdateByRowId, 'update_columns', capture): | |
| self._do_update_by_predicate(table, None, {'city': 'Updated'}) | |
| self.assertEqual([2, 3], group_sizes) | |
| self.assertEqual( | |
| [26, 30, 35, 40, 45], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['Updated'] * 5, | |
| self._read_all(table)['city'].to_pylist(), | |
| ) | |
| @pytest.mark.python_write | |
| def test_predicate_update_preserves_groups_after_later_failure(self): | |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId | |
| table = self._create_seeded_table() | |
| self._do_update( | |
| table, | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}), | |
| ['age'], | |
| ) | |
| before_files = self._list_table_files(table) | |
| calls = 0 | |
| original = TableUpdateByRowId.update_columns | |
| def fail_second_group(updater, data, columns): | |
| nonlocal calls | |
| calls += 1 | |
| if calls == 2: | |
| raise RuntimeError("second group failed") | |
| messages = original(updater, data, columns) | |
| return messages | |
| with mock.patch.object( | |
| TableUpdateByRowId, 'update_columns', fail_second_group): | |
| with self.assertRaisesRegex(RuntimeError, "second group failed"): | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| {'city': 'Updated'}, | |
| ) | |
| self.assertEqual(2, calls) | |
| self.assertLess(before_files, self._list_table_files(table)) | |
| @pytest.mark.python_write | |
| def test_callable_update_preserves_groups_after_later_failure(self): | |
| table = self._create_seeded_table() | |
| before_files = self._list_table_files(table) | |
| calls = 0 | |
| def fail_second_group(rows): | |
| nonlocal calls | |
| calls += 1 | |
| if calls == 2: | |
| raise RuntimeError("second callback failed") | |
| return pa.compute.add(rows['age'], 1) | |
| with self.assertRaisesRegex(RuntimeError, "second callback failed"): | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| {'age': fail_second_group}, | |
| read_columns=['age'], | |
| ) | |
| self.assertEqual(2, calls) | |
| self.assertLess(before_files, self._list_table_files(table)) | |
| def test_array_assignment_spans_file_groups(self): | |
| table = self._create_seeded_table() | |
| self._do_update( | |
| table, | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}), | |
| ['age'], | |
| ) | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| {'age': pa.array([101, 102, 103, 104, 105])}, | |
| ) | |
| self.assertEqual( | |
| [101, 102, 103, 104, 105], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_no_match_is_noop(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_update_by_predicate( | |
| table, | |
| pb.greater_than('age', 100), | |
| {'age': 1}, | |
| ) | |
| self.assertEqual([], msgs) | |
| self.assertEqual( | |
| [25, 30, 35, 40, 45], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| @pytest.mark.python_write | |
| def test_literal_predicate_update_projects_only_row_id(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| projections = [] | |
| with_projection = ReadBuilder.with_projection | |
| def capture_projection(builder, projection): | |
| projections.append(projection) | |
| return with_projection(builder, projection) | |
| with mock.patch.object( | |
| ReadBuilder, 'with_projection', capture_projection): | |
| self._do_update_by_predicate( | |
| table, | |
| pb.greater_or_equal('age', 35), | |
| {'city': 'Updated'}, | |
| ) | |
| self.assertEqual([['_ROW_ID']], projections) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'Updated', 'Updated', 'Updated'], | |
| self._read_all(table).sort_by('id')['city'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_accepts_array_chunked_and_scalar_values(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self._do_update_by_predicate( | |
| table, | |
| pb.greater_or_equal('age', 35), | |
| { | |
| 'age': pa.array([101, 102, 103], type=pa.int64()), | |
| 'city': pa.chunked_array([ | |
| pa.array(['Chicago_v2']), | |
| pa.array(['Houston_v2', 'Phoenix_v2']), | |
| ]), | |
| 'name': pa.scalar('patched'), | |
| }, | |
| ) | |
| result = self._read_all(table) | |
| rows = { | |
| row_id: (name, age, city) | |
| for row_id, name, age, city in zip( | |
| result['id'].to_pylist(), | |
| result['name'].to_pylist(), | |
| result['age'].to_pylist(), | |
| result['city'].to_pylist(), | |
| ) | |
| } | |
| self.assertEqual( | |
| { | |
| 1: ('Alice', 25, 'NYC'), | |
| 2: ('Bob', 30, 'LA'), | |
| 3: ('patched', 101, 'Chicago_v2'), | |
| 4: ('patched', 102, 'Houston_v2'), | |
| 5: ('patched', 103, 'Phoenix_v2'), | |
| }, | |
| rows, | |
| ) | |
| def test_update_by_predicate_accepts_callable_assignments(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| def increment_age(rows): | |
| self.assertEqual( | |
| ['age', 'city', '_ROW_ID'], rows.column_names | |
| ) | |
| return pa.compute.add(rows['age'], 1) | |
| self._do_update_by_predicate( | |
| table, | |
| pb.greater_or_equal('age', 35), | |
| { | |
| 'age': increment_age, | |
| 'city': lambda rows: pa.compute.utf8_upper(rows['city']), | |
| }, | |
| read_columns=['age', 'city'], | |
| ) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([25, 30, 36, 41, 46], result['age'].to_pylist()) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'CHICAGO', 'HOUSTON', 'PHOENIX'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_updates_all_rows_when_predicate_is_none(self): | |
| table = self._create_seeded_table() | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| { | |
| 'age': pa.scalar(7, type=pa.int64()), | |
| 'city': None, | |
| 'name': 'UPDATED', | |
| }, | |
| ) | |
| result = self._read_all(table) | |
| self.assertEqual([7, 7, 7, 7, 7], result['age'].to_pylist()) | |
| self.assertEqual([None, None, None, None, None], | |
| result['city'].to_pylist()) | |
| self.assertEqual( | |
| ['UPDATED'] * 5, | |
| result['name'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_streams_callable_by_file_group(self): | |
| table = self._create_seeded_table() | |
| group_sizes = [] | |
| def increment_age(rows): | |
| group_sizes.append(rows.num_rows) | |
| return pa.compute.add(rows['age'], 1) | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| {'age': increment_age}, | |
| read_columns=['age'], | |
| ) | |
| self.assertEqual([2, 3], group_sizes) | |
| self.assertEqual( | |
| [26, 31, 36, 41, 46], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_requires_callable_read_columns(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaisesRegex( | |
| ValueError, "Callable assignments require read_columns"): | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {'age': lambda rows: rows['age']}, | |
| ) | |
| def test_update_by_predicate_rejects_unused_read_columns(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaisesRegex( | |
| ValueError, "read_columns requires a callable assignment"): | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {'age': 26}, | |
| read_columns=['age'], | |
| ) | |
| def test_update_by_predicate_rejects_callable_with_array_assignment(self): | |
| table = self._create_seeded_table() | |
| with self.assertRaisesRegex( | |
| ValueError, "cannot be combined with Arrow array"): | |
| self._do_update_by_predicate( | |
| table, | |
| None, | |
| { | |
| 'age': lambda rows: rows['age'], | |
| 'city': pa.array(['A', 'B', 'C', 'D', 'E']), | |
| }, | |
| read_columns=['age'], | |
| ) | |
| def test_no_match_does_not_invoke_callable(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| def unexpected(_rows): | |
| self.fail("Callable should not run without matched rows.") | |
| messages = self._do_update_by_predicate( | |
| table, | |
| pb.greater_than('age', 100), | |
| {'age': unexpected}, | |
| read_columns=['age'], | |
| ) | |
| self.assertEqual([], messages) | |
| def test_callable_assignment_rejects_python_list(self): | |
| table_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('values', pa.list_(pa.int32())), | |
| ]) | |
| table = self._create_table(pa_schema=table_schema) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2], | |
| 'values': [[0], [0]], | |
| }, schema=table_schema)) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaisesRegex( | |
| ValueError, "must return a pyarrow.Array"): | |
| self._do_update_by_predicate( | |
| table, | |
| pb.greater_or_equal('id', 1), | |
| {'values': lambda _rows: [1, 2]}, | |
| read_columns=['id'], | |
| ) | |
| def test_update_by_predicate_rejects_assignment_array_length_mismatch(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_update_by_predicate( | |
| table, | |
| pb.greater_or_equal('age', 35), | |
| {'age': pa.array([1, 2], type=pa.int32())}, | |
| ) | |
| self.assertIn('Assignment array length', str(ctx.exception)) | |
| self.assertEqual( | |
| [25, 30, 35, 40, 45], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_rejects_uncastable_assignment(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises((pa.ArrowInvalid, pa.ArrowTypeError, ValueError)): | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {'age': 'not-an-int'}, | |
| ) | |
| self.assertEqual( | |
| [25, 30, 35, 40, 45], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_by_predicate_with_global_index_updates_unindexed_rows(self): | |
| table = self._create_global_indexed_table_for_predicate_update({ | |
| 'scalar-index.search-mode': 'fast', | |
| }) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('name', 'new'), | |
| {'age': 21}, | |
| ) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 10, 2: 15, 3: 21, 4: 30}, ages_by_id) | |
| def test_delete_by_predicate_with_global_index_deletes_unindexed_rows(self): | |
| table = self._create_global_indexed_table_for_predicate_update({ | |
| 'deletion-vectors.enabled': 'true', | |
| 'scalar-index.search-mode': 'fast', | |
| }) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self._do_delete_by_predicate(table, pb.equal('name', 'new')) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 2, 4], result['id'].to_pylist()) | |
| def test_update_by_predicate_with_global_index_falls_back_to_full_scan(self): | |
| table = self._create_global_indexed_table_for_predicate_update() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('city', 'LA'), | |
| {'age': 88}, | |
| ) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 10, 2: 88, 3: 88, 4: 30}, ages_by_id) | |
| def test_update_by_predicate_with_global_index_handles_compound_predicate(self): | |
| table = self._create_global_indexed_table_for_predicate_update() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| predicate = pb.or_predicates([ | |
| pb.equal('name', 'old'), | |
| pb.equal('city', 'SF'), | |
| ]) | |
| self._do_update_by_predicate(table, predicate, {'age': 77}) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 77, 2: 15, 3: 20, 4: 77}, ages_by_id) | |
| def test_update_by_predicate_resolves_time_travel_scan_snapshot(self): | |
| table = self._create_table() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1], | |
| 'name': ['old'], | |
| 'age': [10], | |
| 'city': ['NYC'], | |
| }, schema=self.pa_schema)) | |
| table.create_tag('before_new') | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [2], | |
| 'name': ['new'], | |
| 'age': [20], | |
| 'city': ['LA'], | |
| }, schema=self.pa_schema)) | |
| travel_table = table.copy({'scan.tag-name': 'before_new'}) | |
| pb = travel_table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_update_by_predicate( | |
| travel_table, | |
| pb.equal('name', 'new'), | |
| {'age': 21}, | |
| ) | |
| self.assertEqual([], msgs) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 10, 2: 20}, ages_by_id) | |
| def test_update_by_predicate_resolves_scan_snapshot_id(self): | |
| table = self._create_table() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1], | |
| 'name': ['old'], | |
| 'age': [10], | |
| 'city': ['NYC'], | |
| }, schema=self.pa_schema)) | |
| snapshot = table.snapshot_manager().get_latest_snapshot() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [2], | |
| 'name': ['new'], | |
| 'age': [20], | |
| 'city': ['LA'], | |
| }, schema=self.pa_schema)) | |
| travel_table = table.copy({'scan.snapshot-id': str(snapshot.id)}) | |
| pb = travel_table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_update_by_predicate( | |
| travel_table, | |
| pb.equal('name', 'new'), | |
| {'age': 21}, | |
| ) | |
| self.assertEqual([], msgs) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 10, 2: 20}, ages_by_id) | |
| def test_update_by_predicate_resets_explicit_scan_mode_after_travel(self): | |
| table = self._create_table() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1], | |
| 'name': ['old'], | |
| 'age': [10], | |
| 'city': ['NYC'], | |
| }, schema=self.pa_schema)) | |
| snapshot = table.snapshot_manager().get_latest_snapshot() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [2], | |
| 'name': ['new'], | |
| 'age': [20], | |
| 'city': ['LA'], | |
| }, schema=self.pa_schema)) | |
| travel_table = table.copy({ | |
| 'scan.mode': 'from-timestamp', | |
| 'scan.timestamp-millis': str(snapshot.time_millis), | |
| }) | |
| pb = travel_table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_update_by_predicate( | |
| travel_table, | |
| pb.equal('name', 'new'), | |
| {'age': 21}, | |
| ) | |
| self.assertEqual([], msgs) | |
| result = self._read_all(table) | |
| ages_by_id = dict(zip( | |
| result['id'].to_pylist(), | |
| result['age'].to_pylist(), | |
| )) | |
| self.assertEqual({1: 10, 2: 20}, ages_by_id) | |
| def test_update_by_predicate_rejects_empty_assignments(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {}, | |
| ) | |
| self.assertIn('assignments must not be empty', str(ctx.exception)) | |
| def test_update_by_predicate_rejects_unknown_column(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {'unknown': 1}, | |
| ) | |
| self.assertIn('Column unknown is not in table schema', | |
| str(ctx.exception)) | |
| def test_update_by_predicate_rejects_partition_column(self): | |
| table = self._create_table(partition_keys=['city']) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1], | |
| 'name': ['Alice'], | |
| 'age': [25], | |
| 'city': ['NYC'], | |
| }, schema=self.pa_schema)) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_update_by_predicate( | |
| table, | |
| pb.equal('id', 1), | |
| {'city': 'LA'}, | |
| ) | |
| self.assertIn('partition column', str(ctx.exception)) | |
| def test_delete_by_predicate(self): | |
| table = self._create_seeded_deletion_vector_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_delete_by_predicate( | |
| table, | |
| pb.greater_or_equal('age', 35), | |
| ) | |
| self.assertEqual(1, sum(len(m.index_adds) for m in msgs)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 2], result['id'].to_pylist()) | |
| self.assertEqual(['Alice', 'Bob'], result['name'].to_pylist()) | |
| def test_delete_by_predicate_requires_deletion_vectors(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_delete_by_predicate(table, pb.equal('id', 1)) | |
| self.assertIn('deletion-vectors.enabled', str(ctx.exception)) | |
| def test_delete_by_row_id(self): | |
| table = self._create_seeded_deletion_vector_table() | |
| msgs = self._do_delete_by_row_id(table, [0, 2, 4]) | |
| self.assertEqual(1, sum(len(m.index_adds) for m in msgs)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([2, 4], result['id'].to_pylist()) | |
| self.assertEqual(['Bob', 'David'], result['name'].to_pylist()) | |
| def test_delete_by_row_id_requires_deletion_vectors(self): | |
| table = self._create_seeded_table() | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_delete_by_row_id(table, [0]) | |
| self.assertIn('deletion-vectors.enabled', str(ctx.exception)) | |
| def test_concurrent_row_level_deletes_conflict_on_same_dv_file(self): | |
| table = self._create_seeded_deletion_vector_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| first_wb = self._make_write_builder(table) | |
| first_update = first_wb.new_update() | |
| first_cid = self._next_commit_id() | |
| first_msgs = self._apply_delete_by_predicate( | |
| first_update, | |
| pb.equal('id', 1), | |
| first_cid, | |
| ) | |
| self._do_delete_by_predicate(table, pb.equal('id', 2)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 3, 4, 5], result['id'].to_pylist()) | |
| first_commit = first_wb.new_commit() | |
| with self.assertRaises(RuntimeError) as ctx: | |
| self._apply_commit(first_commit, first_msgs, first_cid) | |
| first_commit.close() | |
| self.assertIn('Deletion vector index conflict', str(ctx.exception)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 3, 4, 5], result['id'].to_pylist()) | |
| def test_row_level_delete_conflicts_when_target_file_removed(self): | |
| table = self._create_seeded_deletion_vector_table(partition_keys=['city']) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| first_wb = self._make_write_builder(table) | |
| first_update = first_wb.new_update() | |
| first_cid = self._next_commit_id() | |
| first_msgs = self._apply_delete_by_predicate( | |
| first_update, | |
| pb.equal('id', 2), | |
| first_cid, | |
| ) | |
| overwrite_wb = table.new_batch_write_builder().overwrite({'city': 'LA'}) | |
| overwrite_write = overwrite_wb.new_write() | |
| overwrite_commit = overwrite_wb.new_commit() | |
| overwrite_write.write_arrow(pa.Table.from_pydict({ | |
| 'id': [6], | |
| 'name': ['Frank'], | |
| 'age': [28], | |
| 'city': ['LA'], | |
| }, schema=self.pa_schema)) | |
| overwrite_commit.commit(overwrite_write.prepare_commit()) | |
| overwrite_write.close() | |
| overwrite_commit.close() | |
| first_commit = first_wb.new_commit() | |
| with self.assertRaises(RuntimeError) as ctx: | |
| self._apply_commit(first_commit, first_msgs, first_cid) | |
| first_commit.close() | |
| self.assertIn('Deletion vector index conflict', str(ctx.exception)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 3, 4, 5, 6], result['id'].to_pylist()) | |
| def test_delete_by_partition_predicate_drops_partition_without_dv(self): | |
| table = self._create_seeded_table(partition_keys=['city']) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| msgs = self._do_delete_by_predicate(table, pb.equal('city', 'LA')) | |
| self.assertGreater(sum(len(m.deleted_files) for m in msgs), 0) | |
| self.assertEqual(0, sum(len(m.index_adds) for m in msgs)) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 3, 4, 5], result['id'].to_pylist()) | |
| self.assertEqual( | |
| ['NYC', 'Chicago', 'Houston', 'Phoenix'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_delete_by_partition_predicate_deletes_matching_indexes(self): | |
| table = self._create_seeded_deletion_vector_table(partition_keys=['city']) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| self._do_delete_by_predicate(table, pb.equal('id', 2)) | |
| msgs = self._do_delete_by_predicate(table, pb.equal('city', 'LA')) | |
| self.assertGreater(sum(len(m.deleted_files) for m in msgs), 0) | |
| self.assertGreater(sum(len(m.index_deletes) for m in msgs), 0) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([1, 3, 4, 5], result['id'].to_pylist()) | |
| self.assertEqual( | |
| ['NYC', 'Chicago', 'Houston', 'Phoenix'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_delete_by_partition_predicate_conflicts_when_partition_changes(self): | |
| table = self._create_seeded_table(partition_keys=['city']) | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| pb = tu.new_predicate_builder() | |
| cid = self._next_commit_id() | |
| msgs = self._apply_delete_by_predicate(tu, pb.equal('city', 'LA'), cid) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [6], | |
| 'name': ['Frank'], | |
| 'age': [28], | |
| 'city': ['LA'], | |
| }, schema=self.pa_schema)) | |
| tc = wb.new_commit() | |
| with self.assertRaises(RuntimeError) as ctx: | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| self.assertIn('Overwrite conflict', str(ctx.exception)) | |
| def test_delete_by_mixed_partition_predicate_still_requires_dv(self): | |
| table = self._create_seeded_table(partition_keys=['city']) | |
| pb = table.new_read_builder().new_predicate_builder() | |
| predicate = pb.and_predicates([ | |
| pb.equal('city', 'LA'), | |
| pb.equal('age', 30), | |
| ]) | |
| with self.assertRaises(ValueError) as ctx: | |
| self._do_delete_by_predicate(table, predicate) | |
| self.assertIn('deletion-vectors.enabled', str(ctx.exception)) | |
| def test_update_multiple_columns(self): | |
| """Update ``age`` + ``city`` together; other columns are untouched.""" | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 0, 2, 3, 4], | |
| 'age': [31, 26, 36, 39, 42], | |
| 'city': ['Los Angeles', 'New York', 'Chicago', 'Phoenix', 'Houston'], | |
| }), ['age', 'city']) | |
| result = self._read_all(table) | |
| self.assertEqual([26, 31, 36, 39, 42], result['age'].to_pylist()) | |
| self.assertEqual( | |
| ['New York', 'Los Angeles', 'Chicago', 'Phoenix', 'Houston'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_update_columns_fall_back_to_data_when_unset(self): | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 1, 2, 3, 4], | |
| 'id': [1, 2, 3, 4, 5], | |
| 'name': ['A', 'B', 'C', 'D', 'E'], | |
| 'age': [1, 2, 3, 4, 5], | |
| 'city': ['c0', 'c1', 'c2', 'c3', 'c4'], | |
| }), ['id', 'name', 'age', 'city']) | |
| result = self._read_all(table) | |
| self.assertEqual(['A', 'B', 'C', 'D', 'E'], result['name'].to_pylist()) | |
| self.assertEqual([1, 2, 3, 4, 5], result['age'].to_pylist()) | |
| self.assertEqual(['c0', 'c1', 'c2', 'c3', 'c4'], result['city'].to_pylist()) | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| cid = self._next_commit_id() | |
| msgs = self._apply_update(tu, pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 1], | |
| 'age': [99, 98], | |
| }), cid) | |
| tc = wb.new_commit() | |
| self._apply_commit(tc, msgs, cid) | |
| tc.close() | |
| result = self._read_all(table) | |
| self.assertEqual([99, 98, 3, 4, 5], result['age'].to_pylist()) | |
| self.assertEqual(['A', 'B', 'C', 'D', 'E'], result['name'].to_pylist()) | |
| def test_update_with_only_row_id_raises(self): | |
| table = self._create_seeded_table() | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update() | |
| cid = self._next_commit_id() | |
| with self.assertRaises(ValueError): | |
| self._apply_update(tu, pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 1], | |
| }), cid) | |
| @pytest.mark.python_plan | |
| def test_partitioned_table_update(self): | |
| """Updates work on a partitioned table the same as a flat one.""" | |
| table = self._create_table(partition_keys=['city']) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2, 3, 4, 5], | |
| 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| 'age': [25, 30, 35, 40, 45], | |
| 'city': ['NYC', 'NYC', 'LA', 'LA', 'Chicago'], | |
| }, schema=self.pa_schema)) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 0, 2, 3, 4], | |
| 'age': [31, 26, 36, 41, 46], | |
| }), ['age']) | |
| self.assertEqual( | |
| [26, 31, 36, 41, 46], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_preserves_other_columns(self): | |
| """Updating ``age`` leaves all other columns untouched.""" | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1], | |
| 'age': [999], | |
| }), ['age']) | |
| result = self._read_all(table) | |
| self.assertEqual([25, 999, 35, 40, 45], result['age'].to_pylist()) | |
| self.assertEqual( | |
| ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| result['name'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_partial_row_updates(self): | |
| """All partial-row patterns share one parameterised test.""" | |
| cases = [ | |
| # name, row_ids, ages, expected ages after update | |
| ('single_first_file', [0], [100], [100, 30, 35, 40, 45]), | |
| ('single_second_file', [1], [31], [25, 31, 35, 40, 45]), | |
| ('one_per_file', [1, 2], [200, 300], [25, 200, 300, 40, 45]), | |
| ('non_consecutive', [0, 2, 4], [100, 300, 500], [100, 30, 300, 40, 500]), | |
| ] | |
| for name, row_ids, ages, expected in cases: | |
| with self.subTest(case=name): | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': row_ids, | |
| 'age': ages, | |
| }), ['age']) | |
| self.assertEqual( | |
| expected, | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_multiple_sequential_single_column_updates(self): | |
| """Two sequential updates on different single columns compose.""" | |
| table = self._create_seeded_table() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 0, 2, 3, 4], | |
| 'age': [31, 26, 36, 41, 46], | |
| }), ['age']) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 0, 2, 3, 4], | |
| 'city': ['Los Angeles', 'New York', 'Chicago', 'Phoenix', 'Houston'], | |
| }), ['city']) | |
| result = self._read_all(table) | |
| self.assertEqual([26, 31, 36, 41, 46], result['age'].to_pylist()) | |
| self.assertEqual( | |
| ['New York', 'Los Angeles', 'Chicago', 'Phoenix', 'Houston'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_sequential_partial_updates(self): | |
| """Sequential single-row updates accumulate correctly.""" | |
| table = self._create_seeded_table() | |
| for row_id, age in [(0, 100), (2, 300), (4, 500)]: | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [row_id], | |
| 'age': [age], | |
| }), ['age']) | |
| self.assertEqual( | |
| [100, 30, 300, 40, 500], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_partial_rows_across_two_files(self): | |
| """A single update modifies partial rows from both data files.""" | |
| table = self._create_table() | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2, 3, 4, 5], | |
| 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| 'age': [20, 25, 30, 35, 40], | |
| 'city': ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'], | |
| }, schema=self.pa_schema)) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [6, 7, 8, 9, 10], | |
| 'name': ['Frank', 'Grace', 'Henry', 'Ivy', 'Jack'], | |
| 'age': [45, 50, 55, 60, 65], | |
| 'city': ['Seattle', 'Boston', 'Denver', 'Miami', 'Atlanta'], | |
| }, schema=self.pa_schema)) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [1, 3, 6, 8], | |
| 'age': [100, 200, 300, 400], | |
| 'name': ['Updated_Bob', 'Updated_David', 'Updated_Grace', 'Updated_Ivy'], | |
| }), ['age', 'name']) | |
| result = self._read_all(table) | |
| self.assertEqual( | |
| [20, 100, 30, 200, 40, 45, 300, 55, 400, 65], | |
| result['age'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['Alice', 'Updated_Bob', 'Charlie', 'Updated_David', 'Eve', | |
| 'Frank', 'Updated_Grace', 'Henry', 'Updated_Ivy', 'Jack'], | |
| result['name'].to_pylist(), | |
| ) | |
| # Untouched columns | |
| self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], | |
| result['id'].to_pylist()) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix', | |
| 'Seattle', 'Boston', 'Denver', 'Miami', 'Atlanta'], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_large_table_partial_column_updates(self): | |
| """4-step update sequence on a 2000-row / 2-file table. | |
| Covers: single-column / multi-column / both-files updates while | |
| verifying that untouched columns stay intact. | |
| """ | |
| num_row = 1000 | |
| table = self._create_table() | |
| # 2 commits, num_row rows each | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': list(range(num_row)), | |
| 'name': [f'Name_{i}' for i in range(num_row)], | |
| 'age': [20 + i for i in range(num_row)], | |
| 'city': [f'City_{i}' for i in range(num_row)], | |
| }, schema=self.pa_schema)) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': list(range(num_row, num_row * 2)), | |
| 'name': [f'Name_{i}' for i in range(num_row, num_row * 2)], | |
| 'age': [20 + num_row + i for i in range(num_row)], | |
| 'city': [f'City_{i}' for i in range(num_row, num_row * 2)], | |
| }, schema=self.pa_schema)) | |
| expected_ids = list(range(num_row * 2)) | |
| expected_names = [f'Name_{i}' for i in range(num_row * 2)] | |
| expected_ages = ([20 + i for i in range(num_row)] | |
| + [20 + num_row + i for i in range(num_row)]) | |
| # 1) update id col, 1 row | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [5], 'id': [999], | |
| }), ['id']) | |
| expected_ids[5] = 999 | |
| # 2) update id+name cols, 2 rows | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [3, 15], | |
| 'id': [888, 777], | |
| 'name': ['Updated_Name_3', 'Updated_Name_15'], | |
| }), ['id', 'name']) | |
| expected_ids[3], expected_ids[15] = 888, 777 | |
| expected_names[3] = 'Updated_Name_3' | |
| expected_names[15] = 'Updated_Name_15' | |
| # 3) update name col, 1 row | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [12], 'name': ['NewName_12'], | |
| }), ['name']) | |
| expected_names[12] = 'NewName_12' | |
| # 4) update age col, 4 rows across both files | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 5, 10, 15], | |
| 'age': [100, 105, 110, 115], | |
| }), ['age']) | |
| for r, a in [(0, 100), (5, 105), (10, 110), (15, 115)]: | |
| expected_ages[r] = a | |
| result = self._read_all(table) | |
| self.assertEqual(expected_ids, result['id'].to_pylist()) | |
| self.assertEqual(expected_names, result['name'].to_pylist()) | |
| self.assertEqual(expected_ages, result['age'].to_pylist()) | |
| # city was never touched | |
| self.assertEqual( | |
| [f'City_{i}' for i in range(num_row * 2)], | |
| result['city'].to_pylist(), | |
| ) | |
| def test_update_with_large_file(self): | |
| """Updates disable both size- and row-based rolling.""" | |
| from pypaimon.schema.schema_change import SetOption | |
| N = 5000 | |
| schema = pa.schema([('id', pa.int64()), ('name', pa.string())]) | |
| opts = { | |
| 'row-tracking.enabled': 'true', | |
| 'data-evolution.enabled': 'true', | |
| 'write-only': 'true', | |
| } | |
| table = self._create_table(pa_schema=schema, options=opts) | |
| table_identifier = table.identifier | |
| self._write_arrow(table, pa.table({ | |
| 'id': list(range(N)), | |
| 'name': [''.join(random.choices(string.ascii_letters, k=200)) | |
| for _ in range(N)], | |
| })) | |
| self.catalog.alter_table( | |
| table_identifier, [ | |
| SetOption('target-file-size', '10kb'), | |
| SetOption('target-file-row-num', '1'), | |
| ] | |
| ) | |
| table = self.catalog.get_table(table_identifier) | |
| msgs = self._do_update(table, pa.table({ | |
| '_ROW_ID': pa.array(list(range(N)), type=pa.int64()), | |
| 'name': [''.join(random.choices(string.ascii_letters, k=200)) | |
| for _ in range(N)], | |
| }), ['name']) | |
| all_files = [f for m in msgs for f in m.new_files] | |
| self.assertEqual(1, len(all_files), | |
| "Update should produce exactly one file per group") | |
| self.assertEqual(0, all_files[0].first_row_id) | |
| self.assertEqual(N, all_files[0].row_count) | |
| # ------------------------------------------------------------------ | |
| # Validation tests | |
| # ------------------------------------------------------------------ | |
| def test_nonexistent_column_raises(self): | |
| table = self._create_seeded_table() | |
| bad = pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 1, 2, 3, 4], | |
| 'nonexistent_column': [100, 200, 300, 400, 500], | |
| }) | |
| with self.assertRaises(ValueError) as ctx: | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(['nonexistent_column']) | |
| self._apply_update(tu, bad, self._next_commit_id()) | |
| self.assertIn('not in table schema', str(ctx.exception)) | |
| def test_missing_row_id_column_raises(self): | |
| table = self._create_seeded_table() | |
| bad = pa.Table.from_pydict({'age': [26, 27, 28, 29, 30]}) | |
| with self.assertRaises(ValueError) as ctx: | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(['age']) | |
| self._apply_update(tu, bad, self._next_commit_id()) | |
| self.assertIn('_ROW_ID column', str(ctx.exception)) | |
| def test_invalid_row_id_raises(self): | |
| """row_id outside valid row_id ranges raises.""" | |
| table = self._create_seeded_table() | |
| cases = [ | |
| ('out_of_range_high', [0, 10], [26, 100]), | |
| ('negative', [-1, 0], [100, 26]), | |
| ] | |
| for name, row_ids, ages in cases: | |
| with self.subTest(case=name): | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(['age']) | |
| bad = pa.Table.from_pydict({'_ROW_ID': row_ids, 'age': ages}) | |
| with self.assertRaises(ValueError) as ctx: | |
| self._apply_update(tu, bad, self._next_commit_id()) | |
| self.assertIn('does not belong to any valid range', str(ctx.exception)) | |
| def test_duplicate_row_id_raises(self): | |
| table = self._create_seeded_table() | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(['age']) | |
| with self.assertRaises(ValueError) as ctx: | |
| self._apply_update( | |
| tu, | |
| pa.Table.from_pydict({ | |
| '_ROW_ID': [0, 0, 1], | |
| 'age': [100, 200, 300], | |
| }), | |
| self._next_commit_id(), | |
| ) | |
| self.assertIn('duplicate _ROW_ID', str(ctx.exception)) | |
| def test_update_deleted_row_id_raises(self): | |
| """Updating a row_id that fell into a hole after truncate raises.""" | |
| partitioned_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('name', pa.string()), | |
| ('age', pa.int32()), | |
| ('region', pa.string()), | |
| ]) | |
| table = self._create_table( | |
| pa_schema=partitioned_schema, | |
| partition_keys=['region'], | |
| ) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': pa.array([1, 2, 3], type=pa.int32()), | |
| 'name': ['A', 'B', 'C'], | |
| 'age': pa.array([10, 20, 30], type=pa.int32()), | |
| 'region': ['US', 'US', 'US'], | |
| }, schema=partitioned_schema)) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': pa.array([4, 5], type=pa.int32()), | |
| 'name': ['D', 'E'], | |
| 'age': pa.array([40, 50], type=pa.int32()), | |
| 'region': ['EU', 'EU'], | |
| }, schema=partitioned_schema)) | |
| wb = table.new_batch_write_builder() | |
| tc = wb.new_commit() | |
| tc.truncate_partitions([{'region': 'US'}]) | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(['age']) | |
| with self.assertRaises(ValueError) as ctx: | |
| self._apply_update( | |
| tu, | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [99]}), | |
| self._next_commit_id(), | |
| ) | |
| self.assertIn('does not belong to any valid range', str(ctx.exception)) | |
| # ------------------------------------------------------------------ | |
| # Concurrency tests | |
| # ------------------------------------------------------------------ | |
| def _run_concurrent_updates(self, table, thread_specs, max_retries): | |
| """Run a batch of concurrent updates with conflict-retry; return the | |
| commit order (``thread_index`` of the winning commit appended last).""" | |
| errors = [] | |
| completion_order = [] | |
| lock = threading.Lock() | |
| retry_marker = ( | |
| "multiple 'MERGE INTO' operations have encountered conflicts" | |
| ) | |
| def worker(idx, spec): | |
| # Tag each thread's commits so the durable winner can be read back from the | |
| # latest snapshot's commit_user (the order threads return in is not the commit order). | |
| worker_table = table.copy({'commit.user-prefix': 'w%d' % idx}) | |
| for _ in range(max_retries): | |
| try: | |
| self._do_update(worker_table, pa.Table.from_pydict({ | |
| '_ROW_ID': spec['row_ids'], | |
| 'age': spec['ages'], | |
| }), ['age']) | |
| with lock: | |
| completion_order.append(idx) | |
| return | |
| except Exception as e: | |
| if retry_marker not in str(e): | |
| errors.append(f"Thread-{idx} unexpected error: {e}") | |
| return | |
| errors.append(f"Thread-{idx} did not succeed in {max_retries} retries") | |
| threads = [threading.Thread(target=worker, args=(i, s)) | |
| for i, s in enumerate(thread_specs)] | |
| for t in threads: | |
| t.start() | |
| for t in threads: | |
| t.join(timeout=120) | |
| self.assertEqual([], errors) | |
| self.assertEqual(len(thread_specs), len(completion_order), | |
| "Not all threads committed successfully") | |
| return completion_order | |
| def test_concurrent_updates_disjoint_rows(self): | |
| """Disjoint per-thread updates all materialise into the final state.""" | |
| table = self._create_seeded_table() | |
| specs = [ | |
| {'row_ids': [0, 1], 'ages': [100, 200]}, | |
| {'row_ids': [2, 3], 'ages': [300, 400]}, | |
| {'row_ids': [4], 'ages': [500]}, | |
| ] | |
| self._run_concurrent_updates(table, specs, max_retries=20) | |
| self.assertEqual( | |
| [100, 200, 300, 400, 500], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_concurrent_updates_overlapping_rows_last_writer_wins(self): | |
| """When threads compete on the same rows, the final commit wins.""" | |
| table = self._create_seeded_table() | |
| specs = [ | |
| {'row_ids': [0, 1, 2], 'ages': [101, 201, 301]}, | |
| {'row_ids': [0, 1, 2], 'ages': [102, 202, 302]}, | |
| {'row_ids': [0, 1, 2], 'ages': [103, 203, 303]}, | |
| ] | |
| self._run_concurrent_updates(table, specs, max_retries=30) | |
| ages = self._read_all(table)['age'].to_pylist() | |
| # The real winner is the thread whose commit produced the latest snapshot. | |
| winner = int(table.snapshot_manager().get_latest_snapshot().commit_user[1:].split('_')[0]) | |
| self.assertEqual(specs[winner]['ages'], ages[:3]) | |
| self.assertEqual([40, 45], ages[3:]) | |
| def test_update_list_and_map_columns(self): | |
| list_map_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('tags', pa.list_(pa.string())), | |
| ('meta', pa.map_(pa.string(), pa.string())), | |
| ]) | |
| table = self._create_table(pa_schema=list_map_schema) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2, 3], | |
| 'tags': [['a', 'b'], ['c'], ['d', 'e']], | |
| 'meta': [[('k1', 'v1')], [('k2', 'v2')], [('k3', 'v3')]], | |
| }, schema=list_map_schema)) | |
| rb = table.new_read_builder().with_projection( | |
| ['id', '_ROW_ID']) | |
| rid_result = rb.new_read().to_arrow( | |
| rb.new_scan().plan().splits()).sort_by('id') | |
| row_ids = rid_result['_ROW_ID'].to_pylist() | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[0], row_ids[2]], | |
| type=pa.int64()), | |
| 'tags': [['x', 'y'], ['z']], | |
| 'meta': [[('k1', 'new1')], [('k3', 'new3')]], | |
| }), ['tags', 'meta']) | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual( | |
| [['x', 'y'], ['c'], ['z']], | |
| result['tags'].to_pylist()) | |
| self.assertEqual( | |
| [[('k1', 'new1')], [('k2', 'v2')], [('k3', 'new3')]], | |
| result['meta'].to_pylist()) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[1]], type=pa.int64()), | |
| 'meta': [{'k2': 'dict_val'}], | |
| }), ['meta']) | |
| result2 = self._read_all(table).sort_by('id') | |
| self.assertEqual( | |
| [[('k1', 'new1')], [('k2', 'dict_val')], [('k3', 'new3')]], | |
| result2['meta'].to_pylist()) | |
| with self.assertRaisesRegex(ValueError, "schema-less dict"): | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[0], row_ids[2]], | |
| type=pa.int64()), | |
| 'meta': [{'a': '1'}, {'b': '2'}], | |
| }), ['meta']) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[0], row_ids[2]], | |
| type=pa.int64()), | |
| 'meta': pa.array( | |
| [[('a', '1')], [('b', '2')]], | |
| type=pa.map_(pa.string(), pa.string())), | |
| }), ['meta']) | |
| result3 = self._read_all(table).sort_by('id') | |
| self.assertEqual( | |
| [[('a', '1')], [('k2', 'dict_val')], [('b', '2')]], | |
| result3['meta'].to_pylist()) | |
| with self.assertRaisesRegex(ValueError, "schema-less dict"): | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[0]], type=pa.int64()), | |
| 'meta': [{'a': None}], | |
| }), ['meta']) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([row_ids[0]], type=pa.int64()), | |
| 'meta': [[('a', None)]], | |
| }, schema=pa.schema([ | |
| ('_ROW_ID', pa.int64()), | |
| ('meta', pa.map_(pa.string(), pa.string())), | |
| ])), ['meta']) | |
| result4 = self._read_all(table).sort_by('id') | |
| self.assertEqual( | |
| [[('a', None)], [('k2', 'dict_val')], [('b', '2')]], | |
| result4['meta'].to_pylist()) | |
| @pytest.mark.python_write | |
| def test_update_by_row_id_preserves_files_after_prepare_commit_failure(self): | |
| table_schema = pa.schema([ | |
| ('id', pa.int32()), | |
| ('age', pa.int32()), | |
| ('picture', pa.large_binary()), | |
| ]) | |
| table = self._create_table(pa_schema=table_schema) | |
| self._write_arrow(table, pa.Table.from_pydict({ | |
| 'id': [1, 2], | |
| 'age': [10, 20], | |
| 'picture': [b'blob-1', b'blob-2'], | |
| }, schema=table_schema)) | |
| rb = table.new_read_builder().with_projection(['id', '_ROW_ID']) | |
| row_ids = rb.new_read().to_arrow( | |
| rb.new_scan().plan().splits()).sort_by('id')['_ROW_ID'] | |
| before_files = self._list_table_files(table) | |
| prepared_files = [] | |
| def fail_after_prepare_commit( | |
| new_files, first_row_id, column_names, blob_columns): | |
| prepared_files.extend(new_files) | |
| raise RuntimeError("forced failure after prepare_commit") | |
| for columns in (['age'], ['picture'], ['age', 'picture']): | |
| with self.subTest(columns=columns): | |
| wb = self._make_write_builder(table) | |
| tu = wb.new_update().with_update_type(columns) | |
| with mock.patch.object( | |
| TableUpdateByRowId, | |
| '_assign_update_file_metadata', | |
| new=staticmethod(fail_after_prepare_commit)): | |
| with self.assertRaisesRegex( | |
| RuntimeError, "forced failure after prepare_commit"): | |
| self._apply_update(tu, pa.Table.from_pydict({ | |
| '_ROW_ID': [row_ids[0].as_py()], | |
| 'age': [99], | |
| 'picture': [b'updated-blob'], | |
| }, schema=pa.schema([ | |
| ('_ROW_ID', pa.int64()), | |
| ('age', pa.int32()), | |
| ('picture', pa.large_binary()), | |
| ])).select(['_ROW_ID'] + columns), self._next_commit_id()) | |
| self.assertLess(before_files, self._list_table_files(table)) | |
| self.assertTrue(prepared_files) | |
| self.assertTrue(all(table.file_io.exists(file.physical_path()) | |
| for file in prepared_files)) | |
| self.assertEqual([10, 20], self._read_all(table)['age'].to_pylist()) | |
| self.assertEqual( | |
| [b'blob-1', b'blob-2'], self._read_all(table)['picture'].to_pylist()) | |
| # ====================================================================== | |
| # Mode-specific mixins (add the ``update_by_arrow_with_row_id`` primitive) | |
| # ====================================================================== | |
| class _BatchModeMixin(BatchModeMixin): | |
| def _apply_update(self, table_update, data, cid): | |
| return table_update.update_by_arrow_with_row_id(data) | |
| def _apply_update_by_predicate( | |
| self, table_update, predicate, assignments, cid, | |
| read_columns=None): | |
| return table_update.update_by_predicate( | |
| predicate, assignments, read_columns | |
| ) | |
| def _apply_delete_by_predicate(self, table_update, predicate, cid): | |
| return table_update.delete_by_predicate(predicate) | |
| def _apply_delete_by_row_id(self, table_update, row_ids, cid): | |
| return table_update.delete_by_row_id(row_ids) | |
| class _StreamModeMixin(StreamModeMixin): | |
| def _apply_update(self, table_update, data, cid): | |
| return table_update.update_by_arrow_with_row_id(data, cid) | |
| def _apply_update_by_predicate( | |
| self, table_update, predicate, assignments, cid, | |
| read_columns=None): | |
| return table_update.update_by_predicate( | |
| predicate, | |
| assignments, | |
| cid, | |
| read_columns, | |
| ) | |
| def _apply_delete_by_predicate(self, table_update, predicate, cid): | |
| return table_update.delete_by_predicate(predicate, cid) | |
| def _apply_delete_by_row_id(self, table_update, row_ids, cid): | |
| return table_update.delete_by_row_id(row_ids, cid) | |
| # ====================================================================== | |
| # Concrete test classes | |
| # ====================================================================== | |
| class TableUpdateBatchTest(_BatchModeMixin, _TableUpdateTestBase, unittest.TestCase): | |
| """All shared update tests under batch (``BatchWriteBuilder``) semantics.""" | |
| @pytest.mark.python_write | |
| def test_update_batches_reject_same_file_and_preserve_staged_files(self): | |
| table = self._create_seeded_table() | |
| before_files = self._list_table_files(table) | |
| update = ( | |
| self._make_write_builder(table) | |
| .new_update() | |
| .with_update_type(['age']) | |
| ) | |
| with self.assertRaisesRegex( | |
| ValueError, "overlapping first_row_ids.*0"): | |
| update.update_by_arrow_batches_with_row_id(iter([ | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}), | |
| pa.Table.from_pydict({'_ROW_ID': [1], 'age': [31]}), | |
| ])) | |
| self.assertLess(before_files, self._list_table_files(table)) | |
| self.assertEqual( | |
| [25, 30, 35, 40, 45], | |
| self._read_all(table)['age'].to_pylist(), | |
| ) | |
| def test_update_batches_allow_same_file_for_different_columns(self): | |
| table = self._create_seeded_table() | |
| builder = self._make_write_builder(table) | |
| messages = ( | |
| builder.new_update() | |
| .update_by_arrow_batches_with_row_id(iter([ | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [26]}), | |
| pa.Table.from_pydict({'_ROW_ID': [1], 'city': ['Seattle']}), | |
| ])) | |
| ) | |
| commit = builder.new_commit() | |
| commit.commit(messages) | |
| commit.close() | |
| rows = self._read_all(table) | |
| self.assertEqual([26, 30, 35, 40, 45], rows['age'].to_pylist()) | |
| self.assertEqual( | |
| ['NYC', 'Seattle', 'Chicago', 'Houston', 'Phoenix'], | |
| rows['city'].to_pylist(), | |
| ) | |
| def test_callable_output_preserves_large_offset_chunks(self): | |
| from pypaimon.write.table_update import TableUpdate | |
| from pypaimon.write.table_update_by_row_id import TableUpdateByRowId | |
| chunks = [ | |
| pa.ListArray.from_arrays( | |
| pa.array([0, 1_100_000_000], type=pa.int32()), | |
| pa.nulls(1_100_000_000), | |
| ) | |
| for _ in range(2) | |
| ] | |
| value = pa.chunked_array(chunks) | |
| result = TableUpdate._assignment_to_array(value, value.type, 2) | |
| self.assertEqual(2, result.num_chunks) | |
| self.assertEqual(2_200_000_000, sum( | |
| len(chunk.values) for chunk in result.chunks | |
| )) | |
| updater = TableUpdateByRowId.__new__(TableUpdateByRowId) | |
| updater.table = mock.Mock(field_names=['payload']) | |
| updater.commit_messages = [] | |
| updater._updated_first_row_ids_by_column = {} | |
| updates = pa.Table.from_arrays( | |
| [pa.array([1, 0], type=pa.int64()), result], | |
| names=['_ROW_ID', 'payload'], | |
| ) | |
| with mock.patch.object( | |
| updater, '_calculate_first_row_id', | |
| side_effect=lambda data: data.append_column( | |
| TableUpdateByRowId.FIRST_ROW_ID_COLUMN, | |
| pa.array([0, 0], type=pa.int64()), | |
| )) as calculate: | |
| with mock.patch.object(updater, '_write_by_first_row_id'): | |
| updater.update_columns(updates, ['payload']) | |
| routed = calculate.call_args[0][0] | |
| self.assertEqual([1, 0], routed['_ROW_ID'].to_pylist()) | |
| self.assertEqual(2, routed['payload'].num_chunks) | |
| def test_callable_predicate_update_allows_concurrent_read_column_update(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| wb = self._make_write_builder(table) | |
| messages = wb.new_update().update_by_predicate( | |
| pb.equal('id', 1), | |
| { | |
| 'name': lambda rows: pa.array([ | |
| 'age-%d' % value for value in rows['age'].to_pylist() | |
| ]), | |
| }, | |
| read_columns=['age'], | |
| ) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([0], type=pa.int64()), | |
| 'age': pa.array([100], type=pa.int32()), | |
| }), ['age']) | |
| commit = wb.new_commit() | |
| commit.commit(messages) | |
| commit.close() | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual(100, result['age'][0].as_py()) | |
| self.assertEqual('age-25', result['name'][0].as_py()) | |
| def test_callable_predicate_update_conflicts_with_concurrent_update(self): | |
| table = self._create_seeded_table() | |
| pb = table.new_read_builder().new_predicate_builder() | |
| wb = self._make_write_builder(table) | |
| update = wb.new_update() | |
| def increment_age(rows): | |
| self.assertEqual([25], rows['age'].to_pylist()) | |
| self._do_update(table, pa.Table.from_pydict({ | |
| '_ROW_ID': pa.array([0], type=pa.int64()), | |
| 'age': pa.array([100], type=pa.int32()), | |
| }), ['age']) | |
| return pa.compute.add(rows['age'], 1) | |
| messages = update.update_by_predicate( | |
| pb.equal('id', 1), | |
| {'age': increment_age}, | |
| read_columns=['age'], | |
| ) | |
| commit = wb.new_commit() | |
| with self.assertRaisesRegex(RuntimeError, "multiple 'MERGE INTO'"): | |
| commit.commit(messages) | |
| commit.close() | |
| result = self._read_all(table).sort_by('id') | |
| self.assertEqual([100, 30, 35, 40, 45], result['age'].to_pylist()) | |
| class TableUpdateStreamTest(_StreamModeMixin, _TableUpdateTestBase, unittest.TestCase): | |
| """All shared update tests under stream (``StreamWriteBuilder``) semantics, | |
| plus stream-only multi-commit scenarios that are impossible to express | |
| under :class:`BatchTableCommit` (which forbids re-committing). | |
| """ | |
| # ------------------------------------------------------------------ | |
| # Stream-only helpers | |
| # ------------------------------------------------------------------ | |
| def _stream_commit_age_updates_by_row_id(self, tu, tc, commit_ids, row_age_pairs): | |
| """Apply one ``age`` update per ``commit_id``; ``tu`` is already | |
| ``with_update_type(['age'])``.""" | |
| for cid, (row_id, age) in zip(commit_ids, row_age_pairs): | |
| msgs = tu.update_by_arrow_with_row_id( | |
| pa.Table.from_pydict({'_ROW_ID': [row_id], 'age': [age]}), | |
| cid, | |
| ) | |
| tc.commit(msgs, cid) | |
| # ------------------------------------------------------------------ | |
| # Stream-only tests | |
| # ------------------------------------------------------------------ | |
| def test_stream_multi_commit_on_one_write_builder(self): | |
| """A single ``StreamWriteBuilder`` drives many update+commit cycles | |
| with reused ``tu`` and ``tc`` instances. Each commit produces its own | |
| snapshot tagged with the caller-supplied ``commit_identifier`` under | |
| a stable ``commit_user`` — the core contract distinguishing stream | |
| from batch mode. | |
| Parameterised over both contiguous and sparse identifier sequences | |
| to catch any accidental coupling between ``commit_identifier`` and | |
| ``snapshot_id`` ordering. | |
| """ | |
| # (row_id, age) for each of the 3 commits — same triplet for both runs. | |
| per_commit_rows = [(0, 100), (2, 300), (4, 500)] | |
| expected_ages = [100, 30, 300, 40, 500] | |
| for case_name, commit_ids in [ | |
| ('contiguous', [1, 2, 3]), | |
| ('sparse', [42, 100, 1000]), | |
| ]: | |
| with self.subTest(case=case_name): | |
| table = self._create_seeded_table() | |
| wb, tc, base_snapshot_id = self._stream_commit_session(table) | |
| tu = wb.new_update().with_update_type(['age']) | |
| self._stream_commit_age_updates_by_row_id( | |
| tu, tc, commit_ids, per_commit_rows, | |
| ) | |
| tc.close() | |
| result = self._read_all(table) | |
| self.assertEqual(expected_ages, result['age'].to_pylist()) | |
| self.assertEqual([1, 2, 3, 4, 5], result['id'].to_pylist()) | |
| self.assertEqual( | |
| ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| result['name'].to_pylist(), | |
| ) | |
| self.assertEqual( | |
| ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'], | |
| result['city'].to_pylist(), | |
| ) | |
| self._assert_stream_builder_snapshots( | |
| table, wb, base_snapshot_id, commit_ids, | |
| ) | |
| def test_stream_commit_instance_accepts_messages_from_different_updates(self): | |
| """``StreamTableCommit`` may be committed many times, and each | |
| commit may carry messages from a *different* ``StreamTableUpdate`` | |
| instance (different ``update_cols``). The one-shot restriction | |
| that applies to :class:`BatchTableCommit` does not apply here. | |
| """ | |
| table = self._create_seeded_table() | |
| wb, tc, base_snapshot_id = self._stream_commit_session(table) | |
| # First update: only 'age' projection, fresh tu. | |
| tu1 = wb.new_update().with_update_type(['age']) | |
| msgs1 = tu1.update_by_arrow_with_row_id( | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'age': [111]}), 1, | |
| ) | |
| tc.commit(msgs1, 1) | |
| # Second update: different 'city' projection, fresh tu, same tc. | |
| tu2 = wb.new_update().with_update_type(['city']) | |
| msgs2 = tu2.update_by_arrow_with_row_id( | |
| pa.Table.from_pydict({'_ROW_ID': [0], 'city': ['Beijing']}), 2, | |
| ) | |
| tc.commit(msgs2, 2) | |
| tc.close() | |
| result = self._read_all(table) | |
| self.assertEqual([111, 30, 35, 40, 45], result['age'].to_pylist()) | |
| self.assertEqual( | |
| ['Beijing', 'LA', 'Chicago', 'Houston', 'Phoenix'], | |
| result['city'].to_pylist(), | |
| ) | |
| self.assertEqual([1, 2, 3, 4, 5], result['id'].to_pylist()) | |
| self.assertEqual( | |
| ['Alice', 'Bob', 'Charlie', 'David', 'Eve'], | |
| result['name'].to_pylist(), | |
| ) | |
| self._assert_stream_builder_snapshots( | |
| table, wb, base_snapshot_id, [1, 2], | |
| ) | |
| def test_stream_interleaved_write_and_update_on_same_builder(self): | |
| """A single stream builder may perform an initial write followed by | |
| an update on the same ``tc``, each tagged with its own | |
| ``commit_identifier`` that propagates to its own snapshot. | |
| """ | |
| table = self._create_table() | |
| wb, tc, base_snapshot_id = self._stream_commit_session(table) | |
| tw = wb.new_write() | |
| tw.write_arrow(pa.Table.from_pydict({ | |
| 'id': [1, 2, 3], | |
| 'name': ['A', 'B', 'C'], | |
| 'age': [10, 20, 30], | |
| 'city': ['X', 'Y', 'Z'], | |
| }, schema=self.pa_schema)) | |
| tc.commit(tw.prepare_commit(1), 1) | |
| tw.close() | |
| tu = wb.new_update().with_update_type(['age']) | |
| msgs = tu.update_by_arrow_with_row_id( | |
| pa.Table.from_pydict({'_ROW_ID': [0, 2], 'age': [111, 333]}), 2, | |
| ) | |
| tc.commit(msgs, 2) | |
| tc.close() | |
| result = self._read_all(table) | |
| self.assertEqual([111, 20, 333], result['age'].to_pylist()) | |
| self.assertEqual([1, 2, 3], result['id'].to_pylist()) | |
| self.assertEqual(['A', 'B', 'C'], result['name'].to_pylist()) | |
| self.assertEqual(['X', 'Y', 'Z'], result['city'].to_pylist()) | |
| self._assert_stream_builder_snapshots( | |
| table, wb, base_snapshot_id, [1, 2], | |
| ) | |
| if __name__ == '__main__': | |
| unittest.main() |