| # 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_composite_primary_key_order(self): | |
| cases = [ | |
| (['customer_id', 'order_id'], []), | |
| (['order_id', 'customer_id'], []), | |
| (['order_id', 'dt', 'customer_id'], ['dt']), | |
| ] | |
| for index, (primary_keys, partition_keys) in enumerate(cases): | |
| with self.subTest(primary_keys=primary_keys): | |
| fields = [('customer_id', pa.int64()), ('order_id', pa.int64()), | |
| ('amount', pa.int64())] | |
| base = {'customer_id': [1, 2, 1, 2], 'order_id': [10, 10, 20, 20], | |
| 'amount': [100, 200, 300, 400]} | |
| update = {'customer_id': [1], 'order_id': [20], 'amount': [350]} | |
| expected = {'customer_id': [1, 1, 2, 2], 'order_id': [10, 20, 10, 20], | |
| 'amount': [100, 350, 200, 400]} | |
| if partition_keys: | |
| fields.append(('dt', pa.string())) | |
| base['dt'] = expected['dt'] = ['p1'] * 4 | |
| update['dt'] = ['p1'] | |
| arrow_schema = pa.schema(fields) | |
| name = 'default.test_composite_pk_order_{}'.format(index) | |
| self.catalog.create_table(name, Schema.from_pyarrow_schema( | |
| arrow_schema, primary_keys=primary_keys, partition_keys=partition_keys, | |
| options={'bucket': '1'}), False) | |
| table = self.catalog.get_table(name) | |
| for data in (base, update): | |
| builder = table.new_batch_write_builder() | |
| writer, committer = builder.new_write(), builder.new_commit() | |
| try: | |
| writer.write_arrow(pa.Table.from_pydict(data, schema=arrow_schema)) | |
| committer.commit(writer.prepare_commit()) | |
| finally: | |
| writer.close() | |
| committer.close() | |
| actual = self._read_test_table(table.new_read_builder()).sort_by( | |
| [('customer_id', 'ascending'), ('order_id', 'ascending')]) | |
| self.assertEqual(actual.to_pylist(), | |
| pa.Table.from_pydict(expected, schema=arrow_schema).to_pylist()) | |
| projected = self._read_test_table( | |
| table.new_read_builder().with_projection(['amount'])) | |
| self.assertEqual(sorted(projected['amount'].to_pylist()), [100, 200, 350, 400]) | |
| 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_nullable_primary_key(self): | |
| nullable_schema = pa.schema([ | |
| pa.field('id', pa.int64()), | |
| pa.field('value', pa.string()), | |
| ]) | |
| schema = Schema.from_pyarrow_schema( | |
| nullable_schema, | |
| primary_keys=['id'], | |
| options={ | |
| 'bucket': '3', | |
| CoreOptions.PRIMARY_KEY_NULLABLE.key(): 'true', | |
| }, | |
| ) | |
| self.catalog.create_table('default.test_nullable_pk', schema, False) | |
| table = self.catalog.get_table('default.test_nullable_pk') | |
| for rows in [ | |
| [{'id': None, 'value': 'old'}], | |
| [{'id': None, 'value': 'new'}, | |
| {'id': 1, 'value': 'one'}, | |
| {'id': 2, 'value': 'two'}]]: | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| writer.write_arrow(pa.Table.from_pylist(rows, schema=nullable_schema)) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| actual = self._read_test_table(table.new_read_builder()).to_pylist() | |
| actual.sort(key=lambda row: (-1 if row['id'] is None else row['id'])) | |
| self.assertEqual( | |
| actual, | |
| [{'id': None, 'value': 'new'}, | |
| {'id': 1, 'value': 'one'}, | |
| {'id': 2, 'value': 'two'}], | |
| ) | |
| 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 reads still merge the two versions of user_id=2. | |
| snapshot_rows = self._read_test_table(table.new_read_builder()).sort_by('user_id') | |
| self.assertEqual(self.expected, snapshot_rows) | |
| 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) | |
| # The full incremental window retains both committed versions of user_id=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', 'ascending'), ('behavior', 'ascending')]) | |
| expected = pa.Table.from_pydict({ | |
| 'user_id': [1, 2, 2, 3, 4, 5, 7, 8], | |
| 'item_id': [1001, 1002, 1002, 1003, 1004, 1005, 1007, 1008], | |
| 'behavior': ['a', 'b', 'b-new', 'c', None, 'e', 'g', 'h'], | |
| 'dt': ['p1', 'p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2'], | |
| }, schema=self.pa_schema) | |
| self.assertEqual(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") |