| # 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 time |
| import unittest |
| |
| import pyarrow as pa |
| |
| from pypaimon import CatalogFactory, Schema |
| from pypaimon.common.options.core_options import CoreOptions |
| |
| |
| class PkReaderTest(unittest.TestCase): |
| @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', False) |
| |
| cls.pa_schema = pa.schema([ |
| pa.field('user_id', pa.int32(), nullable=False), |
| ('item_id', pa.int64()), |
| ('behavior', pa.string()), |
| pa.field('dt', pa.string(), nullable=False) |
| ]) |
| cls.expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 3, 4, 5, 7, 8], |
| 'item_id': [1001, 1002, 1003, 1004, 1005, 1007, 1008], |
| 'behavior': ['a', 'b-new', 'c', None, 'e', 'g', 'h'], |
| 'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2'], |
| }, schema=cls.pa_schema) |
| |
| @classmethod |
| def tearDownClass(cls): |
| shutil.rmtree(cls.tempdir, ignore_errors=True) |
| |
| def test_pk_parquet_reader(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_pk_parquet', schema, False) |
| table = self.catalog.get_table('default.test_pk_parquet') |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| self.assertEqual(actual, self.expected) |
| |
| # Verify _VALUE_KIND field type is int8 in the written parquet file |
| table_scan = read_builder.new_scan() |
| splits = table_scan.plan().splits() |
| value_kind_field_found = False |
| for split in splits: |
| for file in split.files: |
| file_path = file.file_path |
| table_path = os.path.join(self.warehouse, 'default.db', 'test_pk_parquet') |
| full_path = os.path.join(table_path, file_path) |
| if os.path.exists(full_path) and file_path.endswith('.parquet'): |
| import pyarrow.parquet as pq |
| parquet_file = pq.ParquetFile(full_path) |
| # Use schema_arrow to get Arrow schema instead of ParquetSchema |
| file_schema = parquet_file.schema_arrow |
| for i in range(len(file_schema)): |
| field = file_schema.field(i) |
| if field.name == '_VALUE_KIND': |
| value_kind_field_found = True |
| self.assertEqual( |
| field.type, pa.int8(), |
| f"_VALUE_KIND field type should be int8, got {field.type}") |
| break |
| if value_kind_field_found: |
| break |
| if value_kind_field_found: |
| break |
| self.assertTrue( |
| value_kind_field_found, |
| "_VALUE_KIND field should exist in the written parquet file") |
| |
| def test_pk_orc_reader(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={ |
| 'bucket': '1', |
| 'file.format': 'orc' |
| }) |
| self.catalog.create_table('default.test_pk_orc', schema, False) |
| table = self.catalog.get_table('default.test_pk_orc') |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder() |
| actual: pa.Table = self._read_test_table(read_builder).sort_by('user_id') |
| |
| # when bucket=1, actual field name will contain 'not null', so skip comparing field name |
| for i in range(len(actual.columns)): |
| col_a = actual.column(i) |
| col_b = self.expected.column(i) |
| self.assertEqual(col_a, col_b) |
| |
| def test_pk_avro_reader(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={ |
| 'bucket': '2', |
| 'file.format': 'avro' |
| }) |
| self.catalog.create_table('default.test_pk_avro', schema, False) |
| table = self.catalog.get_table('default.test_pk_avro') |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| self.assertEqual(actual, self.expected) |
| |
| def test_pk_lance_reader(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={ |
| 'bucket': '2', |
| 'file.format': 'lance' |
| }) |
| self.catalog.create_table('default.test_pk_lance', schema, False) |
| table = self.catalog.get_table('default.test_pk_lance') |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| splits = table_scan.plan().splits() |
| |
| for split in splits: |
| for file in split.files: |
| file_path = file.file_path |
| table_path = os.path.join(self.warehouse, 'default.db', 'test_pk_lance') |
| full_path = os.path.join(table_path, file_path) |
| if os.path.exists(full_path): |
| self.assertTrue(os.path.exists(full_path)) |
| self.assertTrue( |
| file_path.endswith('.lance'), |
| f"Expected file path to end with .lance, got {file_path}") |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| self.assertEqual(actual, self.expected) |
| |
| def test_pk_lance_reader_with_filter(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={ |
| 'bucket': '2', |
| 'file.format': 'lance' |
| }) |
| self.catalog.create_table('default.test_pk_lance_filter', schema, False) |
| table = self.catalog.get_table('default.test_pk_lance_filter') |
| self._write_test_table(table) |
| |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| p1 = predicate_builder.is_in('dt', ['p1']) |
| p2 = predicate_builder.between('user_id', 2, 7) |
| p3 = predicate_builder.is_not_null('behavior') |
| g1 = predicate_builder.and_predicates([p1, p2, p3]) |
| read_builder = table.new_read_builder().with_filter(g1) |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| expected = pa.concat_tables([ |
| self.expected.slice(1, 1), # 2/b |
| self.expected.slice(5, 1) # 7/g |
| ]) |
| self.assertEqual(actual, expected) |
| |
| def test_pk_multi_write_once_commit(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_pk_multi', schema, False) |
| table = self.catalog.get_table('default.test_pk_multi') |
| write_builder = table.new_batch_write_builder() |
| |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'behavior': ['a', 'b', 'c', None], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table1 = pa.Table.from_pydict(data1, schema=self.pa_schema) |
| data2 = { |
| 'user_id': [5, 2, 7, 8], |
| 'item_id': [1005, 1002, 1007, 1008], |
| 'behavior': ['e', 'b-new', 'g', 'h'], |
| 'dt': ['p2', 'p1', 'p1', 'p2'] |
| } |
| pa_table2 = pa.Table.from_pydict(data2, schema=self.pa_schema) |
| |
| table_write.write_arrow(pa_table1) |
| table_write.write_arrow(pa_table2) |
| |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| # The in-memory merge buffer in KeyValueDataWriter folds the |
| # two writes for user_id=2 down to the latest row before flush |
| # (default merge engine is deduplicate), so the PK appears once |
| # with the second batch's value. |
| expected = pa.Table.from_pydict({ |
| 'user_id': [1, 2, 3, 4, 5, 7, 8], |
| 'item_id': [1001, 1002, 1003, 1004, 1005, 1007, 1008], |
| 'behavior': ['a', 'b-new', 'c', None, 'e', 'g', 'h'], |
| 'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2'], |
| }, schema=self.pa_schema) |
| self.assertEqual(actual, expected) |
| |
| def test_pk_reader_with_filter(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_pk_filter', schema, False) |
| table = self.catalog.get_table('default.test_pk_filter') |
| self._write_test_table(table) |
| |
| predicate_builder = table.new_read_builder().new_predicate_builder() |
| p1 = predicate_builder.is_in('dt', ['p1']) |
| p2 = predicate_builder.between('user_id', 2, 7) |
| p3 = predicate_builder.is_not_null('behavior') |
| g1 = predicate_builder.and_predicates([p1, p2, p3]) |
| read_builder = table.new_read_builder().with_filter(g1) |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| expected = pa.concat_tables([ |
| self.expected.slice(1, 1), # 2/b |
| self.expected.slice(5, 1) # 7/g |
| ]) |
| self.assertEqual(actual, expected) |
| |
| def test_pk_reader_with_projection(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_pk_projection', schema, False) |
| table = self.catalog.get_table('default.test_pk_projection') |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder().with_projection(['dt', 'user_id', 'behavior']) |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| expected = self.expected.select(['dt', 'user_id', 'behavior']) |
| self.assertEqual(actual, expected) |
| |
| def _assert_value_only_projection_works(self, file_format: str, table_suffix: str): |
| # Two commits force the split through the merge path. The merge |
| # reader still needs the PK column to assemble its key, even |
| # though the user-visible projection drops it — regress the |
| # case where narrowing to value-only fields broke the file |
| # column lookup. |
| schema = Schema.from_pyarrow_schema( |
| self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2', 'file.format': file_format}) |
| self.catalog.create_table( |
| 'default.test_pk_projection_no_pk_' + table_suffix, schema, False) |
| table = self.catalog.get_table( |
| 'default.test_pk_projection_no_pk_' + table_suffix) |
| self._write_test_table(table) |
| |
| read_builder = table.new_read_builder().with_projection(['behavior']) |
| actual = self._read_test_table(read_builder) |
| expected = self.expected.select(['behavior']) |
| # Projection drops PKs so we can only compare bag semantics. |
| self.assertEqual( |
| sorted([r['behavior'] for r in actual.to_pylist()], |
| key=lambda v: '' if v is None else v), |
| sorted([r['behavior'] for r in expected.to_pylist()], |
| key=lambda v: '' if v is None else v)) |
| |
| def test_pk_reader_with_projection_excluding_pk(self): |
| self._assert_value_only_projection_works('parquet', 'parquet') |
| |
| def test_pk_reader_with_projection_excluding_pk_orc(self): |
| self._assert_value_only_projection_works('orc', 'orc') |
| |
| def test_pk_reader_with_projection_excluding_pk_avro(self): |
| # Avro path resolves DataField names through ``full_fields_map`` |
| # built from ``self.read_fields``; the alias-safe lookup must also |
| # cover the bare PK name (``user_id``) the file actually stores. |
| self._assert_value_only_projection_works('avro', 'avro') |
| |
| def test_pk_reader_with_limit(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '1'}) |
| self.catalog.create_table('default.test_pk_limit', schema, False) |
| table = self.catalog.get_table('default.test_pk_limit') |
| |
| num_commits = 5 |
| rows_per_commit = 25 |
| |
| for commit_idx in range(num_commits): |
| write_builder = table.new_batch_write_builder() |
| writer = write_builder.new_write() |
| |
| start_id = commit_idx * rows_per_commit |
| end_id = start_id + rows_per_commit |
| |
| if commit_idx > 0: |
| start_id -= 10 |
| |
| batch = pa.Table.from_pydict({ |
| 'user_id': list(range(start_id, end_id)), |
| 'item_id': [1000 + i for i in range(start_id, end_id)], |
| 'behavior': [f'val-{i}' for i in range(start_id, end_id)], |
| 'dt': ['p1' if i % 2 == 0 else 'p2' for i in range(start_id, end_id)], |
| }, schema=self.pa_schema) |
| |
| writer.write_arrow(batch) |
| commit_messages = writer.prepare_commit() |
| commit = write_builder.new_commit() |
| commit.commit(commit_messages) |
| writer.close() |
| |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| all_splits = table_scan.plan().splits() |
| merge_splits = [s for s in all_splits if not s.raw_convertible] |
| self.assertGreater( |
| len(merge_splits), 0, |
| "Should have at least one merge split to test limit with merge scenario") |
| |
| for limit in [5, 10, 20, 50]: |
| read_builder = table.new_read_builder().with_limit(limit) |
| table_read = read_builder.new_read() |
| table_scan = read_builder.new_scan() |
| splits = table_scan.plan().splits() |
| self.assertGreater( |
| len(splits), 0, |
| f"with_limit({limit}) should not produce empty splits for PK table") |
| result = table_read.to_arrow(splits) |
| row_count = result.num_rows if result is not None else 0 |
| self.assertEqual( |
| row_count, limit, |
| f"with_limit({limit}) on PK table must return exactly " |
| f"{limit} rows now that the read-level limit is wired") |
| |
| def test_incremental_timestamp(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_incremental_parquet', schema, False) |
| table = self.catalog.get_table('default.test_incremental_parquet') |
| timestamp = int(time.time() * 1000) |
| self._write_test_table(table) |
| |
| snapshot_manager = table.snapshot_manager() |
| t1 = snapshot_manager.get_snapshot_by_id(1).time_millis |
| t2 = snapshot_manager.get_snapshot_by_id(2).time_millis |
| # test 1 |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp - 1) + ',' + str(timestamp)}) |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder) |
| self.assertEqual(len(actual), 0) |
| # test 2 |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp) + ',' + str(t2)}) |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| self.assertEqual(self.expected, actual) |
| # test 3 |
| table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(t1) + ',' + str(t2)}) |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| expected = pa.Table.from_pydict({ |
| "user_id": [2, 5, 7, 8], |
| "item_id": [1002, 1005, 1007, 1008], |
| "behavior": ["b-new", "e", "g", "h"], |
| "dt": ["p1", "p2", "p1", "p2"] |
| }, schema=self.pa_schema) |
| self.assertEqual(expected, actual) |
| |
| def test_incremental_read_multi_snapshots(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_incremental_read_multi_snapshots', schema, False) |
| table = self.catalog.get_table('default.test_incremental_read_multi_snapshots') |
| write_builder = table.new_batch_write_builder() |
| for i in range(1, 101): |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| pa_table = pa.Table.from_pydict({ |
| 'user_id': [i], |
| 'item_id': [1000 + i], |
| 'behavior': [f'snap{i}'], |
| 'dt': ['p1' if i % 2 == 1 else 'p2'], |
| }, schema=self.pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| snapshot_manager = table.snapshot_manager() |
| t10 = snapshot_manager.get_snapshot_by_id(10).time_millis |
| t20 = snapshot_manager.get_snapshot_by_id(20).time_millis |
| |
| table_inc = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): f"{t10},{t20}"}) |
| read_builder = table_inc.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| |
| expected = pa.Table.from_pydict({ |
| 'user_id': list(range(11, 21)), |
| 'item_id': [1000 + i for i in range(11, 21)], |
| 'behavior': [f'snap{i}' for i in range(11, 21)], |
| 'dt': ['p1' if i % 2 == 1 else 'p2' for i in range(11, 21)], |
| }, schema=self.pa_schema).sort_by('user_id') |
| self.assertEqual(expected, actual) |
| |
| def test_manifest_creation_time_timestamp(self): |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table('default.test_manifest_creation_time', schema, False) |
| table = self.catalog.get_table('default.test_manifest_creation_time') |
| |
| self._write_test_table(table) |
| |
| snapshot_manager = table.snapshot_manager() |
| latest_snapshot = snapshot_manager.get_latest_snapshot() |
| read_builder = table.new_read_builder() |
| table_scan = read_builder.new_scan() |
| manifest_list_manager = table_scan.file_scanner.manifest_list_manager |
| manifest_files = manifest_list_manager.read_all(latest_snapshot) |
| |
| manifest_file_manager = table_scan.file_scanner.manifest_file_manager |
| creation_times_found = [] |
| for manifest_file_meta in manifest_files: |
| entries = manifest_file_manager.read(manifest_file_meta.file_name, drop_stats=False) |
| for entry in entries: |
| if entry.file.creation_time is not None: |
| creation_time = entry.file.creation_time |
| self.assertIsNotNone(creation_time) |
| epoch_millis = entry.file.creation_time_epoch_millis() |
| self.assertIsNotNone(epoch_millis) |
| self.assertGreater(epoch_millis, 0) |
| import time |
| expected_epoch_millis = creation_time.get_millisecond() |
| local_dt = creation_time.to_local_date_time() |
| local_time_struct = local_dt.timetuple() |
| local_timestamp = time.mktime(local_time_struct) |
| local_time_struct_utc = time.gmtime(local_timestamp) |
| utc_timestamp = time.mktime(local_time_struct_utc) |
| expected_epoch_millis = int(utc_timestamp * 1000) |
| self.assertEqual(epoch_millis, expected_epoch_millis) |
| creation_times_found.append(epoch_millis) |
| |
| self.assertGreater( |
| len(creation_times_found), 0, |
| "At least one manifest entry should have creation_time") |
| |
| def _write_test_table(self, table): |
| write_builder = table.new_batch_write_builder() |
| |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [1, 2, 3, 4], |
| 'item_id': [1001, 1002, 1003, 1004], |
| 'behavior': ['a', 'b', 'c', None], |
| 'dt': ['p1', 'p1', 'p2', 'p1'], |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| data1 = { |
| 'user_id': [5, 2, 7, 8], |
| 'item_id': [1005, 1002, 1007, 1008], |
| 'behavior': ['e', 'b-new', 'g', 'h'], |
| 'dt': ['p2', 'p1', 'p1', 'p2'] |
| } |
| pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema) |
| table_write.write_arrow(pa_table) |
| table_commit.commit(table_write.prepare_commit()) |
| table_write.close() |
| table_commit.close() |
| |
| def _read_test_table(self, read_builder): |
| table_read = read_builder.new_read() |
| splits = read_builder.new_scan().plan().splits() |
| return table_read.to_arrow(splits) |
| |
| def test_concurrent_writes_with_retry(self): |
| """Test concurrent writes to verify retry mechanism works correctly for PK tables.""" |
| import threading |
| |
| # Run the test 3 times to verify stability |
| iter_num = 3 |
| for test_iteration in range(iter_num): |
| # Create a unique table for each iteration |
| table_name = f'default.test_pk_concurrent_writes_{test_iteration}' |
| schema = Schema.from_pyarrow_schema(self.pa_schema, |
| partition_keys=['dt'], |
| primary_keys=['user_id', 'dt'], |
| options={'bucket': '2'}) |
| self.catalog.create_table(table_name, schema, False) |
| table = self.catalog.get_table(table_name) |
| |
| write_results = [] |
| write_errors = [] |
| |
| def write_data(thread_id, start_user_id): |
| """Write data in a separate thread.""" |
| try: |
| threading.current_thread().name = f"Iter{test_iteration}-Thread-{thread_id}" |
| write_builder = table.new_batch_write_builder() |
| table_write = write_builder.new_write() |
| table_commit = write_builder.new_commit() |
| |
| # Create unique data for this thread |
| data = { |
| 'user_id': list(range(start_user_id, start_user_id + 5)), |
| 'item_id': [1000 + i for i in range(start_user_id, start_user_id + 5)], |
| 'behavior': [f'thread{thread_id}_{i}' for i in range(5)], |
| 'dt': ['p1' if i % 2 == 0 else 'p2' for i in range(5)], |
| } |
| pa_table = pa.Table.from_pydict(data, schema=self.pa_schema) |
| |
| table_write.write_arrow(pa_table) |
| commit_messages = table_write.prepare_commit() |
| |
| table_commit.commit(commit_messages) |
| table_write.close() |
| table_commit.close() |
| |
| write_results.append({ |
| 'thread_id': thread_id, |
| 'start_user_id': start_user_id, |
| 'success': True |
| }) |
| except Exception as e: |
| write_errors.append({ |
| 'thread_id': thread_id, |
| 'error': str(e) |
| }) |
| |
| # Create and start multiple threads |
| threads = [] |
| num_threads = 10 |
| for i in range(num_threads): |
| thread = threading.Thread( |
| target=write_data, |
| args=(i, i * 10) |
| ) |
| threads.append(thread) |
| thread.start() |
| |
| # Wait for all threads to complete |
| for thread in threads: |
| thread.join() |
| |
| # Verify all writes succeeded (retry mechanism should handle conflicts) |
| self.assertEqual(num_threads, len(write_results), |
| f"Iteration {test_iteration}: Expected {num_threads} successful writes, " |
| f"got {len(write_results)}. Errors: {write_errors}") |
| self.assertEqual(0, len(write_errors), |
| f"Iteration {test_iteration}: Expected no errors, but got: {write_errors}") |
| |
| read_builder = table.new_read_builder() |
| actual = self._read_test_table(read_builder).sort_by('user_id') |
| |
| # Verify data rows (PK table should have unique user_id+dt combinations) |
| self.assertEqual(num_threads * 5, actual.num_rows, |
| f"Iteration {test_iteration}: Expected {num_threads * 5} rows") |
| |
| # Verify user_id |
| user_ids = actual.column('user_id').to_pylist() |
| expected_user_ids = [] |
| for i in range(num_threads): |
| expected_user_ids.extend(range(i * 10, i * 10 + 5)) |
| expected_user_ids.sort() |
| |
| self.assertEqual(user_ids, expected_user_ids, |
| f"Iteration {test_iteration}: User IDs mismatch") |
| |
| # Verify snapshot count (should have num_threads snapshots) |
| snapshot_manager = table.snapshot_manager() |
| latest_snapshot = snapshot_manager.get_latest_snapshot() |
| self.assertIsNotNone(latest_snapshot, |
| f"Iteration {test_iteration}: Latest snapshot should not be None") |
| self.assertEqual(latest_snapshot.id, num_threads, |
| f"Iteration {test_iteration}: Expected snapshot ID {num_threads}, " |
| f"got {latest_snapshot.id}") |
| |
| print(f"✓ PK Table Iteration {test_iteration + 1}/{iter_num} completed successfully") |
