From e721d6fb8793b9fedee35ef5fc6b9cbda79422b6 Mon Sep 17 00:00:00 2001 From: "wenchao.wu" Date: Tue, 22 Sep 2026 15:29:57 +0800 Subject: [PATCH] [python] Track managed BLOB packs and delete their sidecars on abort. The .blobref file lists packs referenced by one data file. Commit abort removes those sidecars and .row files, and leaves .blob packs in place when a BlobConsumer owns them. prepare_commit still keeps file ownership on the writer. --- paimon-python/pypaimon/blob/__init__.py | 16 ++ .../blob/managed_blob_reference_collector.py | 168 +++++++++++++++++ .../blob/managed_blob_reference_file.py | 173 ++++++++++++++++++ .../tests/managed_blob_lifecycle_test.py | 116 ++++++++++++ .../tests/managed_blob_reference_file_test.py | 162 ++++++++++++++++ .../managed_blob_write_ownership_test.py | 65 +++++++ .../pypaimon/write/commit_message.py | 1 + .../pypaimon/write/file_store_commit.py | 46 ++++- .../pypaimon/write/file_store_write.py | 5 + 9 files changed, 750 insertions(+), 2 deletions(-) create mode 100644 paimon-python/pypaimon/blob/__init__.py create mode 100644 paimon-python/pypaimon/blob/managed_blob_reference_collector.py create mode 100644 paimon-python/pypaimon/blob/managed_blob_reference_file.py create mode 100644 paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py create mode 100644 paimon-python/pypaimon/tests/managed_blob_reference_file_test.py create mode 100644 paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py diff --git a/paimon-python/pypaimon/blob/__init__.py b/paimon-python/pypaimon/blob/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/paimon-python/pypaimon/blob/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/paimon-python/pypaimon/blob/managed_blob_reference_collector.py b/paimon-python/pypaimon/blob/managed_blob_reference_collector.py new file mode 100644 index 000000000000..99a5faceefb2 --- /dev/null +++ b/paimon-python/pypaimon/blob/managed_blob_reference_collector.py @@ -0,0 +1,168 @@ +# 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. + +from typing import List, Set + +import pyarrow as pa + +from pypaimon.blob.managed_blob_reference_file import ( + ManagedBlobReferenceFile, + Reference, +) +from pypaimon.schema.data_types import ( + DataField, + is_array_blob_type, + is_blob_file_field, + is_blob_type, + is_map_blob_type, +) +from pypaimon.table.row.blob import Blob, BlobDescriptor, BlobRef +from pypaimon.table.row.row_kind import RowKind + + +class ManagedBlobReferenceCollector: + """Collect managed BLOB pack references from rows written to one data file.""" + + _RETRACT_KINDS = frozenset({ + RowKind.UPDATE_BEFORE.value, + RowKind.DELETE.value, + }) + + def __init__( + self, + file_io, + data_file_path: str, + value_fields: List[DataField], + managed_blob_fields: Set[str]): + self._file_io = file_io + self._sidecar_path = ManagedBlobReferenceFile.sidecar_path(data_file_path) + self._field_specs = [] + for field in value_fields: + if field.name not in managed_blob_fields: + continue + if not is_blob_file_field(field): + continue + if is_blob_type(field.type): + kind = "scalar" + elif is_array_blob_type(field.type): + kind = "array" + elif is_map_blob_type(field.type): + kind = "map" + else: + continue + self._field_specs.append((field.name, kind)) + self._descriptor_uris: Set[str] = set() + self._closed = False + self._aborted = False + + def collect_table(self, data: pa.Table, value_kind_column: str = "_VALUE_KIND") -> None: + if self._closed: + raise RuntimeError("Managed BLOB reference collector is already closed.") + if not self._field_specs: + return + if value_kind_column not in data.schema.names: + raise ValueError("Missing value kind column %r." % value_kind_column) + + columns = [] + names = set(data.schema.names) + for field_name, field_kind in self._field_specs: + if field_name not in names: + continue + columns.append((data.column(field_name).to_pylist(), field_kind)) + if not columns: + return + + kinds = data.column(value_kind_column).to_pylist() + for row_idx, kind in enumerate(kinds): + if kind in self._RETRACT_KINDS: + continue + for values, field_kind in columns: + self._collect_value(values[row_idx], field_kind) + + def close(self) -> str: + if self._aborted: + raise RuntimeError("Managed BLOB reference collector was aborted.") + if self._closed: + return self._sidecar_path.rsplit("/", 1)[-1] + references: List[Reference] = [] + for descriptor_uri in self._descriptor_uris: + reference = ManagedBlobReferenceFile.from_descriptor_uri(descriptor_uri) + if reference is not None: + references.append(reference) + self._descriptor_uris.clear() + try: + ManagedBlobReferenceFile.write(self._file_io, self._sidecar_path, references) + except Exception: + self.abort() + raise + self._closed = True + return self._sidecar_path.rsplit("/", 1)[-1] + + def abort(self) -> None: + self._file_io.delete_quietly(self._sidecar_path) + self._descriptor_uris.clear() + self._aborted = True + self._closed = True + + def _collect_value(self, value, field_kind: str) -> None: + if value is None: + return + if hasattr(value, "as_py"): + value = value.as_py() + if value is None: + return + if field_kind == "scalar": + self._collect_descriptor_bytes(value) + return + if field_kind == "array": + if not hasattr(value, "__iter__") or isinstance(value, (bytes, bytearray, str)): + return + for element in value: + if element is not None: + self._collect_descriptor_bytes(element) + return + if field_kind == "map": + for _, map_value in _map_items(value): + if map_value is not None: + self._collect_descriptor_bytes(map_value) + + def _collect_descriptor_bytes(self, raw) -> None: + if isinstance(raw, BlobRef): + self._descriptor_uris.add(raw.to_descriptor().uri) + return + if isinstance(raw, Blob): + descriptor = raw.to_descriptor() + if descriptor.uri.endswith(ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX): + self._descriptor_uris.add(descriptor.uri) + return + if isinstance(raw, (bytes, bytearray)): + # Values in a managed BLOB data-file column are descriptors by + # contract. Parse them non-heuristically so legacy v1 descriptors + # (which have no magic header) contribute their pack reference. + descriptor = BlobDescriptor.deserialize(bytes(raw)) + if descriptor.uri.endswith(ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX): + self._descriptor_uris.add(descriptor.uri) + + +def _map_items(value): + if value is None: + return [] + if isinstance(value, dict): + return value.items() + if hasattr(value, "items"): + return value.items() + return list(value) diff --git a/paimon-python/pypaimon/blob/managed_blob_reference_file.py b/paimon-python/pypaimon/blob/managed_blob_reference_file.py new file mode 100644 index 000000000000..118d0f6fae73 --- /dev/null +++ b/paimon-python/pypaimon/blob/managed_blob_reference_file.py @@ -0,0 +1,173 @@ +# 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 io +import struct +import zlib +from typing import List, Optional, Tuple + +from pypaimon.index.pk.primary_key_index_source_meta import ( + _decode_modified_utf8, + _encode_modified_utf8, +) + + +class ManagedBlobReferenceFile: + """Versioned metadata listing managed BLOB packs referenced by one data file.""" + + MAGIC = 0x50424C52 + VERSION = 1 + MANAGED_BLOB_SUFFIX = ".managed.blob" + REFERENCE_FILE_SUFFIX = ".blobref" + + @staticmethod + def from_descriptor_uri(descriptor_uri: str) -> Optional["Reference"]: + if not descriptor_uri.endswith(ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX): + return None + parent, _, name = descriptor_uri.rpartition("/") + if not parent or not name: + return None + return Reference(parent, name) + + @staticmethod + def sidecar_path(data_file_path: str) -> str: + return data_file_path + ManagedBlobReferenceFile.REFERENCE_FILE_SUFFIX + + @staticmethod + def sidecar_name(data_file_name: str) -> str: + return data_file_name + ManagedBlobReferenceFile.REFERENCE_FILE_SUFFIX + + @staticmethod + def write(file_io, path: str, references: List["Reference"]) -> None: + normalized = sorted( + references, + key=lambda ref: (ref.storage_root_id, ref.relative_path), + ) + unique: List[Reference] = [] + for reference in normalized: + if not unique or reference != unique[-1]: + unique.append(reference) + + payload = io.BytesIO() + payload.write(struct.pack(">B", ManagedBlobReferenceFile.VERSION)) + payload.write(struct.pack(">i", len(unique))) + for reference in unique: + _write_modified_utf(payload, reference.storage_root_id) + _write_modified_utf(payload, reference.relative_path) + + payload_bytes = payload.getvalue() + checksum = zlib.crc32(payload_bytes) & 0xFFFFFFFF + checksum_signed = ( + checksum if checksum < 0x80000000 else checksum - 0x100000000 + ) + try: + with file_io.new_output_stream(path) as out: + out.write(struct.pack(">i", ManagedBlobReferenceFile.MAGIC)) + out.write(payload_bytes) + out.write(struct.pack(">i", checksum_signed)) + except Exception: + file_io.delete_quietly(path) + raise + + @staticmethod + def read(file_io, path: str) -> List["Reference"]: + with file_io.new_input_stream(path) as stream: + data = stream.read() + # magic 4 + version 1 + count 4 + crc 4. Shorter input must not + # surface as struct.error from the checksum unpack. + if len(data) < 13: + raise IOError("Invalid managed BLOB reference file: too short") + magic = struct.unpack_from(">i", data, 0)[0] + if magic != ManagedBlobReferenceFile.MAGIC: + raise IOError("Invalid managed BLOB reference file magic: %s" % magic) + + offset = 4 + version = data[offset] + offset += 1 + if version != ManagedBlobReferenceFile.VERSION: + raise IOError("Unsupported managed BLOB reference file version: %s" % version) + + count = struct.unpack_from(">i", data, offset)[0] + offset += 4 + if count < 0: + raise IOError("Invalid managed BLOB reference count: %s" % count) + + references: List[Reference] = [] + for _ in range(count): + root, offset = _read_modified_utf(data, offset) + rel, offset = _read_modified_utf(data, offset) + references.append(Reference(root, rel)) + + if offset + 4 > len(data): + raise IOError("Invalid managed BLOB reference file checksum") + expected_checksum = struct.unpack_from(">i", data, offset)[0] & 0xFFFFFFFF + actual_checksum = zlib.crc32(data[4:offset]) & 0xFFFFFFFF + if expected_checksum != actual_checksum: + raise IOError( + "Invalid managed BLOB reference file checksum. Expected %s but computed %s." + % (expected_checksum, actual_checksum)) + if offset + 4 != len(data): + raise IOError("Unexpected trailing bytes in managed BLOB reference file.") + return references + + +class Reference: + """Exact identity of a managed BLOB payload pack.""" + + __slots__ = ("storage_root_id", "relative_path") + + def __init__(self, storage_root_id: str, relative_path: str): + if not storage_root_id: + raise ValueError("Managed BLOB storage root must not be empty.") + if not relative_path or "/" in relative_path or relative_path in (".", ".."): + raise ValueError( + "Managed BLOB relative path must be a file name: %s." % relative_path) + self.storage_root_id = storage_root_id + self.relative_path = relative_path + + def __eq__(self, other) -> bool: + if not isinstance(other, Reference): + return False + return ( + self.storage_root_id == other.storage_root_id + and self.relative_path == other.relative_path + ) + + def __hash__(self) -> int: + return hash((self.storage_root_id, self.relative_path)) + + def __repr__(self) -> str: + return "%s/%s" % (self.storage_root_id, self.relative_path) + + +def _write_modified_utf(out: io.BytesIO, value: str) -> None: + encoded = _encode_modified_utf8(value) + if len(encoded) > 65535: + raise ValueError("Modified UTF-8 string is too long.") + out.write(struct.pack(">H", len(encoded))) + out.write(encoded) + + +def _read_modified_utf(data: bytes, offset: int) -> Tuple[str, int]: + if offset + 2 > len(data): + raise IOError("Truncated modified UTF-8 length.") + length = struct.unpack_from(">H", data, offset)[0] + offset += 2 + end = offset + length + if end > len(data): + raise IOError("Truncated modified UTF-8 payload.") + return _decode_modified_utf8(data[offset:end]), end diff --git a/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py b/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py new file mode 100644 index 000000000000..11249af8cfb9 --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py @@ -0,0 +1,116 @@ +# 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 +from unittest.mock import Mock, call + +from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.file_store_commit import FileStoreCommit + + +class ManagedBlobLifecycleTest(unittest.TestCase): + + def test_commit_abort_deletes_sidecars_and_unpreserved_packs(self): + file_io = Mock() + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io) + + data_file = Mock( + external_path=None, + file_path="/warehouse/table/bucket-0/data.avro", + extra_files=[ + "data.avro.blobref", + "data.avro.row", + "pack.managed.blob", + ], + ) + message = CommitMessage( + partition=(), + bucket=0, + new_files=[data_file], + preserve_blob_files_on_abort=False, + ) + + commit.abort([message]) + + file_io.delete_quietly.assert_has_calls([ + call("/warehouse/table/bucket-0/data.avro"), + call("/warehouse/table/bucket-0/data.avro.blobref"), + call("/warehouse/table/bucket-0/data.avro.row"), + call("/warehouse/table/bucket-0/pack.managed.blob"), + ], any_order=True) + + def test_commit_abort_preserves_consumer_owned_blob_packs(self): + file_io = Mock() + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io) + + data_file = Mock( + external_path=None, + file_path="/warehouse/table/bucket-0/data.avro", + extra_files=["data.avro.blobref", "pack.managed.blob"], + ) + blob_file = Mock( + external_path=None, + file_path="/warehouse/table/bucket-0/payload.blob", + extra_files=[], + ) + message = CommitMessage( + partition=(), + bucket=0, + new_files=[data_file, blob_file], + preserve_blob_files_on_abort=True, + ) + + commit.abort([message]) + + file_io.delete_quietly.assert_has_calls([ + call("/warehouse/table/bucket-0/data.avro"), + call("/warehouse/table/bucket-0/data.avro.blobref"), + ]) + deleted = [args[0] for args, _ in file_io.delete_quietly.call_args_list] + self.assertNotIn("/warehouse/table/bucket-0/payload.blob", deleted) + self.assertNotIn("/warehouse/table/bucket-0/pack.managed.blob", deleted) + + def test_commit_abort_deletes_resolved_data_file_when_extras_fail(self): + file_io = Mock() + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io) + + class _BoomExtras(object): + external_path = None + file_path = "/warehouse/table/bucket-0/data.avro" + + @property + def extra_files(self): + raise RuntimeError("boom resolving extras") + + message = CommitMessage( + partition=(), + bucket=0, + new_files=[_BoomExtras()], + preserve_blob_files_on_abort=False, + ) + + commit.abort([message]) + + file_io.delete_quietly.assert_called_once_with( + "/warehouse/table/bucket-0/data.avro") + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/pypaimon/tests/managed_blob_reference_file_test.py b/paimon-python/pypaimon/tests/managed_blob_reference_file_test.py new file mode 100644 index 000000000000..9ac3f2388460 --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_reference_file_test.py @@ -0,0 +1,162 @@ +# 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 io +import struct +import tempfile +import unittest + +from pypaimon.blob.managed_blob_reference_file import ( + ManagedBlobReferenceFile, + Reference, + _read_modified_utf, + _write_modified_utf, +) +from pypaimon.blob.managed_blob_reference_collector import ( + ManagedBlobReferenceCollector, +) +from pypaimon.common.file_io import FileIO +from pypaimon.common.options.options import Options + + +class ManagedBlobReferenceFileTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.file_io = FileIO.get(self.temp_dir, Options({})) + + def test_round_trip_and_deduplicate_references(self): + path = f"{self.temp_dir}/data.avro.blobref" + first = Reference(f"{self.temp_dir}/bucket-0", "data-a.managed.blob") + second = Reference(f"{self.temp_dir}/bucket-0", "data-b.managed.blob") + + ManagedBlobReferenceFile.write(self.file_io, path, [second, first, second]) + self.assertEqual( + ManagedBlobReferenceFile.read(self.file_io, path), + [first, second], + ) + + empty_path = f"{self.temp_dir}/empty.avro.blobref" + ManagedBlobReferenceFile.write(self.file_io, empty_path, []) + self.assertEqual(ManagedBlobReferenceFile.read(self.file_io, empty_path), []) + + def test_classify_managed_blob_path(self): + managed = f"{self.temp_dir}/bucket-0/data-b.managed.blob" + ordinary = f"{self.temp_dir}/bucket-0/data-d.blob" + reference = ManagedBlobReferenceFile.from_descriptor_uri(managed) + self.assertEqual( + reference, + Reference(f"{self.temp_dir}/bucket-0", "data-b.managed.blob"), + ) + self.assertIsNone(ManagedBlobReferenceFile.from_descriptor_uri(ordinary)) + self.assertEqual( + ManagedBlobReferenceFile.sidecar_name("data-a.avro"), + "data-a.avro.blobref", + ) + + def test_reference_collector_close_is_idempotent(self): + collector = ManagedBlobReferenceCollector( + self.file_io, + f"{self.temp_dir}/data.avro", + [], + set(), + ) + + first = collector.close() + second = collector.close() + + self.assertEqual(first, "data.avro.blobref") + self.assertEqual(second, first) + + def test_reference_is_hashable(self): + reference = Reference(f"{self.temp_dir}/bucket-0", "data-a.managed.blob") + self.assertIn(reference, {reference}) + + def test_close_after_abort_raises(self): + collector = ManagedBlobReferenceCollector( + self.file_io, + f"{self.temp_dir}/data.avro", + [], + set(), + ) + collector.abort() + with self.assertRaisesRegex(RuntimeError, "aborted"): + collector.close() + + def test_reject_truncated_reference_file(self): + path = f"{self.temp_dir}/short.avro.blobref" + with self.file_io.new_output_stream(path) as out: + out.write(struct.pack(">i", ManagedBlobReferenceFile.MAGIC)) + out.write(struct.pack(">i", 0)) + with self.assertRaisesRegex(IOError, "too short"): + ManagedBlobReferenceFile.read(self.file_io, path) + + def test_modified_utf8_round_trip(self): + payload = io.BytesIO() + value = "storage/root/测试" + _write_modified_utf(payload, value) + decoded, offset = _read_modified_utf(payload.getvalue(), 0) + self.assertEqual(decoded, value) + self.assertEqual(offset, len(payload.getvalue())) + + def test_reference_file_matches_fixed_binary_fixture(self): + # Fixed bytes guard the on-disk layout independently of this module's + # writer/reader round trip. The fixture covers big-endian framing, + # modified UTF-8, and the CRC payload boundary. + fixture = bytes.fromhex( + "50424c520100000001001866696c653a2f2f2f77617265686f7573652f" + "eda0bdedb8800013646174612d612e6d616e616765642e626c6f62f2cb93c4" + ) + expected = [Reference("file:///warehouse/😀", "data-a.managed.blob")] + fixture_path = f"{self.temp_dir}/fixture.blobref" + with self.file_io.new_output_stream(fixture_path) as out: + out.write(fixture) + + self.assertEqual(ManagedBlobReferenceFile.read( + self.file_io, fixture_path), expected) + + written_path = f"{self.temp_dir}/written.blobref" + ManagedBlobReferenceFile.write(self.file_io, written_path, expected) + with self.file_io.new_input_stream(written_path) as stream: + self.assertEqual(stream.read(), fixture) + + def test_reject_unsupported_version(self): + path = f"{self.temp_dir}/unsupported.avro.blobref" + with self.file_io.new_output_stream(path) as out: + out.write(struct.pack(">i", ManagedBlobReferenceFile.MAGIC)) + out.write(struct.pack(">B", 99)) + out.write(struct.pack(">i", 0)) + out.write(struct.pack(">i", 0)) + with self.assertRaisesRegex(IOError, "Unsupported managed BLOB reference file version"): + ManagedBlobReferenceFile.read(self.file_io, path) + + def test_reject_corrupt_checksum(self): + path = f"{self.temp_dir}/corrupt.avro.blobref" + payload = io.BytesIO() + payload.write(struct.pack(">B", ManagedBlobReferenceFile.VERSION)) + payload.write(struct.pack(">i", 0)) + payload_bytes = payload.getvalue() + with self.file_io.new_output_stream(path) as out: + out.write(struct.pack(">i", ManagedBlobReferenceFile.MAGIC)) + out.write(payload_bytes) + out.write(struct.pack(">i", 12345)) + with self.assertRaisesRegex(IOError, "Invalid managed BLOB reference file checksum"): + ManagedBlobReferenceFile.read(self.file_io, path) + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py b/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py new file mode 100644 index 000000000000..95819dfe4460 --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py @@ -0,0 +1,65 @@ +# 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 + +from pypaimon.write.file_store_write import FileStoreWrite + + +class _KeepingWriter: + """Returns file lists without giving up the writer's own copies.""" + + def __init__(self, files): + self.committed_files = list(files) + self.committed_changelog_files = [] + + def prepare_commit(self): + return list(self.committed_files) + + def prepare_changelog_commit(self): + return list(self.committed_changelog_files) + + +def _file_store(writer, blob_consumer=None): + file_store = object.__new__(FileStoreWrite) + file_store.data_writers = {((), 0): writer} + file_store.commit_identifier = 0 + file_store._runtime_total_buckets = {} + file_store.blob_consumer = blob_consumer + return file_store + + +class ManagedBlobWriteOwnershipTest(unittest.TestCase): + + def test_prepare_commit_keeps_files_for_writer_abort(self): + writer = _KeepingWriter(["data"]) + messages = _file_store(writer).prepare_commit(1) + + self.assertEqual(messages[0].new_files, ["data"]) + self.assertFalse(messages[0].preserve_blob_files_on_abort) + self.assertEqual(writer.committed_files, ["data"]) + + def test_blob_consumer_marks_messages_to_preserve_packs(self): + writer = _KeepingWriter(["pack.blob"]) + messages = _file_store(writer, blob_consumer=object()).prepare_commit(1) + + self.assertTrue(messages[0].preserve_blob_files_on_abort) + self.assertEqual(writer.committed_files, ["pack.blob"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/pypaimon/write/commit_message.py b/paimon-python/pypaimon/write/commit_message.py index 1325260cd213..3c52d1d7f70b 100644 --- a/paimon-python/pypaimon/write/commit_message.py +++ b/paimon-python/pypaimon/write/commit_message.py @@ -41,6 +41,7 @@ class CommitMessage: compact_changelog_files: List[DataFileMeta] = field(default_factory=list) compact_index_adds: List['IndexManifestEntry'] = field(default_factory=list) compact_index_deletes: List['IndexManifestEntry'] = field(default_factory=list) + preserve_blob_files_on_abort: bool = False def is_empty(self): return ( diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index afe747662053..3b0dde986b7b 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -85,12 +85,43 @@ def _reject_compact_increment(messages: List[CommitMessage]): 'Committing a compact increment requires a separate COMPACT snapshot.') +def _preserve_blob_files(message) -> bool: + return message.preserve_blob_files_on_abort + + +def _is_preserved_blob_pack(path) -> bool: + return str(path).endswith("." + CoreOptions.FILE_FORMAT_BLOB) + + +def _aligned_extra_file_path(file, extra_file: str, resolved_path=None) -> str: + if "://" in extra_file or extra_file.startswith("/"): + return extra_file + file_path = file.external_path or file.file_path or resolved_path + if not file_path or "/" not in str(file_path): + return extra_file + return "{}/{}".format(str(file_path).rsplit("/", 1)[0], extra_file) + + +def _delete_abort_paths(table, paths): + for path_to_delete in paths: + try: + table.file_io.delete_quietly(str(path_to_delete)) + except Exception as error: + logger.warning( + "Failed to clean up file %s during abort: %s", + path_to_delete, + error, + ) + + def _abort_commit_messages(table, commit_messages: List[CommitMessage]): """Delete files created by messages known to be uncommitted.""" for message in commit_messages: + preserve_blob_files = _preserve_blob_files(message) for file in (list(message.new_files) + list(message.changelog_files) + list(message.compact_after) + list(message.compact_changelog_files)): + paths = [] path = None try: path = file.external_path or file.file_path @@ -98,14 +129,25 @@ def _abort_commit_messages(table, commit_messages: List[CommitMessage]): bucket_path = table.path_factory().bucket_path( tuple(message.partition), message.bucket) path = '%s/%s' % (bucket_path.rstrip('/'), file.file_name) - if path: - table.file_io.delete_quietly(str(path)) + if path and not ( + preserve_blob_files and _is_preserved_blob_pack(path)): + paths.append(path) + for extra_file in (getattr(file, "extra_files", None) or []): + extra_path = _aligned_extra_file_path(file, extra_file, path) + if preserve_blob_files and _is_preserved_blob_pack(extra_path): + continue + paths.append(extra_path) except Exception as error: + # Extras failed after the data path was resolved. Delete that + # path before logging, then skip the rest of this file. + _delete_abort_paths(table, paths) logger.warning( "Failed to clean up file %s during abort: %s", path, error, ) + continue + _delete_abort_paths(table, paths) for entry in message.index_adds + message.compact_index_adds: file_name = None try: diff --git a/paimon-python/pypaimon/write/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index a73b9df77dbf..7bffc76755a5 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/paimon-python/pypaimon/write/file_store_write.py @@ -322,6 +322,10 @@ def _has_vector_columns(self) -> bool: def prepare_commit(self, commit_identifier) -> List[CommitMessage]: self.commit_identifier = commit_identifier commit_messages = [] + # A BlobConsumer owns the pack bytes. Commit abort must leave those + # ``.blob`` files alone; the writer still owns every other file and + # ``writer.abort()`` can still delete them. + preserve_blob_files = self.blob_consumer is not None for (partition, bucket), writer in self.data_writers.items(): committed_files = writer.prepare_commit() changelog_files = writer.prepare_changelog_commit() @@ -332,6 +336,7 @@ def prepare_commit(self, commit_identifier) -> List[CommitMessage]: new_files=committed_files, changelog_files=changelog_files, total_buckets=self._runtime_total_buckets.get(partition), + preserve_blob_files_on_abort=preserve_blob_files, ) commit_messages.append(commit_message) return commit_messages