| # 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 unittest |
| import uuid |
| from datetime import datetime |
| from unittest.mock import MagicMock, Mock, patch |
| |
| from pypaimon.manifest.schema.data_file_meta import DataFileMeta |
| from pypaimon.manifest.schema.manifest_entry import ManifestEntry |
| from pypaimon.manifest.schema.manifest_file_meta import ManifestFileMeta |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| from pypaimon.snapshot.snapshot_commit import PartitionStatistics |
| from pypaimon.table.row.binary_row import BinaryRow |
| from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer |
| from pypaimon.write.commit.row_id_conflict_rewriter import RowIdRewriteResult |
| from pypaimon.write.commit_message import CommitMessage |
| from pypaimon.write.file_store_commit import ( |
| FileStoreCommit, |
| ManifestMergeResult, |
| RetryResult, |
| RewriteResult, |
| _try_replace_manifest_files, |
| ) |
| |
| |
| @patch('pypaimon.write.file_store_commit.ManifestFileManager') |
| @patch('pypaimon.write.file_store_commit.ManifestListManager') |
| class TestFileStoreCommit(unittest.TestCase): |
| """Test cases for FileStoreCommit class.""" |
| |
| def setUp(self): |
| """Set up test fixtures.""" |
| # Mock table with required attributes |
| self.mock_table = Mock() |
| self.mock_table.partition_keys = ['dt', 'region'] |
| self.mock_table.current_branch.return_value = 'main' |
| self.mock_table.table_path = '/test/table/path' |
| self.mock_table.file_io = Mock() |
| self.mock_table.options.manifest_target_size.return_value = 8 * 1024 * 1024 |
| self.mock_table.options.manifest_merge_min_count.return_value = 30 |
| |
| # Mock snapshot commit |
| self.mock_snapshot_commit = Mock() |
| |
| def _create_file_store_commit(self): |
| """Helper method to create FileStoreCommit instance.""" |
| return FileStoreCommit( |
| snapshot_commit=self.mock_snapshot_commit, |
| table=self.mock_table, |
| commit_user='test_user' |
| ) |
| |
| @staticmethod |
| def _manifest_meta(name): |
| row = GenericRowSerializer.to_bytes(GenericRow([], [])) |
| return ManifestFileMeta( |
| file_name=name, |
| file_size=10, |
| num_added_files=1, |
| num_deleted_files=0, |
| partition_stats=SimpleStats( |
| BinaryRow(row, []), BinaryRow(row, []), []), |
| schema_id=0, |
| ) |
| |
| def test_replace_manifest_files_uses_stable_value_equality( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| previous = [self._manifest_meta('a'), self._manifest_meta('b')] |
| current = [ |
| self._manifest_meta('prefix'), |
| self._manifest_meta('a'), |
| self._manifest_meta('b'), |
| self._manifest_meta('suffix'), |
| ] |
| replacement = [self._manifest_meta('merged')] |
| |
| result = _try_replace_manifest_files( |
| current, previous, replacement) |
| |
| self.assertEqual( |
| ['prefix', 'merged', 'suffix'], |
| [manifest.file_name for manifest in result], |
| ) |
| self.assertIsNot(current[1], previous[0]) |
| |
| def test_replace_manifest_files_preserves_order_and_empty_semantics( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| a = self._manifest_meta('a') |
| b = self._manifest_meta('b') |
| merged = self._manifest_meta('merged') |
| |
| self.assertIsNone(_try_replace_manifest_files( |
| [a, self._manifest_meta('x'), b], [a, b], [merged])) |
| self.assertEqual( |
| ['a', 'merged'], |
| [manifest.file_name for manifest in _try_replace_manifest_files( |
| [a, self._manifest_meta('a'), b], [a, b], [merged])], |
| ) |
| self.assertEqual( |
| [merged], _try_replace_manifest_files([], [], [merged])) |
| self.assertIsNone(_try_replace_manifest_files([a], [], [merged])) |
| |
| def test_manifest_merge_result_copies_and_freezes_lists( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| before = [self._manifest_meta('before')] |
| after = [self._manifest_meta('after')] |
| |
| result = ManifestMergeResult(before, after) |
| before.clear() |
| after.clear() |
| |
| self.assertIsInstance(result.merge_before_manifests, tuple) |
| self.assertIsInstance(result.merge_after_manifests, tuple) |
| self.assertEqual( |
| ['before'], |
| [manifest.file_name |
| for manifest in result.merge_before_manifests], |
| ) |
| self.assertEqual( |
| ['after'], |
| [manifest.file_name |
| for manifest in result.merge_after_manifests], |
| ) |
| |
| def _run_manifest_commit_attempt(self, commit_side_effect=None, |
| commit_result=None, retry_result=None, |
| existing_manifests=None, |
| merged_manifests=None, |
| latest_watermark=None): |
| file_store_commit = self._create_file_store_commit() |
| self.mock_table.identifier = 'default.test_table' |
| self.mock_table.table_schema.id = 7 |
| self.mock_table.options.row_tracking_enabled.return_value = False |
| |
| snapshot_commit = MagicMock() |
| snapshot_commit.__enter__.return_value = snapshot_commit |
| snapshot_commit.__exit__.return_value = False |
| snapshot_commit.commit.side_effect = commit_side_effect |
| snapshot_commit.commit.return_value = commit_result |
| file_store_commit.snapshot_commit = snapshot_commit |
| |
| before = self._manifest_meta('before') |
| after = self._manifest_meta('after') |
| existing_manifests = ( |
| [before] if existing_manifests is None |
| else existing_manifests) |
| merged_manifests = ( |
| [after] if merged_manifests is None |
| else merged_manifests) |
| delta = self._manifest_meta('delta') |
| file_store_commit._write_manifest_files = Mock( |
| return_value=[delta]) |
| file_store_commit._generate_partition_statistics = Mock( |
| return_value=[]) |
| file_store_commit.manifest_list_manager.read_all.return_value = ( |
| existing_manifests) |
| file_store_commit.manifest_file_merger = Mock() |
| file_store_commit.manifest_file_merger.merge.return_value = ( |
| merged_manifests, merged_manifests) |
| file_store_commit._clean_up_reuse_tmp_manifests = Mock() |
| file_store_commit._clean_up_no_reuse_tmp_manifests = Mock() |
| |
| latest_snapshot = Mock( |
| id=3, |
| uuid='base-snapshot-uuid', |
| total_record_count=10, |
| index_manifest=None, |
| watermark=latest_watermark, |
| ) |
| commit_entry = Mock(kind=0) |
| commit_entry.file.row_count = 2 |
| result = file_store_commit._try_commit_once( |
| retry_result=retry_result, |
| commit_kind='APPEND', |
| commit_entries=[commit_entry], |
| changelog_entries=[], |
| commit_identifier=11, |
| latest_snapshot=latest_snapshot, |
| ) |
| return file_store_commit, result |
| |
| def test_append_commit_inherits_watermark( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| file_store_commit, result = self._run_manifest_commit_attempt( |
| commit_result=True, |
| latest_watermark=123, |
| ) |
| |
| self.assertTrue(result.is_success()) |
| committed_snapshot = ( |
| file_store_commit.snapshot_commit.commit.call_args[0][1]) |
| self.assertEqual(123, committed_snapshot.watermark) |
| |
| def test_false_atomic_commit_retains_manifest_merge_result( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| file_store_commit, result = self._run_manifest_commit_attempt( |
| commit_result=False) |
| |
| self.assertIsInstance(result, RetryResult) |
| self.assertIsNone(result.exception) |
| self.assertEqual( |
| ['before'], |
| [manifest.file_name for manifest |
| in result.manifest_merge_result.merge_before_manifests], |
| ) |
| self.assertEqual( |
| ['after'], |
| [manifest.file_name for manifest |
| in result.manifest_merge_result.merge_after_manifests], |
| ) |
| file_store_commit.manifest_file_merger.merge.assert_called_once() |
| |
| def test_atomic_commit_exception_does_not_retain_manifest_merge_result( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| failure = TimeoutError('lost commit response') |
| file_store_commit, result = self._run_manifest_commit_attempt( |
| commit_side_effect=failure) |
| |
| self.assertIsInstance(result, RetryResult) |
| self.assertIs(failure, result.exception) |
| self.assertTrue(result.commit_result_may_be_uncertain) |
| self.assertIsNone(result.manifest_merge_result) |
| file_store_commit._clean_up_reuse_tmp_manifests.assert_not_called() |
| file_store_commit._clean_up_no_reuse_tmp_manifests.assert_not_called() |
| |
| def test_retry_reuses_manifest_merge_and_preserves_surrounding_files( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| previous_before = [ |
| self._manifest_meta('before-a'), |
| self._manifest_meta('before-b'), |
| ] |
| previous_after = [self._manifest_meta('merged')] |
| retry_result = RetryResult( |
| Mock(id=3), |
| manifest_merge_result=ManifestMergeResult( |
| previous_before, previous_after), |
| ) |
| current = [ |
| self._manifest_meta('prefix'), |
| self._manifest_meta('before-a'), |
| self._manifest_meta('before-b'), |
| self._manifest_meta('suffix'), |
| ] |
| |
| file_store_commit, result = self._run_manifest_commit_attempt( |
| commit_result=False, |
| retry_result=retry_result, |
| existing_manifests=current, |
| ) |
| |
| self.assertIsInstance(result, RetryResult) |
| self.assertEqual( |
| ['prefix', 'before-a', 'before-b', 'suffix'], |
| [manifest.file_name for manifest |
| in result.manifest_merge_result.merge_before_manifests], |
| ) |
| self.assertEqual( |
| ['prefix', 'merged', 'suffix'], |
| [manifest.file_name for manifest |
| in result.manifest_merge_result.merge_after_manifests], |
| ) |
| file_store_commit.manifest_file_merger.merge.assert_not_called() |
| base_manifests = ( |
| file_store_commit.manifest_list_manager.write |
| .call_args_list[-1].args[1]) |
| self.assertEqual( |
| ['prefix', 'merged', 'suffix'], |
| [manifest.file_name for manifest in base_manifests], |
| ) |
| |
| def test_retry_skips_manifest_merge_when_previous_input_is_not_contiguous( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| previous_before = [ |
| self._manifest_meta('before-a'), |
| self._manifest_meta('before-b'), |
| ] |
| retry_result = RetryResult( |
| Mock(id=3), |
| manifest_merge_result=ManifestMergeResult( |
| previous_before, [self._manifest_meta('merged')]), |
| ) |
| current = [ |
| self._manifest_meta('before-a'), |
| self._manifest_meta('between'), |
| self._manifest_meta('before-b'), |
| ] |
| |
| file_store_commit, result = self._run_manifest_commit_attempt( |
| commit_result=False, |
| retry_result=retry_result, |
| existing_manifests=current, |
| ) |
| |
| self.assertIsInstance(result, RetryResult) |
| self.assertIsNone(result.manifest_merge_result) |
| file_store_commit.manifest_file_merger.merge.assert_not_called() |
| base_manifests = ( |
| file_store_commit.manifest_list_manager.write |
| .call_args_list[-1].args[1]) |
| self.assertEqual(current, base_manifests) |
| |
| def test_manifest_merge_runs_once_across_multiple_retries( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| first_commit, retry_result = self._run_manifest_commit_attempt( |
| commit_result=False, |
| existing_manifests=[self._manifest_meta('before')], |
| merged_manifests=[self._manifest_meta('merged')], |
| ) |
| first_commit.manifest_file_merger.merge.assert_called_once() |
| |
| unchanged_retry, retry_result = self._run_manifest_commit_attempt( |
| commit_result=False, |
| retry_result=retry_result, |
| existing_manifests=[self._manifest_meta('before')], |
| ) |
| retry_commits = [unchanged_retry] |
| current_names = ['before'] |
| for suffix in ['concurrent-1', 'concurrent-2']: |
| current_names.append(suffix) |
| retry_commit, retry_result = self._run_manifest_commit_attempt( |
| commit_result=False, |
| retry_result=retry_result, |
| existing_manifests=[ |
| self._manifest_meta(name) for name in current_names |
| ], |
| ) |
| retry_commits.append(retry_commit) |
| |
| final_names = current_names + ['concurrent-3'] |
| final_commit, result = self._run_manifest_commit_attempt( |
| commit_result=True, |
| retry_result=retry_result, |
| existing_manifests=[ |
| self._manifest_meta(name) for name in final_names |
| ], |
| ) |
| |
| self.assertTrue(result.is_success()) |
| for retry_commit in retry_commits + [final_commit]: |
| retry_commit.manifest_file_merger.merge.assert_not_called() |
| base_manifests = ( |
| final_commit.manifest_list_manager.write |
| .call_args_list[-1].args[1]) |
| self.assertEqual( |
| ['merged', 'concurrent-1', 'concurrent-2', 'concurrent-3'], |
| [manifest.file_name for manifest in base_manifests], |
| ) |
| |
| def test_generate_partition_statistics_single_partition_single_file( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation with single partition and single file.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| # Create test data |
| creation_time_dt = datetime(2024, 1, 15, 10, 30, 0) |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| creation_time = Timestamp.from_local_date_time(creation_time_dt) |
| file_meta = DataFileMeta.create( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, # 1MB |
| row_count=10000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time, |
| external_path=None, |
| first_row_id=None, |
| write_cols=None |
| ) |
| |
| commit_message = CommitMessage( |
| partition=('2024-01-15', 'us-east-1'), |
| bucket=0, |
| new_files=[file_meta] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message])) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 1) |
| |
| stat = statistics[0] |
| self.assertIsInstance(stat, PartitionStatistics) |
| self.assertEqual(stat.spec, {'dt': '2024-01-15', 'region': 'us-east-1'}) |
| self.assertEqual(stat.record_count, 10000) |
| self.assertEqual(stat.file_count, 1) |
| self.assertEqual(stat.file_size_in_bytes, 1024 * 1024) |
| expected_time = file_meta.creation_time_epoch_millis() |
| self.assertEqual(stat.last_file_creation_time, expected_time) |
| |
| def test_generate_partition_statistics_multiple_files_same_partition( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation with multiple files in same partition.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| creation_time_1 = Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0)) |
| creation_time_2 = Timestamp.from_local_date_time(datetime(2024, 1, 15, 11, 30, 0)) # Later time |
| |
| file_meta_1 = DataFileMeta.create( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, # 1MB |
| row_count=10000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time_1 |
| ) |
| |
| file_meta_2 = DataFileMeta.create( |
| file_name="test_file_2.parquet", |
| file_size=2 * 1024 * 1024, # 2MB |
| row_count=15000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=101, |
| max_sequence_number=200, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time_2 |
| ) |
| |
| commit_message = CommitMessage( |
| partition=('2024-01-15', 'us-east-1'), |
| bucket=0, |
| new_files=[file_meta_1, file_meta_2] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message])) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 1) |
| |
| stat = statistics[0] |
| self.assertEqual(stat.spec, {'dt': '2024-01-15', 'region': 'us-east-1'}) |
| self.assertEqual(stat.record_count, 25000) # 10000 + 15000 |
| self.assertEqual(stat.file_count, 2) |
| self.assertEqual(stat.file_size_in_bytes, 3 * 1024 * 1024) # 1MB + 2MB |
| expected_time = file_meta_2.creation_time_epoch_millis() |
| self.assertEqual(stat.last_file_creation_time, expected_time) |
| |
| def test_partition_statistics_use_replacement_bucket_count( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| file_store_commit = self._create_file_store_commit() |
| file_meta = Mock(row_count=1, file_size=10, creation_time=None) |
| |
| for old_buckets, new_buckets in [(-2, 2), (2, 3)]: |
| with self.subTest(old=old_buckets, new=new_buckets): |
| partition = GenericRow(['2024-01-15', 'us-east-1'], None) |
| entries = [ |
| ManifestEntry( |
| kind=1, |
| partition=partition, |
| bucket=0, |
| total_buckets=old_buckets, |
| file=file_meta, |
| ), |
| ManifestEntry( |
| kind=0, |
| partition=partition, |
| bucket=0, |
| total_buckets=new_buckets, |
| file=file_meta, |
| ), |
| ] |
| |
| statistics = ( |
| file_store_commit._generate_partition_statistics(entries)) |
| self.assertEqual(new_buckets, statistics[0].total_buckets) |
| |
| def test_generate_partition_statistics_multiple_partitions( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation with multiple different partitions.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| creation_time_dt = datetime(2024, 1, 15, 10, 30, 0) |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| creation_time = Timestamp.from_local_date_time(creation_time_dt) |
| |
| # File for partition 1 |
| file_meta_1 = DataFileMeta.create( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, |
| row_count=10000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time, |
| external_path=None, |
| first_row_id=None, |
| write_cols=None |
| ) |
| |
| # File for partition 2 |
| file_meta_2 = DataFileMeta.create( |
| file_name="test_file_2.parquet", |
| file_size=2 * 1024 * 1024, |
| row_count=20000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=101, |
| max_sequence_number=200, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time, |
| external_path=None, |
| first_row_id=None, |
| write_cols=None |
| ) |
| |
| commit_message_1 = CommitMessage( |
| partition=('2024-01-15', 'us-east-1'), |
| bucket=0, |
| new_files=[file_meta_1] |
| ) |
| |
| commit_message_2 = CommitMessage( |
| partition=('2024-01-15', 'us-west-2'), |
| bucket=0, |
| new_files=[file_meta_2] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics( |
| self._to_entries([commit_message_1, commit_message_2])) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 2) |
| |
| # Sort statistics by partition spec for consistent testing |
| statistics.sort(key=lambda s: s.spec['region']) |
| |
| # Check first partition (us-east-1) |
| stat_1 = statistics[0] |
| self.assertEqual(stat_1.spec, {'dt': '2024-01-15', 'region': 'us-east-1'}) |
| self.assertEqual(stat_1.record_count, 10000) |
| self.assertEqual(stat_1.file_count, 1) |
| self.assertEqual(stat_1.file_size_in_bytes, 1024 * 1024) |
| |
| # Check second partition (us-west-2) |
| stat_2 = statistics[1] |
| self.assertEqual(stat_2.spec, {'dt': '2024-01-15', 'region': 'us-west-2'}) |
| self.assertEqual(stat_2.record_count, 20000) |
| self.assertEqual(stat_2.file_count, 1) |
| self.assertEqual(stat_2.file_size_in_bytes, 2 * 1024 * 1024) |
| |
| def test_generate_partition_statistics_unpartitioned_table( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation for unpartitioned table.""" |
| # Update mock table to have no partition keys |
| self.mock_table.partition_keys = [] |
| |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| creation_time_dt = datetime(2024, 1, 15, 10, 30, 0) |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| creation_time = Timestamp.from_local_date_time(creation_time_dt) |
| file_meta = DataFileMeta.create( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, |
| row_count=10000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time, |
| external_path=None, |
| first_row_id=None, |
| write_cols=None |
| ) |
| |
| commit_message = CommitMessage( |
| partition=(), # Empty partition for unpartitioned table |
| bucket=0, |
| new_files=[file_meta] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message])) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 1) |
| |
| stat = statistics[0] |
| self.assertEqual(stat.spec, {}) # Empty spec for unpartitioned table |
| self.assertEqual(stat.record_count, 10000) |
| self.assertEqual(stat.file_count, 1) |
| self.assertEqual(stat.file_size_in_bytes, 1024 * 1024) |
| |
| def test_generate_partition_statistics_no_creation_time( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation when file has no creation time.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| file_meta = DataFileMeta( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, |
| row_count=10000, |
| min_key=None, |
| max_key=None, |
| key_stats=None, |
| value_stats=None, |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=None, |
| ) |
| |
| commit_message = CommitMessage( |
| partition=('2024-01-15', 'us-east-1'), |
| bucket=0, |
| new_files=[file_meta] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message])) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 1) |
| |
| stat = statistics[0] |
| # Should have a valid timestamp (current time) |
| self.assertGreater(stat.last_file_creation_time, 0) |
| |
| def test_generate_partition_statistics_mismatched_partition_keys( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation when partition tuple doesn't match partition keys.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| # Table has 2 partition keys but partition tuple has 3 values |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.table.row.generic_row import GenericRow |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| file_meta = DataFileMeta.create( |
| file_name="test_file_1.parquet", |
| file_size=1024 * 1024, |
| row_count=10000, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=100, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0)) |
| ) |
| |
| commit_message = CommitMessage( |
| partition=('2024-01-15', 'us-east-1', 'extra-value'), # 3 values but table has 2 keys |
| bucket=0, |
| new_files=[file_meta] |
| ) |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics(self._to_entries([commit_message])) |
| |
| # Verify results - should fallback to index-based naming |
| self.assertEqual(len(statistics), 1) |
| |
| stat = statistics[0] |
| expected_spec = { |
| 'partition_0': '2024-01-15', |
| 'partition_1': 'us-east-1', |
| 'partition_2': 'extra-value' |
| } |
| self.assertEqual(stat.spec, expected_spec) |
| |
| def test_generate_partition_statistics_empty_commit_messages( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| """Test partition statistics generation with empty commit messages list.""" |
| # Create FileStoreCommit instance |
| file_store_commit = self._create_file_store_commit() |
| |
| # Test method |
| statistics = file_store_commit._generate_partition_statistics([]) |
| |
| # Verify results |
| self.assertEqual(len(statistics), 0) |
| |
| def test_append_commit_inherits_index_manifest( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| file_store_commit = self._create_file_store_commit() |
| |
| self.mock_table.identifier = 'default.test_table' |
| self.mock_table.table_schema = Mock() |
| self.mock_table.table_schema.id = 7 |
| self.mock_table.options.row_tracking_enabled.return_value = False |
| |
| snapshot_commit = MagicMock() |
| snapshot_commit.__enter__.return_value = snapshot_commit |
| snapshot_commit.__exit__.return_value = False |
| snapshot_commit.commit.return_value = True |
| file_store_commit.snapshot_commit = snapshot_commit |
| |
| file_store_commit._write_manifest_files = Mock(return_value=[Mock()]) |
| file_store_commit._generate_partition_statistics = Mock(return_value=[]) |
| file_store_commit.manifest_list_manager.read_all.return_value = [] |
| |
| latest_snapshot = Mock() |
| latest_snapshot.id = 3 |
| latest_snapshot.uuid = "base-snapshot-uuid" |
| latest_snapshot.total_record_count = 10 |
| latest_snapshot.index_manifest = "index-manifest-existing" |
| latest_snapshot.watermark = None |
| |
| commit_entry = Mock() |
| commit_entry.kind = 0 |
| commit_entry.file = Mock() |
| commit_entry.file.row_count = 2 |
| |
| result = file_store_commit._try_commit_once( |
| retry_result=None, |
| commit_kind="APPEND", |
| commit_entries=[commit_entry], |
| changelog_entries=[], |
| commit_identifier=11, |
| latest_snapshot=latest_snapshot |
| ) |
| |
| self.assertTrue(result.is_success()) |
| self.assertEqual( |
| "base-snapshot-uuid", |
| snapshot_commit.commit.call_args[0][0], |
| ) |
| committed_snapshot = snapshot_commit.commit.call_args[0][1] |
| self.assertEqual( |
| "index-manifest-existing", |
| committed_snapshot.index_manifest |
| ) |
| self.assertEqual(str(uuid.UUID(committed_snapshot.uuid)), committed_snapshot.uuid) |
| |
| def test_null_partition_value( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| from pypaimon.data.timestamp import Timestamp |
| from pypaimon.manifest.schema.simple_stats import SimpleStats |
| from pypaimon.schema.data_types import DataField, AtomicType |
| |
| file_store_commit = self._create_file_store_commit() |
| self.mock_table.partition_keys = ['dt'] |
| self.mock_table.partition_keys_fields = [ |
| DataField(0, 'dt', AtomicType('STRING')) |
| ] |
| self.mock_table.table_schema = Mock() |
| self.mock_table.table_schema.id = 0 |
| |
| file_store_commit.manifest_file_manager = Mock() |
| file_store_commit.manifest_file_manager.manifest_path = '/test/manifest' |
| self.mock_table.file_io.get_file_size.return_value = 1024 |
| |
| creation_time = Timestamp.from_local_date_time(datetime(2024, 1, 15, 10, 30, 0)) |
| |
| def make_file(name): |
| return DataFileMeta.create( |
| file_name=name, |
| file_size=1024, |
| row_count=100, |
| min_key=GenericRow([], []), |
| max_key=GenericRow([], []), |
| key_stats=SimpleStats.empty_stats(), |
| value_stats=SimpleStats.empty_stats(), |
| min_sequence_number=1, |
| max_sequence_number=10, |
| schema_id=0, |
| level=0, |
| extra_files=[], |
| creation_time=creation_time, |
| ) |
| |
| entries = [ |
| ManifestEntry(kind=0, partition=GenericRow([None], None), bucket=0, total_buckets=None, |
| file=make_file("f1.parquet")), |
| ManifestEntry(kind=0, partition=GenericRow(['2024-01-15'], None), bucket=0, total_buckets=None, |
| file=make_file("f2.parquet")), |
| ] |
| |
| result = file_store_commit._write_manifest_files(entries, "manifest-test") |
| self.assertIsNotNone(result) |
| |
| def test_row_id_rewrite_respects_commit_retry_limit( |
| self, mock_manifest_list_manager, mock_manifest_file_manager): |
| file_store_commit = self._create_file_store_commit() |
| file_store_commit.commit_max_retries = 1 |
| file_store_commit.commit_timeout = 10 ** 9 |
| file_store_commit._commit_retry_wait = Mock() |
| |
| latest_snapshot = Mock() |
| latest_snapshot.id = 7 |
| file_store_commit.snapshot_manager.get_latest_snapshot.return_value = ( |
| latest_snapshot |
| ) |
| |
| commit_entry = Mock() |
| rewrite = RewriteResult(RowIdRewriteResult([commit_entry], 1)) |
| file_store_commit._try_commit_once = Mock(side_effect=[ |
| rewrite, |
| rewrite, |
| AssertionError("rewrite retry budget was not enforced"), |
| ]) |
| |
| with self.assertRaises(RuntimeError) as ctx: |
| file_store_commit._try_commit( |
| commit_kind="APPEND", |
| commit_identifier=11, |
| commit_entries_plan=lambda snapshot: [commit_entry], |
| ) |
| |
| self.assertIn("with 1 retries", str(ctx.exception)) |
| self.assertEqual(2, file_store_commit._try_commit_once.call_count) |
| file_store_commit._commit_retry_wait.assert_called_once_with(0) |
| |
| @staticmethod |
| def _to_entries(commit_messages): |
| commit_entries = [] |
| for msg in commit_messages: |
| partition = GenericRow(list(msg.partition), None) |
| for file in msg.new_files: |
| commit_entries.append(ManifestEntry( |
| kind=0, |
| partition=partition, |
| bucket=msg.bucket, |
| total_buckets=None, |
| file=file |
| )) |
| return commit_entries |
