| # 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. | |
| """Native write/commit routing and failure ownership contracts.""" | |
| import json | |
| from types import ModuleType | |
| from unittest.mock import Mock, patch | |
| import pyarrow as pa | |
| import pytest | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.execution.native.runtime import (native_runtime_available) | |
| from pypaimon.execution.native.write.native_commit import _rest_catalog_supported | |
| from pypaimon.write.write_dispatch import NativeTableWrite | |
| pytestmark = [pytest.mark.python_write, pytest.mark.python_commit] | |
| native = pytest.mark.skipif(not native_runtime_available(), reason='pypaimon-rust not installed') | |
| _SCHEMA = pa.schema([('id', pa.int64()), ('pt', pa.string())]) | |
| def _table(tmp_path, catalog=None, options=None): | |
| catalog = catalog or CatalogFactory.create({'warehouse': str(tmp_path / 'warehouse')}) | |
| catalog.create_database('default', True) | |
| catalog.create_table('default.t', Schema.from_pyarrow_schema( | |
| _SCHEMA, partition_keys=['pt'], options=dict({ | |
| 'commit.native.enabled': 'true', 'write.native.enabled': 'false' | |
| }, **(options or {}))), False) | |
| return catalog.get_table('default.t') | |
| def _write(table, rows): | |
| builder = table.new_batch_write_builder() | |
| writer, commit = builder.new_write(), builder.new_commit() | |
| try: | |
| writer.write_arrow(pa.Table.from_pylist(rows, schema=_SCHEMA)) | |
| commit.commit(writer.prepare_commit()) | |
| finally: | |
| writer.close() | |
| commit.close() | |
| def _rows(table): | |
| builder = table.new_read_builder() | |
| return sorted(builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist(), | |
| key=lambda row: row['id']) | |
| @pytest.mark.parametrize('operation,args', [ | |
| ('truncate_table', ()), ('truncate_partitions', ([{'pt': 'a'}],)), | |
| ]) | |
| def test_native_truncate_publication_error_never_retries_or_aborts(tmp_path, operation, args): | |
| table = _table(tmp_path) | |
| commit = table.new_batch_write_builder().new_commit() | |
| native_commit = Mock() | |
| cause = RuntimeError('snapshot may have published before response failed') | |
| getattr(native_commit, operation).side_effect = cause | |
| try: | |
| with patch.object(commit, '_prepare_native_commit', return_value=(native_commit, [])), \ | |
| patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.truncate_table') as truncate, \ | |
| patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.drop_partitions') as drop, \ | |
| patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.abort') as abort: | |
| with pytest.raises(RuntimeError) as error: | |
| getattr(commit, operation)(*args) | |
| assert error.value is cause | |
| truncate.assert_not_called() | |
| drop.assert_not_called() | |
| abort.assert_not_called() | |
| native_commit.abort.assert_not_called() | |
| finally: | |
| commit.close() | |
| def test_abort_then_truncate_partitions_then_commit_keeps_batch_lifecycle(tmp_path): | |
| table = _table(tmp_path) | |
| commit = table.new_batch_write_builder().new_commit() | |
| stream, batch = Mock(), Mock() | |
| try: | |
| with patch( | |
| 'pypaimon.execution.native.write.native_commit.create_native_commit', | |
| side_effect=[stream, batch]) as create, \ | |
| patch('pypaimon.execution.native.write.native_commit.to_native_commit_messages', return_value=[]), \ | |
| patch('pypaimon.execution.native.write.native_commit.native_messages_supported', return_value=True), \ | |
| patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.commit') as fallback: | |
| commit.abort([]) | |
| commit.truncate_partitions([{'pt': 'a'}]) | |
| # Java truncatePartitions does not consume the one-shot batch commit. | |
| commit.table = table.copy({'snapshot.ignore-empty-commit': 'false'}) | |
| commit.commit([]) | |
| stream.abort.assert_called_once_with([]) | |
| stream.close.assert_called_once_with() | |
| batch.truncate_partitions.assert_called_once_with([{'pt': 'a'}]) | |
| batch.commit.assert_called_once_with([]) | |
| assert create.call_args_list[1][1] == {'batch': True} | |
| fallback.assert_not_called() | |
| with pytest.raises(RuntimeError, match='one-time'): | |
| commit.commit([]) | |
| finally: | |
| commit.close() | |
| def _bind_fuse(table, root, cached=False): | |
| from pypaimon.catalog.rest.fuse_support import FusePathResolver | |
| from pypaimon.filesystem.caching_file_io import CachingFileIO | |
| from pypaimon.filesystem.local_file_io import FuseLocalFileIO, LocalFileIO | |
| options = table.catalog_environment.catalog_loader.context().options | |
| options.data.update({'fuse.enabled': 'true', 'fuse.root': root.as_uri(), | |
| 'fuse.mode': 'pvfs', 'fuse.validation-mode': 'none'}) | |
| mount = LocalFileIO()._to_file(FusePathResolver(options, None).resolve_local_path( | |
| table.table_path, table.identifier)).absolute() | |
| mount.parent.mkdir(parents=True, exist_ok=True) | |
| mount.symlink_to(LocalFileIO()._to_file(table.table_path), target_is_directory=True) | |
| io = FuseLocalFileIO(table.table_path, str(mount), options) | |
| table.file_io = CachingFileIO(io, Mock()) if cached else io | |
| return mount | |
| @pytest.mark.parametrize('cached', [False, True]) | |
| def test_builtin_fuse_is_accepted_without_relaxing_custom_file_io(tmp_path, native_rest_catalog, cached): | |
| from pypaimon.filesystem.local_file_io import FuseLocalFileIO | |
| table = _table(tmp_path, native_rest_catalog) | |
| mount = _bind_fuse(table, tmp_path / 'fuse', cached) | |
| assert _rest_catalog_supported(table) | |
| class CustomFuse(FuseLocalFileIO): | |
| pass | |
| table.file_io = CustomFuse(table.table_path, str(mount)) | |
| assert not _rest_catalog_supported(table) | |
| @pytest.mark.parametrize('cached', [False, True]) | |
| @pytest.mark.parametrize('invalid', [None, 'source-uuid', 'source-path', 'target-path', 'target-branch']) | |
| def test_copied_fuse_write_uses_existing_copy_and_checks_identity( | |
| tmp_path, native_rest_catalog, cached, invalid): | |
| from pypaimon.execution.native.write.native_commit import create_native_write_table | |
| original = _table(tmp_path, native_rest_catalog) | |
| _bind_fuse(original, tmp_path / 'fuse', cached) | |
| table = original.copy_without_time_travel({'branch': 'dev'}) | |
| assert table.catalog_environment.uuid is None | |
| assert _rest_catalog_supported(table) | |
| module = ModuleType('pypaimon_rust.datafusion') | |
| module.Table = Mock() | |
| module.PaimonCatalog = Mock() | |
| source = Mock(spec=['location', 'rest_table_uuid', 'copy']) | |
| target = Mock(spec=['location', 'branch', 'copy_with_resolved_schema']) | |
| source.location.return_value = original.table_path if invalid != 'source-path' else '/different' | |
| source.rest_table_uuid.return_value = ( | |
| original.catalog_environment.uuid if invalid != 'source-uuid' else 'replacement-uuid') | |
| source.copy.return_value = target | |
| target.location.return_value = table.table_path if invalid != 'target-path' else '/different' | |
| target.branch.return_value = 'dev' if invalid != 'target-branch' else 'main' | |
| target.copy_with_resolved_schema.return_value = target | |
| module.PaimonCatalog.return_value.get_table.return_value = source | |
| with patch.dict('sys.modules', {'pypaimon_rust.datafusion': module}): | |
| result = create_native_write_table(table) | |
| module.PaimonCatalog.return_value.get_table.assert_called_once_with(('default', 't')) | |
| if invalid is None: | |
| assert result is target | |
| source.copy.assert_called_once_with({'branch': 'dev'}) | |
| target.copy_with_resolved_schema.assert_called_once() | |
| schema = json.loads(target.copy_with_resolved_schema.call_args[0][0]) | |
| assert schema['options']['branch'] == 'dev' | |
| else: | |
| assert result is None | |
| target.copy_with_resolved_schema.assert_not_called() | |
| if invalid.startswith('source-'): | |
| source.copy.assert_not_called() | |
| @pytest.mark.parametrize('identity', [None, '', 123]) | |
| def test_copied_fuse_write_requires_original_rest_uuid(tmp_path, native_rest_catalog, identity): | |
| table = _table(tmp_path, native_rest_catalog) | |
| _bind_fuse(table, tmp_path / 'fuse') | |
| table = table.copy_without_time_travel({'branch': 'dev'}) | |
| response = json.loads(table.catalog_environment.rest_table_response) | |
| response['id'] = identity | |
| table.catalog_environment.rest_table_response = json.dumps(response) | |
| assert not _rest_catalog_supported(table) | |
| @native | |
| @pytest.mark.parametrize('copied_branch', [False, True]) | |
| def test_fuse_native_write_commit_and_truncate_preserve_bound_mount( | |
| tmp_path, native_rest_catalog, copied_branch): | |
| table = _table(tmp_path, native_rest_catalog, {'write.native.enabled': 'true'}) | |
| _bind_fuse(table, tmp_path / 'fuse') | |
| if copied_branch: | |
| table.branch_manager().create_branch('dev') | |
| table = table.copy({'branch': 'dev'}) | |
| builder = table.new_batch_write_builder() | |
| writer, commit = builder.new_write(), builder.new_commit() | |
| assert isinstance(writer, NativeTableWrite) | |
| try: | |
| writer.write_arrow(pa.Table.from_pylist([{'id': 1, 'pt': 'a'}, {'id': 2, 'pt': 'b'}], schema=_SCHEMA)) | |
| with patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.commit', | |
| side_effect=AssertionError('Python commit')): | |
| commit.commit(writer.prepare_commit()) | |
| finally: | |
| writer.close() | |
| commit.close() | |
| assert _rows(table) == [{'id': 1, 'pt': 'a'}, {'id': 2, 'pt': 'b'}] | |
| commit = table.new_batch_write_builder().new_commit() | |
| try: | |
| with patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.drop_partitions', | |
| side_effect=AssertionError('Python truncate')): | |
| commit.truncate_partitions([{'pt': 'a'}]) | |
| assert _rows(table) == [{'id': 2, 'pt': 'b'}] | |
| with patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.truncate_table', | |
| side_effect=AssertionError('Python truncate')): | |
| commit.truncate_table() | |
| assert _rows(table) == [] | |
| finally: | |
| commit.close() | |
| def _merge_table(tmp_path, use_native): | |
| schema = pa.schema([('id', pa.int32()), ('value', pa.int32()), ('payload', pa.string())]) | |
| catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'merge')}) | |
| catalog.create_database('default', True) | |
| catalog.create_table('default.t', Schema.from_pyarrow_schema(schema, options={ | |
| 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', | |
| 'write.native.enabled': str(use_native).lower()}), False) | |
| table = catalog.get_table('default.t') | |
| builder = table.new_batch_write_builder() | |
| writer, commit = builder.new_write(), builder.new_commit() | |
| try: | |
| writer.write_arrow(pa.Table.from_pylist([ | |
| {'id': 1, 'value': 10, 'payload': 'old'}, {'id': 2, 'value': 20, 'payload': 'keep'} | |
| ], schema=schema)) | |
| commit.commit(writer.prepare_commit()) | |
| finally: | |
| writer.close() | |
| commit.close() | |
| return table | |
| @pytest.mark.parametrize('use_native', [False, pytest.param(True, marks=native)]) | |
| @pytest.mark.parametrize('source_key', ['id', 'source_id']) | |
| @pytest.mark.parametrize('dictionary_key', [False, True]) | |
| def test_merge_callable_reads_target_source_and_coalesced_keys(tmp_path, use_native, source_key, dictionary_key): | |
| import pyarrow.compute as pc | |
| import pypaimon.execution.python.write.merge_into as merge_impl | |
| from pypaimon.table.data_evolution_merge_into import WhenMatched, WhenNotMatched | |
| table = _merge_table(tmp_path, use_native) | |
| keys = pa.array([1, 3], pa.int32()) | |
| source = pa.table({source_key: keys.dictionary_encode() if dictionary_key else keys, | |
| 'delta': pa.array([4, 5], pa.int32()), 'label': ['unused', 'new']}) | |
| calls = [] | |
| def value(batch): | |
| calls.append(batch.column_names) | |
| assert batch['s.' + source_key].type == pa.int32() | |
| return pc.add(pc.add(batch['t.value'], batch['s.delta']), batch['s.' + source_key]) | |
| update = table.new_batch_write_builder().new_update() | |
| with patch('pypaimon.execution.python.write.merge_into._build_tables', | |
| wraps=merge_impl._build_tables) as python_merge: | |
| messages = update.merge_into( | |
| source, on={'id': source_key}, | |
| when_matched=[WhenMatched.update({ | |
| 'value': value, 'payload': lambda batch: pc.utf8_upper(batch['t.payload'])})], | |
| when_not_matched=[WhenNotMatched({ | |
| 'value': lambda batch: pc.multiply(batch['s.delta'], 2), | |
| 'payload': lambda batch: batch['s.label']})]) | |
| if use_native: | |
| python_merge.assert_not_called() | |
| table.new_batch_write_builder().new_commit().commit(messages) | |
| assert _rows(table) == [{'id': 1, 'value': 15, 'payload': 'OLD'}, | |
| {'id': 2, 'value': 20, 'payload': 'keep'}, | |
| {'id': 3, 'value': 10, 'payload': 'new'}] | |
| assert len(calls) == 1 | |
| assert calls[0] == sorted(['t._ROW_ID', 't.id', 't.value', 't.payload', | |
| 's.' + source_key, 's.delta', 's.label']) | |
| @pytest.mark.parametrize('use_native', [False, pytest.param(True, marks=native)]) | |
| def test_merge_callable_is_lazy_and_failure_preserves_committed_rows(tmp_path, use_native): | |
| from pypaimon.table.data_evolution_merge_into import WhenMatched | |
| table = _merge_table(tmp_path, use_native) | |
| cause = RuntimeError('callable failure') | |
| calls = [] | |
| def fail(batch): | |
| calls.append(batch.num_rows) | |
| raise cause | |
| update = table.new_batch_write_builder().new_update() | |
| clause = [WhenMatched.update({'value': fail})] | |
| assert update.merge_into(pa.table({'id': pa.array([99], pa.int32())}), on=['id'], | |
| when_matched=clause, when_not_matched=[]) == [] | |
| assert calls == [] | |
| with pytest.raises(RuntimeError) as error: | |
| update.merge_into(pa.table({'id': pa.array([1], pa.int32())}), on=['id'], | |
| when_matched=clause, when_not_matched=[]) | |
| assert error.value is cause | |
| assert calls == [1] | |
| assert _rows(table) == [{'id': 1, 'value': 10, 'payload': 'old'}, | |
| {'id': 2, 'value': 20, 'payload': 'keep'}] | |
| @pytest.mark.parametrize('use_native', [False, pytest.param(True, marks=native)]) | |
| @pytest.mark.parametrize('result_kind', ['chunked', 'wrong_length', 'not_arrow']) | |
| def test_merge_callable_validates_complete_clause_result(tmp_path, use_native, result_kind): | |
| from pypaimon.table.data_evolution_merge_into import WhenMatched | |
| table = _merge_table(tmp_path, use_native) | |
| source = pa.table({'id': pa.array([1, 2], pa.int32())}) | |
| calls = [] | |
| def values(batch): | |
| calls.append(batch.num_rows) | |
| if result_kind == 'chunked': | |
| return pa.chunked_array([pa.array([101]), pa.array([202])]) | |
| if result_kind == 'wrong_length': | |
| return pa.array([1]) | |
| return [101, 202] | |
| update = table.new_batch_write_builder().new_update() | |
| if result_kind == 'chunked': | |
| messages = update.merge_into(source, on=['id'], when_matched=[WhenMatched.update({'value': values})]) | |
| table.new_batch_write_builder().new_commit().commit(messages) | |
| assert [row['value'] for row in _rows(table)] == [101, 202] | |
| else: | |
| with pytest.raises((ValueError, TypeError)): | |
| update.merge_into(source, on=['id'], when_matched=[WhenMatched.update({'value': values})]) | |
| assert [row['value'] for row in _rows(table)] == [10, 20] | |
| assert calls == [2] | |
| def test_merge_callback_aliases_use_coalesced_key_type(): | |
| from pypaimon.ray.data_evolution_merge_transform import _merge_callable_input | |
| dictionary = pa.array([1, 2], pa.int32()).dictionary_encode() | |
| batch = pa.table({'t.id': pa.array([1, 2], pa.int32()), 's.source_id': dictionary, | |
| 's.value': [10, 20], '_scratch': [0, 0]}) | |
| actual = _merge_callable_input(batch, [('source_id', 'id')]) | |
| assert actual.column_names == ['s.source_id', 's.value', 't.id'] | |
| assert actual['s.source_id'].equals(actual['t.id']) | |
| # INSERT callbacks retain source layout because no target key exists. | |
| unmatched = _merge_callable_input(batch.drop(['t.id']), [('source_id', 'id')]) | |
| assert unmatched['s.source_id'].type == dictionary.type | |
| def test_explicit_callable_read_columns_keep_self_merge_projection_narrow(): | |
| from pypaimon.ray.data_evolution_merge_into import _resolve_target_projection | |
| from pypaimon.ray.data_evolution_merge_transform import _NormalizedClause | |
| clauses = [_NormalizedClause({'value': lambda rows: rows['value']})] | |
| fields = ['id', 'value', 'large_blob'] | |
| assert _resolve_target_projection(clauses, ['_ROW_ID'], ['value'], fields, | |
| callable_columns=['value']) == ['value'] | |
| assert _resolve_target_projection(clauses, ['id'], ['value'], fields) == fields | |
| def test_consumed_blob_rows_are_released_without_losing_unplanned_inputs(): | |
| from pypaimon.schema.data_types import AtomicType | |
| from pypaimon.table.row.blob import Blob, BlobDescriptor | |
| from pypaimon.execution.native.write.native_blob_rows import NativeBlobRows | |
| rows = NativeBlobRows() | |
| blob_type = AtomicType('BLOB') | |
| known = BlobDescriptor.deserialize(rows.encode(Blob.from_data(b'known'), blob_type)).uri | |
| unknown = BlobDescriptor.deserialize(rows.encode(Blob.from_data(b'unknown'), blob_type)).uri | |
| reader = rows.create(known) | |
| assert rows.create(known) is reader | |
| rows.release_consumed() | |
| assert not rows._supports_uri(known) | |
| assert rows._supports_uri(unknown) | |
| with rows.create(unknown).new_input_stream(unknown) as stream: | |
| assert stream.read() == b'unknown' | |
| rows.clear() | |
| assert not rows.readers | |
| assert not rows._consumed | |
| @pytest.mark.parametrize('use_native', [False, pytest.param(True, marks=native)]) | |
| def test_automatic_postpone_plan_counts_all_batches_per_partition(tmp_path, use_native): | |
| from pypaimon.write.write_dispatch import NativePostponeFixedBucketTableWrite | |
| schema = pa.schema([('id', pa.int64()), ('pt', pa.string())]) | |
| catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'postpone')}) | |
| catalog.create_database('default', True) | |
| catalog.create_table('default.t', Schema.from_pyarrow_schema( | |
| schema, primary_keys=['id', 'pt'], partition_keys=['pt'], options={ | |
| 'bucket': '-2', 'write.native.enabled': str(use_native).lower(), | |
| 'postpone.target-row-num-per-bucket': '2', | |
| 'postpone.batch-write-fixed-bucket.max-parallelism': '3'}), False) | |
| table = catalog.get_table('default.t') | |
| builder = table.new_postpone_fixed_bucket_write_builder() | |
| writer = builder.new_write() | |
| if use_native: | |
| assert isinstance(writer, NativePostponeFixedBucketTableWrite) | |
| try: | |
| writer.write_arrow(pa.Table.from_pylist([{'id': 1, 'pt': 'a'}, {'id': 2, 'pt': 'b'}], schema=schema)) | |
| writer.write_arrow(pa.Table.from_pylist([{'id': n, 'pt': 'a'} for n in [3, 4, 5, 6]], schema=schema)) | |
| messages = writer.prepare_commit() | |
| assert {m.partition: m.total_buckets for m in messages} == {('a',): 3, ('b',): 1} | |
| builder.new_commit().commit(messages) | |
| finally: | |
| writer.close() | |
| assert _rows(table) == [{'id': 1, 'pt': 'a'}, {'id': 2, 'pt': 'b'}] + [ | |
| {'id': n, 'pt': 'a'} for n in [3, 4, 5, 6]] | |
| @native | |
| @pytest.mark.parametrize('kind', ['boolean', 'date', 'time', 'timestamp', 'decimal', 'float', | |
| 'unicode_int', 'decimal_overflow']) | |
| def test_native_truncate_partition_specs_use_java_typed_strings(tmp_path, native_rest_catalog, kind): | |
| from datetime import date, datetime, time | |
| from decimal import Decimal | |
| cases = { | |
| 'boolean': (pa.bool_(), [True, False], 'TRUE'), | |
| 'date': (pa.date32(), [date(2024, 1, 2), date(2024, 1, 3)], '2024-01-02'), | |
| 'time': (pa.time32('ms'), [time(12, 34, 56, 789000), time(12, 35)], '12:34:56.789'), | |
| 'timestamp': (pa.timestamp('us'), [datetime(2024, 1, 2, 12, 34), datetime(2024, 1, 3)], | |
| '2024-01-02 12:34:00'), | |
| 'decimal': (pa.decimal128(6, 2), [Decimal('12.35'), Decimal('13.00')], '12.345'), | |
| 'float': (pa.float32(), [1.2, 2.0], '1.2'), | |
| 'unicode_int': (pa.int32(), [12, 13], '１２'), | |
| 'decimal_overflow': (pa.decimal128(3, 0), [None, Decimal('12')], '1000'), | |
| } | |
| data_type, values, text = cases[kind] | |
| schema = pa.schema([('id', pa.int64()), ('pt', data_type)]) | |
| native_rest_catalog.create_table('default.typed', Schema.from_pyarrow_schema( | |
| schema, partition_keys=['pt'], options={ | |
| 'write.native.enabled': 'false', 'commit.native.enabled': 'true'}), False) | |
| table = native_rest_catalog.get_table('default.typed') | |
| builder = table.new_batch_write_builder() | |
| writer, commit = builder.new_write(), builder.new_commit() | |
| try: | |
| writer.write_arrow(pa.table({'id': [1, 2], 'pt': values}, schema=schema)) | |
| commit.commit(writer.prepare_commit()) | |
| finally: | |
| writer.close() | |
| commit.close() | |
| commit = table.new_batch_write_builder().new_commit() | |
| try: | |
| with patch('pypaimon.execution.python.write.file_store_commit.' | |
| 'FileStoreCommit.drop_partitions', | |
| side_effect=AssertionError('Python truncate')): | |
| commit.truncate_partitions([{'pt': text}]) | |
| assert [row['id'] for row in _rows(table)] == [2] | |
| finally: | |
| commit.close() | |
| @pytest.mark.python_plan | |
| @pytest.mark.python_read | |
| @pytest.mark.parametrize('fail', [False, True]) | |
| def test_postpone_blob_row_bridge_keeps_sources_until_prepare(fail): | |
| from types import SimpleNamespace | |
| from pypaimon.common.options.core_options import CoreOptions | |
| from pypaimon.common.options.options import Options | |
| from pypaimon.schema.data_types import AtomicType, DataField | |
| from pypaimon.table.row.blob import Blob, BlobDescriptor | |
| from pypaimon.table.row.generic_row import GenericRow | |
| from pypaimon.write.write_dispatch import NativePostponeFixedBucketTableWrite | |
| # Exercise the bridge's deferred input lifetime without creating a table: | |
| # Python catalog validation does not allow a primary-key BLOB schema. | |
| fields = [DataField(0, 'id', AtomicType('INT')), DataField(1, 'pt', AtomicType('STRING')), | |
| DataField(2, 'payload', AtomicType('BLOB'))] | |
| table = SimpleNamespace( | |
| table_schema=SimpleNamespace(fields=fields, primary_keys=['id', 'pt']), fields=fields, | |
| field_dict={field.name: field for field in fields}, partition_keys=['pt'], | |
| is_primary_key_table=True, cross_partition_update=False, | |
| options=CoreOptions(Options({'bucket': '-2', 'postpone.target-row-num-per-bucket': '1'}))) | |
| backend = Mock() | |
| staged, payloads = [], [] | |
| backend.write_arrow.side_effect = lambda batch, bucket: staged.append(batch) | |
| writer = NativePostponeFixedBucketTableWrite(table, 'job', None, None, backend) | |
| factory = backend.with_blob_uri_reader_factory.call_args[0][0] | |
| error = RuntimeError('prepare failed') | |
| def prepare(): | |
| for batch in staged: | |
| for encoded in batch.column('payload').to_pylist(): | |
| uri = BlobDescriptor.deserialize(encoded).uri | |
| with factory.create(uri).new_input_stream(uri) as stream: | |
| payloads.append(stream.read()) | |
| if fail: | |
| raise error | |
| return [] | |
| backend.prepare_commit.side_effect = prepare | |
| try: | |
| for identifier in [1, 2]: | |
| writer.write_row(GenericRow( | |
| [identifier, 'new', Blob.from_data(str(identifier).encode())], fields)) | |
| assert len(factory.readers) == 2 | |
| assert not payloads | |
| if fail: | |
| with pytest.raises(RuntimeError) as raised: | |
| writer.prepare_commit() | |
| assert raised.value is error | |
| else: | |
| assert writer.prepare_commit() == [] | |
| assert payloads == [b'1', b'2'] | |
| assert not factory.readers | |
| backend.abort.assert_not_called() | |
| finally: | |
| writer.close() |