[python] Track managed BLOB packs and delete their sidecars on abort.

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.

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)

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