From 38d3a064289852c1f428dee6ef001019fcc225ea Mon Sep 17 00:00:00 2001 From: "wenchao.wu" Date: Thu, 8 Oct 2026 10:10:20 +0800 Subject: [PATCH 1/2] [python] Track managed BLOB packs and delete their sidecars on abort. The .blobref file lists packs referenced by one data file. Explicit commit abort deletes newly written data, BLOB packs, sidecars, and .row files. Writer abort still leaves packs in place when a BlobConsumer has been given them. prepare_commit still keeps file ownership on the writer. --- paimon-python/README.md | 7 +- paimon-python/pypaimon/blob/__init__.py | 16 ++ .../blob/managed_blob_reference_collector.py | 168 +++++++++++ .../blob/managed_blob_reference_file.py | 178 ++++++++++++ .../python/write/file_store_commit.py | 26 +- .../index/pk/primary_key_index_source_meta.py | 17 +- .../pypaimon/tests/blob_table_test.py | 44 +++ .../tests/managed_blob_lifecycle_test.py | 120 ++++++++ .../tests/managed_blob_reference_file_test.py | 260 ++++++++++++++++++ .../managed_blob_write_ownership_test.py | 68 +++++ 10 files changed, 896 insertions(+), 8 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/README.md b/paimon-python/README.md index 4bcdd29995ad..6b8d91c1f055 100644 --- a/paimon-python/README.md +++ b/paimon-python/README.md @@ -300,9 +300,10 @@ mixing Arrow batches and Python `Blob` objects in one writer. A successful `prepare_commit()` transfers file ownership to the caller. Writer `abort()` and `close()` preserve those prepared files and may clean only -outputs still owned by the writer. Explicit `commit.abort(messages)` deletes -newly written data, changelog, sidecars and added indexes, matching Java's -Paimon Table committer. Call it only when those messages are known to be +outputs still owned by the writer. A BlobConsumer keeps that writer from +deleting packs it has already exposed. Explicit `commit.abort(messages)` deletes +newly written data, BLOB packs, changelog, sidecars and added indexes, matching +Java's Paimon Table committer. Call it only when those messages are known to be uncommitted and will never be submitted. Never call it after a commit with an unknown outcome: publication may have succeeded before its response failed. Internal failure paths preserve prepared files instead of calling commit abort. 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..341a013fcbc1 --- /dev/null +++ b/paimon-python/pypaimon/blob/managed_blob_reference_file.py @@ -0,0 +1,178 @@ +# 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.") + try: + value = _decode_modified_utf8(data[offset:end]) + except UnicodeDecodeError as error: + raise IOError( + "Invalid modified UTF-8 in managed BLOB reference file.") from error + return value, end diff --git a/paimon-python/pypaimon/execution/python/write/file_store_commit.py b/paimon-python/pypaimon/execution/python/write/file_store_commit.py index de76eff182b8..ebebf63d1bce 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_commit.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_commit.py @@ -86,6 +86,18 @@ def _reject_compact_increment(messages: List[CommitMessage]): 'Committing a compact increment requires a separate COMPACT snapshot.') +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, + ) + + class CommitResult: """Base class for commit results.""" @@ -972,19 +984,29 @@ def abort(self, commit_messages: List[CommitMessage]): for file in (list(message.new_files) + list(message.changelog_files) + list(message.compact_after) + list(message.compact_changelog_files)): + paths = [] path = None + data_path = None try: if not file.physical_path(): file.file_path = canonical_data_file_path( self.table, tuple(message.partition), message.bucket, file.file_name) + data_path = file.physical_path() for path in file.collect_files(): - self.table.file_io.delete_quietly(path) + paths.append(path) except Exception as error: + # Collection failed after the data path was resolved. Delete + # that path before logging, then skip the rest of this file. + if data_path and data_path not in paths: + paths.append(data_path) + _delete_abort_paths(self.table, paths) logger.warning( "Failed to clean up file %s during abort: %s", - path, + path or data_path, error, ) + continue + _delete_abort_paths(self.table, paths) for entry in message.index_adds + message.compact_index_adds: file_name = None try: diff --git a/paimon-python/pypaimon/index/pk/primary_key_index_source_meta.py b/paimon-python/pypaimon/index/pk/primary_key_index_source_meta.py index 3766efedcdca..fd1eb5a0d715 100644 --- a/paimon-python/pypaimon/index/pk/primary_key_index_source_meta.py +++ b/paimon-python/pypaimon/index/pk/primary_key_index_source_meta.py @@ -112,11 +112,22 @@ def _decode_modified_utf8(data: bytes) -> str: code = first offset += 1 elif first & 0xE0 == 0xC0 and offset + 1 < len(data): - code = ((first & 0x1F) << 6) | (data[offset + 1] & 0x3F) + second = data[offset + 1] + # Java DataInputStream.readUTF requires continuation bytes to be 10xxxxxx. + if second & 0xC0 != 0x80: + raise UnicodeDecodeError( + "modified-utf8", data, offset, offset + 2, + "invalid continuation byte") + code = ((first & 0x1F) << 6) | (second & 0x3F) offset += 2 elif first & 0xF0 == 0xE0 and offset + 2 < len(data): - code = ((first & 0x0F) << 12) | ((data[offset + 1] & 0x3F) << 6) \ - | (data[offset + 2] & 0x3F) + second = data[offset + 1] + third = data[offset + 2] + if second & 0xC0 != 0x80 or third & 0xC0 != 0x80: + raise UnicodeDecodeError( + "modified-utf8", data, offset, offset + 3, + "invalid continuation byte") + code = ((first & 0x0F) << 12) | ((second & 0x3F) << 6) | (third & 0x3F) offset += 3 else: raise UnicodeDecodeError("modified-utf8", data, offset, offset + 1, diff --git a/paimon-python/pypaimon/tests/blob_table_test.py b/paimon-python/pypaimon/tests/blob_table_test.py index 53f050bc0d1e..6cbaaf01a89e 100755 --- a/paimon-python/pypaimon/tests/blob_table_test.py +++ b/paimon-python/pypaimon/tests/blob_table_test.py @@ -6291,6 +6291,50 @@ def my_consumer(field_name, descriptor): data = Blob.from_descriptor(uri_reader, desc).to_data() self.assertEqual(data, blob_bytes) + def test_flush_only_consumer_commit_abort_deletes_new_files(self): + """Explicit commit abort deletes prepared data and BLOB packs.""" + pa_schema = pa.schema([ + ('id', pa.int32()), + ('blob_data', pa.large_binary()), + ]) + schema = Schema.from_pyarrow_schema(pa_schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + 'blob.target-file-size': '1KB', + }) + self.catalog.create_table('test_db.blob_consumer_commit_abort', schema, False) + table = self.catalog.get_table('test_db.blob_consumer_commit_abort') + + def _files(root): + found = set() + for dirpath, _, filenames in os.walk(root): + for name in filenames: + found.add(os.path.join(dirpath, name)) + return found + + write_builder = table.new_batch_write_builder() + writer = write_builder.new_write() + commit = write_builder.new_commit() + try: + # Flush-only: request a roll and do not retain the descriptor. + writer.with_blob_consumer(lambda field_name, descriptor: True) + before = _files(str(table.table_path)) + writer.write_arrow(pa.Table.from_pydict({ + 'id': list(range(5)), + 'blob_data': [b'X' * 2048] * 5, + }, schema=pa_schema)) + messages = writer.prepare_commit() + written = _files(str(table.table_path)) - before + self.assertTrue(any(path.endswith('.blob') for path in written)) + self.assertTrue(any(not path.endswith('.blob') for path in written)) + commit.abort(messages) + remaining = _files(str(table.table_path)) + for path in written: + self.assertNotIn(path, remaining) + finally: + writer.close() + commit.close() + def test_blob_consumer_changes_apply_to_new_partition_writers(self): """Existing writers retain their consumer when the factory changes.""" from pypaimon.common.uri_reader import FileUriReader 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..1b49fa19c431 --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py @@ -0,0 +1,120 @@ +# 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.manifest.schema.data_file_meta import DataFileMeta +from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.file_store_commit import FileStoreCommit + + +def _file(path, extra_files): + return DataFileMeta( + file_name=path.rsplit("/", 1)[-1], + file_size=1, + row_count=0, + min_key=None, + max_key=None, + key_stats=None, + value_stats=None, + min_sequence_number=0, + max_sequence_number=0, + schema_id=0, + level=0, + extra_files=extra_files, + file_path=path, + ) + + +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 = _file( + "/warehouse/table/bucket-0/data.avro", + ["data.avro.blobref", "data.avro.row", "pack.managed.blob"], + ) + message = CommitMessage( + partition=(), + bucket=0, + new_files=[data_file], + ) + + 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_deletes_blob_packs(self): + file_io = Mock() + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io) + + data_file = _file( + "/warehouse/table/bucket-0/data.avro", + ["data.avro.blobref", "pack.managed.blob"], + ) + blob_file = _file("/warehouse/table/bucket-0/payload.blob", []) + message = CommitMessage( + partition=(), + bucket=0, + new_files=[data_file, blob_file], + ) + + 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/pack.managed.blob"), + call("/warehouse/table/bucket-0/payload.blob"), + ], any_order=True) + + 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): + def physical_path(self): + return "/warehouse/table/bucket-0/data.avro" + + def collect_files(self, bucket_path=None): + raise RuntimeError("boom resolving extras") + + message = CommitMessage( + partition=(), + bucket=0, + new_files=[_BoomExtras()], + ) + + 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..9f0641fbbfdb --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_reference_file_test.py @@ -0,0 +1,260 @@ +# 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 +import zlib + +import pyarrow as pa + +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 +from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, MapType +from pypaimon.table.row.blob import BlobDescriptor +from pypaimon.table.row.row_kind import RowKind + + +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) + + def test_reject_crc_valid_malformed_modified_utf8(self): + # Java DataInputStream.readUTF rejects b"\xc0A": the second byte is not + # a 10xxxxxx continuation. The CRC covers that payload, so a loose + # decoder would accept it as U+0001 instead of failing the checksum. + root = b"file:///warehouse" + payload = bytearray() + payload.append(ManagedBlobReferenceFile.VERSION) + payload.extend(struct.pack(">i", 1)) + payload.extend(struct.pack(">H", len(root))) + payload.extend(root) + payload.extend(struct.pack(">H", 2)) + payload.extend(b"\xc0A") + checksum = zlib.crc32(payload) & 0xFFFFFFFF + if checksum >= 0x80000000: + checksum -= 0x100000000 + path = f"{self.temp_dir}/malformed.avro.blobref" + with self.file_io.new_output_stream(path) as out: + out.write(struct.pack(">i", ManagedBlobReferenceFile.MAGIC)) + out.write(payload) + out.write(struct.pack(">i", checksum)) + with self.assertRaisesRegex(IOError, "Invalid modified UTF-8"): + ManagedBlobReferenceFile.read(self.file_io, path) + + def test_collect_table_keeps_add_rows_and_dedupes_packs(self): + blob = AtomicType("BLOB") + fields = [ + DataField(0, "scalar_blob", blob), + DataField(1, "array_blob", ArrayType(True, blob)), + DataField(2, "map_blob", MapType(True, AtomicType("STRING"), blob)), + ] + data_path = f"{self.temp_dir}/data.avro" + collector = ManagedBlobReferenceCollector( + self.file_io, data_path, fields, + {"scalar_blob", "array_blob", "map_blob"}, + ) + + def descriptor(name): + return BlobDescriptor( + "%s/%s.managed.blob" % (self.temp_dir, name), 0, 4).serialize() + + ordinary = BlobDescriptor( + "%s/plain.blob" % self.temp_dir, 0, 4).serialize() + table = pa.table({ + "scalar_blob": [ + descriptor("a"), + descriptor("a"), + descriptor("retract-scalar"), + descriptor("b"), + descriptor("deleted"), + ordinary, + ], + "array_blob": [ + [descriptor("arr"), None], + [descriptor("arr")], + [descriptor("retract-arr")], + [descriptor("arr2")], + [descriptor("deleted-arr")], + None, + ], + "map_blob": [ + {"k": descriptor("map")}, + {"k": descriptor("map"), "again": descriptor("a")}, + {"k": descriptor("retract-map")}, + {"k2": descriptor("map2"), "empty": None}, + {"k": descriptor("deleted-map")}, + {"k": ordinary}, + ], + "_VALUE_KIND": [ + RowKind.INSERT.value, + RowKind.UPDATE_AFTER.value, + RowKind.UPDATE_BEFORE.value, + RowKind.INSERT.value, + RowKind.DELETE.value, + RowKind.INSERT.value, + ], + }) + + collector.collect_table(table) + self.assertEqual(collector.close(), "data.avro.blobref") + self.assertEqual( + ManagedBlobReferenceFile.read( + self.file_io, data_path + ".blobref"), + [ + Reference(self.temp_dir, "a.managed.blob"), + Reference(self.temp_dir, "arr.managed.blob"), + Reference(self.temp_dir, "arr2.managed.blob"), + Reference(self.temp_dir, "b.managed.blob"), + Reference(self.temp_dir, "map.managed.blob"), + Reference(self.temp_dir, "map2.managed.blob"), + ], + ) + + +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..4d985feccf6b --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py @@ -0,0 +1,68 @@ +# 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, then drops them when the parent accepts the commit.""" + + 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 _release_prepared_files(self): + self.committed_files.clear() + self.committed_changelog_files.clear() + + +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_returns_files_and_releases_writer(self): + writer = _KeepingWriter(["data"]) + messages = _file_store(writer).prepare_commit(1) + + self.assertEqual(messages[0].new_files, ["data"]) + self.assertEqual(writer.committed_files, []) + + def test_blob_consumer_prepare_still_releases_files(self): + writer = _KeepingWriter(["pack.blob"]) + messages = _file_store(writer, blob_consumer=object()).prepare_commit(1) + + self.assertEqual(messages[0].new_files, ["pack.blob"]) + self.assertEqual(writer.committed_files, []) + + +if __name__ == "__main__": + unittest.main() From 5e26b63b919315566fa4d2053544fc193bdc9360 Mon Sep 17 00:00:00 2001 From: "wenchao.wu" Date: Wed, 14 Oct 2026 12:04:17 +0800 Subject: [PATCH 2/2] [python] Externalize primary-key BLOBs and keep shared packs reachable. Primary-key writes store payloads in .managed.blob packs and preserve the pack writer URIs in descriptors and reference sidecars for Java reader and GC compatibility. Data files register only their exclusively owned .blobref sidecars in extra_files, keeping shared packs safe from cross-engine abort cleanup. Python abort discovers pack candidates through sidecars and preserves packs when surviving references exist or reference discovery cannot complete. Resolve managed payloads in the reader chain before payload predicates, limits, and output projection. Retain predicate dependencies and resolved BLOB metadata, identify projected MAP values by source field ID, and preserve every batch produced by filtering. Finalize prepare and close by cleaning up unreferenced packs left by failed writes without discarding retry state or deleting packs already handed to successful data files. Add regression coverage for shared-pack cleanup, projected predicates and MAP keys, parallel limits, large batches, mixed BLOB iteration, and failed-write retry, close and abort paths. Cover Python writes followed by Java table reads and pack GC, then Python readback, in the mixed-language test workflow. --- .../java/org/apache/paimon/JavaPyE2ETest.java | 62 ++ paimon-python/dev/run_mixed_tests.sh | 12 +- .../pypaimon/blob/managed_blob_cleanup.py | 132 +++ .../pypaimon/blob/managed_blob_payload.py | 129 +++ .../blob/primary_key_blob_externalizer.py | 282 +++++ .../python/casting/schema_evolution.py | 80 +- .../python/read/reader/auth_masking_reader.py | 52 +- .../reader/blob_descriptor_convert_reader.py | 19 +- .../read/reader/blob_view_read_support.py | 1 + .../read/reader/filter_record_batch_reader.py | 34 +- .../reader/outer_projection_record_reader.py | 9 +- .../execution/python/read/split_read.py | 31 +- .../execution/python/read/table_read.py | 62 +- .../python/write/file_store_commit.py | 6 +- .../python/write/file_store_write.py | 22 +- .../python/write/writer/data_writer.py | 29 +- .../write/writer/dedicated_format_writer.py | 62 +- .../write/writer/key_value_data_writer.py | 117 ++- .../read/reader/iface/record_batch_reader.py | 19 +- .../read/reader/iface/record_reader.py | 2 + .../managed_blob_record_batch_reader.py | 73 ++ paimon-python/pypaimon/schema/data_types.py | 10 + .../pypaimon/schema/schema_manager.py | 151 ++- paimon-python/pypaimon/table/row/blob.py | 14 +- .../pypaimon/table/row/offset_row.py | 11 +- .../tests/auth_masking_reader_test.py | 67 ++ paimon-python/pypaimon/tests/blob_test.py | 61 ++ .../tests/e2e/java_py_read_write_test.py | 56 + .../tests/managed_blob_lifecycle_test.py | 125 ++- .../tests/managed_blob_pk_write_test.py | 987 ++++++++++++++++++ .../managed_blob_write_ownership_test.py | 7 +- .../py36/filter_record_batch_reader_test.py | 38 +- .../tests/schema_evolution_cast_test.py | 34 +- .../pypaimon/tests/schema_manager_test.py | 141 ++- paimon-python/pypaimon/utils/arrow_utils.py | 28 + paimon-python/pypaimon/utils/path.py | 17 + paimon-python/pypaimon/write/row_utils.py | 52 + 37 files changed, 2858 insertions(+), 176 deletions(-) create mode 100644 paimon-python/pypaimon/blob/managed_blob_cleanup.py create mode 100644 paimon-python/pypaimon/blob/managed_blob_payload.py create mode 100644 paimon-python/pypaimon/blob/primary_key_blob_externalizer.py create mode 100644 paimon-python/pypaimon/read/reader/managed_blob_record_batch_reader.py create mode 100644 paimon-python/pypaimon/tests/managed_blob_pk_write_test.py diff --git a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 0d00ebf96b16..c2d0abfa1858 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -59,6 +59,7 @@ import org.apache.paimon.manifest.FileKind; import org.apache.paimon.manifest.IndexManifestEntry; import org.apache.paimon.manifest.ManifestEntry; +import org.apache.paimon.operation.LocalManagedBlobOrphanFilesClean; import org.apache.paimon.options.MemorySize; import org.apache.paimon.options.Options; import org.apache.paimon.predicate.Predicate; @@ -70,6 +71,7 @@ import org.apache.paimon.table.AppendOnlyFileStoreTable; import org.apache.paimon.table.CatalogEnvironment; import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; import org.apache.paimon.table.PrimaryKeyFileStoreTable; import org.apache.paimon.table.Table; import org.apache.paimon.table.sink.BatchTableCommit; @@ -104,6 +106,7 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Paths; +import java.nio.file.attribute.FileTime; import java.time.Instant; import java.time.LocalDate; import java.time.LocalDateTime; @@ -116,6 +119,7 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Collectors; @@ -732,6 +736,64 @@ public void testReadPythonDynamicBucketHashIndex() throws Exception { "hello-java, 42, java-new", "python-only, 7, python-only"); } + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testReadPythonManagedBlobAndClean() throws Exception { + java.nio.file.Path tablePath = + tempDir.resolve("warehouse/default.db/managed_blob_python_to_java"); + FileStoreTable table = + FileStoreTableFactory.create(LocalFileIO.create(), new Path(tablePath.toString())); + assertPythonManagedBlobRows(table); + Long snapshotId = table.snapshotManager().latestSnapshotId(); + + List packs; + try (Stream files = Files.walk(tablePath)) { + packs = + files.filter(path -> path.toString().endsWith(".managed.blob")) + .collect(Collectors.toList()); + } + assertThat(packs).hasSize(3); + long now = System.currentTimeMillis(); + FileTime oldTime = FileTime.fromMillis(now - TimeUnit.HOURS.toMillis(72)); + for (java.nio.file.Path pack : packs) { + Files.setLastModifiedTime(pack, oldTime); + } + java.nio.file.Path orphan = packs.get(0).resolveSibling("orphan.managed.blob"); + Files.copy(packs.get(0), orphan); + Files.setLastModifiedTime(orphan, oldTime); + + try (LocalManagedBlobOrphanFilesClean cleaner = + new LocalManagedBlobOrphanFilesClean( + table, now - TimeUnit.HOURS.toMillis(24), false)) { + assertThat(cleaner.clean().getDeletedFileCount()).isEqualTo(1); + } + assertThat(Files.exists(orphan)).isFalse(); + for (java.nio.file.Path pack : packs) { + assertThat(Files.exists(pack)).isTrue(); + } + assertThat(table.snapshotManager().latestSnapshotId()).isEqualTo(snapshotId); + assertPythonManagedBlobRows(table); + } + + private void assertPythonManagedBlobRows(FileStoreTable table) throws Exception { + List rows = + getResult( + table.newRead(), + table.newScan().plan().splits(), + row -> { + assertThat(row.getBlob(1).toData()) + .isEqualTo("python-payload".getBytes(StandardCharsets.UTF_8)); + assertThat(row.getArray(2).getBlob(0).toData()) + .isEqualTo("python-item".getBytes(StandardCharsets.UTF_8)); + InternalMap map = row.getMap(3); + assertThat(map.keyArray().getString(0).toString()).isEqualTo("k"); + assertThat(map.valueArray().getBlob(0).toData()) + .isEqualTo("python-map".getBytes(StandardCharsets.UTF_8)); + return Integer.toString(row.getInt(0)); + }); + assertThat(rows).containsExactly("1"); + } + private int assignDynamicBucket(FileStoreTable table, InternalRow row) { InternalRowSerializer serializer = new InternalRowSerializer(DataTypes.STRING(), DataTypes.BIGINT()); diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index d95afa77467f..af5cb8f3446e 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -242,7 +242,7 @@ run_python_write_test() { # Run the parameterized Python test method for writing data (pk table, includes bucket num assertion) echo "Running Python test for JavaPyReadWriteTest (test_py_write_read_pk_table)..." - if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_py_write_read_pk_table or test_py_write_read_append_table or test_py_write_dynamic_bucket_hash_index or test_py_write_floating_sequence" -v; then + if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_py_write_read_pk_table or test_py_write_read_append_table or test_py_write_dynamic_bucket_hash_index or test_py_write_floating_sequence or test_py_write_managed_blob" -v; then echo -e "${GREEN}✓ Python write test completed successfully${NC}" return 0 else @@ -261,13 +261,21 @@ run_java_read_test() { echo "Running Maven test for JavaPyE2ETest.testReadPkTable (Java Read Parquet/Orc/Avro)..." echo "Note: Maven may download dependencies on first run, this may take a while..." local parquet_result=0 - if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testReadPkTable+testReadAppendTable+testReadPythonDynamicBucketHashIndex+testReadPythonFloatingSequence -pl paimon-core -Drun.e2e.tests=true -Dpython.version="$PYTHON_VERSION"; then + if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testReadPkTable+testReadAppendTable+testReadPythonDynamicBucketHashIndex+testReadPythonFloatingSequence+testReadPythonManagedBlobAndClean -pl paimon-core -Drun.e2e.tests=true -Dpython.version="$PYTHON_VERSION"; then echo -e "${GREEN}✓ Java read Parquet/Orc/Avro test completed successfully${NC}" else echo -e "${RED}✗ Java read Parquet/Orc/Avro test failed${NC}" parquet_result=1 fi + if [[ $parquet_result -eq 0 ]]; then + cd "$PAIMON_PYTHON_DIR" + if ! python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_read_python_managed_blob_after_java_gc -v; then + parquet_result=1 + fi + cd "$PROJECT_ROOT" + fi + echo "" if [[ $parquet_result -eq 0 ]]; then diff --git a/paimon-python/pypaimon/blob/managed_blob_cleanup.py b/paimon-python/pypaimon/blob/managed_blob_cleanup.py new file mode 100644 index 000000000000..b18aec747a92 --- /dev/null +++ b/paimon-python/pypaimon/blob/managed_blob_cleanup.py @@ -0,0 +1,132 @@ +# 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. + +"""Reference-aware cleanup of shared primary-key BLOB packs.""" + +import logging + +from pypaimon.blob.managed_blob_payload import resolve_table_blob_uri +from pypaimon.blob.managed_blob_reference_file import ManagedBlobReferenceFile +from pypaimon.utils.path import path_identity, to_file_io_path + +logger = logging.getLogger(__name__) + + +def _reference_identities(table, reference): + uri = reference.storage_root_id + '/' + reference.relative_path + # Java Path strings contain literal URI characters. Older Python references + # may instead contain encoded file URIs or table-relative paths. A surviving + # reference must protect every possible interpretation of that path. + paths = {uri, resolve_table_blob_uri(table.table_path, uri)} + return {path_identity(candidate, table.file_io) + for path in paths for candidate in (path, to_file_io_path(path))} + + +def managed_blob_deletion_paths(table, paths): + """Keep shared packs while deleting exclusively owned data files and sidecars.""" + deleting = list(paths) + deleting.extend(_managed_pack_paths(table, deleting)) + protected = _protected_managed_packs(table, deleting) + return [path for path in deleting if path_identity(path, table.file_io) not in protected] + + +def _managed_pack_paths(table, paths): + """Collect colocated pack candidates before deleting aborted sidecars.""" + packs = set() + for path in paths: + path = str(path) + if not path.endswith(ManagedBlobReferenceFile.REFERENCE_FILE_SUFFIX): + continue + try: + references = ManagedBlobReferenceFile.read(table.file_io, path) + parent = path_identity(path, table.file_io).rpartition('/')[0] + for reference in references: + for pack in _reference_identities(table, reference): + if (pack.rpartition('/')[0] == parent + and pack.endswith(ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX)): + packs.add(pack) + except Exception as error: + logger.warning( + "Failed to read managed BLOB references from %s during abort: %s", + path, error) + return packs + + +def _protected_managed_packs(table, paths): + """Keep a pack when a blobref outside this abort still names it. + + One pack can be referenced by several data files. Deleting it + with the first aborted file would make the other files unreadable. + """ + pack_suffix = ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX + ref_suffix = ManagedBlobReferenceFile.REFERENCE_FILE_SUFFIX + deleting = {path_identity(path, table.file_io) for path in paths} + packs_by_dir = {} + for path in deleting: + if not path.endswith(pack_suffix): + continue + parent, _, name = path.rpartition("/") + if parent and name: + packs_by_dir.setdefault(parent, set()).add(path) + if not packs_by_dir: + return set() + + list_status = getattr(table.file_io, "list_status", None) + if list_status is None: + return set(path for packs in packs_by_dir.values() for path in packs) + + protected = set() + for parent, packs in packs_by_dir.items(): + try: + statuses = list(list_status(parent)) + except Exception as error: + logger.warning( + "Failed to list %s while checking managed BLOB packs: %s", + parent, + error, + ) + # Discovery did not finish. Deleting the packs could remove + # payloads that a committed file still references. + protected.update(packs) + continue + for status in statuses: + name = getattr(status, "base_name", None) + if not name: + path = getattr(status, "path", None) or getattr(status, "original_path", None) + name = str(path).replace("\\", "/").rsplit("/", 1)[-1] if path else None + if not name: + protected.update(packs) + break + if not str(name).endswith(ref_suffix): + continue + blobref_path = parent + "/" + str(name) + if blobref_path in deleting: + continue + try: + references = ManagedBlobReferenceFile.read(table.file_io, blobref_path) + referenced = {identity for reference in references + for identity in _reference_identities(table, reference)} + except Exception as error: + logger.warning( + "Failed to read %s while checking managed BLOB packs: %s", + blobref_path, + error, + ) + protected.update(packs) + break + protected.update(packs & referenced) + return protected diff --git a/paimon-python/pypaimon/blob/managed_blob_payload.py b/paimon-python/pypaimon/blob/managed_blob_payload.py new file mode 100644 index 000000000000..7731916cda93 --- /dev/null +++ b/paimon-python/pypaimon/blob/managed_blob_payload.py @@ -0,0 +1,129 @@ +# 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. + +"""Resolve primary-key managed BLOB descriptors to payload bytes.""" + +from pypaimon.blob.managed_blob_reference_collector import _map_items +from pypaimon.schema.data_types import ( + is_array_blob_type, + is_blob_type, + is_map_blob_type, + managed_blob_field_names, +) +from pypaimon.table.row.blob import Blob, BlobDescriptor +from pypaimon.utils.path import resolve_path + + +def resolve_table_blob_uri(table_path, uri): + """Join a table-relative pack URI onto the table directory. + + Absolute paths and URIs that already carry a scheme are left unchanged so + descriptors written with a filesystem path still open. + """ + if uri is None or table_path is None: + return uri + normalized = uri.replace('\\', '/') + if (normalized.startswith('/') + or normalized.startswith('file:') + or '://' in normalized): + return uri + return resolve_path(str(table_path), normalized) + + +def managed_blob_output_specs(table, fields, blob_as_descriptor): + """``(index, name, kind)`` for primary-key managed BLOB columns. + + Append tables store these columns in dedicated ``.blob`` files and the + reader has already turned them into payload bytes. Managed packs exist + only on primary-key tables. + """ + if blob_as_descriptor or not table.is_primary_key_table: + return [] + inline = (table.options.blob_descriptor_fields() + | table.options.blob_view_fields()) + managed = managed_blob_field_names(table.fields, inline) + managed_ids = {field.id for field in table.fields if field.name in managed} + specs = [] + for index, field in enumerate(fields): + # MAP-key projections rename the field but retain its source ID. + if field.id not in managed_ids: + 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 + specs.append((index, field.name, kind)) + return specs + + +def materialize_managed_blob(value, kind, file_io, table_path): + """Replace one managed BLOB cell with its pack payload.""" + if hasattr(value, 'as_py'): + value = value.as_py() + if value is None: + return None + if kind == 'scalar': + return _descriptor_payload(value, file_io, table_path) + if kind == 'array': + return [ + None if element is None + else _descriptor_payload(element, file_io, table_path) + for element in value + ] + if kind == 'map': + entries = _map_entries(value) + externalized = [] + for key, map_value in entries: + if map_value is None: + externalized.append((key, None)) + else: + externalized.append( + (key, _descriptor_payload(map_value, file_io, table_path))) + if isinstance(value, dict): + return dict(externalized) + if (isinstance(value, list) and value + and isinstance(value[0], dict)): + return [{'key': key, 'value': map_value} for key, map_value in externalized] + return externalized + return value + + +def _map_entries(value): + if value is None: + return [] + if isinstance(value, dict): + return list(value.items()) + if isinstance(value, list) and value and isinstance(value[0], dict): + return [(entry.get('key'), entry.get('value')) for entry in value] + return list(_map_items(value)) + + +def _descriptor_payload(raw, file_io, table_path): + if hasattr(raw, 'as_py'): + raw = raw.as_py() + if raw is None: + return None + if isinstance(raw, str): + raw = raw.encode('utf-8') + descriptor = BlobDescriptor.deserialize(bytes(raw)) + uri = resolve_table_blob_uri(table_path, descriptor.uri) + return Blob.from_file( + file_io, uri, descriptor.offset, descriptor.length).to_data() diff --git a/paimon-python/pypaimon/blob/primary_key_blob_externalizer.py b/paimon-python/pypaimon/blob/primary_key_blob_externalizer.py new file mode 100644 index 000000000000..64894c1b2dfc --- /dev/null +++ b/paimon-python/pypaimon/blob/primary_key_blob_externalizer.py @@ -0,0 +1,282 @@ +# 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. + +"""Externalize primary-key BLOB values into managed packs before the data file.""" + +import pyarrow as pa + +from pypaimon.blob.managed_blob_reference_collector import _map_items +from pypaimon.schema.data_types import ( + AtomicType, + DataField, + is_array_blob_type, + is_blob_type, + is_map_blob_type, +) +from pypaimon.table.row.generic_row import GenericRow +from pypaimon.table.row.row_kind import RowKind +from pypaimon.execution.python.write.blob_format_writer import BlobFormatWriter + + +class PrimaryKeyBlobExternalizer: + """Write managed BLOB payloads to ``.managed.blob`` packs and store descriptors. + + Retract rows drop their payload, matching ``PrimaryKeyBlobExternalizer`` in + Java. Packs stay owned by this writer until ``prepare_commit``. + """ + + _RETRACT_KINDS = frozenset({ + RowKind.UPDATE_BEFORE.value, + RowKind.DELETE.value, + }) + + def __init__( + self, + file_io, + new_pack_path, + value_fields, + managed_blob_fields, + target_file_size, + copy_buffer_size, + uri_reader_factory=None): + if target_file_size <= 0: + raise ValueError( + "Managed BLOB target file size must be positive, but was %s." + % target_file_size) + self._file_io = file_io + self._uri_reader_factory = uri_reader_factory + self._uncommitted = [] + self._referenced_pack_names = [] + self._handed_off = set() + self._pending_hand_off = [] + self._writers = {} + self._field_specs = [] + for field in value_fields: + if field.name not in managed_blob_fields: + 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._writers[field.name] = _ManagedBlobPackWriter( + file_io, new_pack_path, target_file_size, copy_buffer_size, + self._uncommitted, self._handed_off) + + def externalize_table(self, data: pa.Table) -> pa.Table: + self._referenced_pack_names = [] + if not self._field_specs: + return data + try: + if "_VALUE_KIND" in data.schema.names: + kinds = data.column("_VALUE_KIND").to_pylist() + else: + kinds = [RowKind.INSERT.value] * data.num_rows + for field_name, field_kind in self._field_specs: + if field_name not in data.schema.names: + continue + # TODO: to_pylist() materializes the whole BLOB column. + values = data.column(field_name).to_pylist() + replaced = [] + for index, value in enumerate(values): + if kinds[index] in self._RETRACT_KINDS: + replaced.append(None) + else: + replaced.append( + self._externalize_value(value, field_kind, field_name)) + field = data.schema.field(field_name) + column_index = data.schema.get_field_index(field_name) + data = data.set_column( + column_index, field, pa.array(replaced, type=field.type)) + return data + except Exception: + self.abort() + raise + + def finish_current(self): + for writer in self._writers.values(): + writer.close_current() + + def prepare_commit(self): + try: + self.finish_current() + # Successful data files already took ownership of their packs. + # A failed write followed by a retry can leave closed, unused packs. + self._delete_uncommitted_packs() + except Exception: + self.abort() + raise + + def take_referenced_pack_names(self): + """Pack file names referenced by the data file just externalized. + + Ownership stays with this writer until ``hand_off_pending``. A failure + while finishing the data file must still be able to delete these packs. + """ + names = self._referenced_pack_names + self._referenced_pack_names = [] + self._pending_hand_off = list(names) + return names + + def hand_off_pending(self): + """The data file that references these packs is now recorded.""" + names = self._pending_hand_off + self._pending_hand_off = [] + self._hand_off_pack_names(names) + + def discard_pending_hand_off(self): + """The data file was not recorded. Keep the packs deletable.""" + self._pending_hand_off = [] + + def _hand_off_pack_names(self, names): + handed = set(names) + if not handed: + return + self._handed_off.update(handed) + self._uncommitted[:] = [ + path for path in self._uncommitted + if _pack_basename(path) not in handed + ] + + def abort(self): + for writer in self._writers.values(): + writer.abort_current() + self._delete_uncommitted_packs() + + def _delete_uncommitted_packs(self): + for path in list(self._uncommitted): + if _pack_basename(path) in self._handed_off: + continue + self._file_io.delete_quietly(path) + self._uncommitted[:] = [ + path for path in self._uncommitted + if _pack_basename(path) in self._handed_off + ] + + def _externalize_value(self, value, field_kind, field_name): + if hasattr(value, "as_py"): + value = value.as_py() + if value is None: + return None + if field_kind == "scalar": + return self._write_descriptor(field_name, value) + if field_kind == "array": + if isinstance(value, (bytes, bytearray, str)) or not hasattr(value, "__iter__"): + raise ValueError( + "ARRAY field '%s' must be a list of blob values." % field_name) + return [ + None if element is None else self._write_descriptor(field_name, element) + for element in value + ] + if field_kind == "map": + externalized = [] + for key, map_value in _map_items(value): + if map_value is None: + externalized.append((key, None)) + else: + externalized.append( + (key, self._write_descriptor(field_name, map_value))) + if isinstance(value, dict): + return dict(externalized) + return externalized + return value + + def _write_descriptor(self, field_name, raw): + blob = BlobFormatWriter._to_blob(raw, self._uri_reader_factory) + descriptor = self._writers[field_name].write(blob) + name = _pack_basename(descriptor.uri) + if name not in self._referenced_pack_names: + self._referenced_pack_names.append(name) + # Java readers and pack GC resolve this URI without a table base. + return descriptor.serialize() + + +def _pack_basename(path): + return str(path).replace('\\', '/').rsplit('/', 1)[-1] + + +class _ManagedBlobPackWriter: + + def __init__(self, file_io, new_pack_path, target_file_size, copy_buffer_size, + uncommitted, handed_off): + self._file_io = file_io + self._new_pack_path = new_pack_path + self._target_file_size = target_file_size + self._copy_buffer_size = copy_buffer_size + self._uncommitted = uncommitted + self._handed_off = handed_off + self._format = None + self._stream = None + self._path = None + self._last = None + self._blob_field = DataField(0, "blob", AtomicType("BLOB")) + + def write(self, blob): + if self._format is None: + self._open() + self._last = None + self._format.add_element(GenericRow([blob], [self._blob_field])) + if self._last is None: + raise IOError("Managed BLOB writer did not produce a descriptor.") + descriptor = self._last + if self._format.reach_target_size(self._target_file_size): + self.close_current() + return descriptor + + def close_current(self): + if self._format is None: + return + try: + self._format.close() + finally: + self._format = None + self._stream = None + self._path = None + + def abort_current(self): + path = self._path + stream = self._stream + self._format = None + self._stream = None + self._path = None + if stream is not None: + try: + stream.close() + except Exception: + pass + if path is not None and _pack_basename(path) not in self._handed_off: + self._file_io.delete_quietly(path) + + def _open(self): + path = self._new_pack_path() + self._path = path + self._uncommitted.append(path) + self._stream = self._file_io.new_output_stream(path) + self._format = BlobFormatWriter( + self._stream, + blob_consumer=self._consume, + file_path=path, + copy_buffer_size=self._copy_buffer_size, + ) + + def _consume(self, field_name, descriptor): + self._last = descriptor + return False diff --git a/paimon-python/pypaimon/execution/python/casting/schema_evolution.py b/paimon-python/pypaimon/execution/python/casting/schema_evolution.py index b14f28e8cb50..2348c01363c5 100644 --- a/paimon-python/pypaimon/execution/python/casting/schema_evolution.py +++ b/paimon-python/pypaimon/execution/python/casting/schema_evolution.py @@ -156,4 +156,82 @@ def cast_array_for_schema_evolution(array, target_type, source_paimon_type=None, if (pa.types.is_signed_integer(target_type) and (pa.types.is_floating(array.type) or pa.types.is_decimal(array.type))): return _numeric_to_integer(array, target_type) - return array.cast(target_type, safe=False) + try: + return array.cast(target_type, safe=False) + except pa.ArrowNotImplementedError: + # PyArrow 6 cannot cast a map or list as a whole when only the child + # binary type changes. Child arrays themselves still cast. + rebuilt = _cast_nested_container(array, target_type) + if rebuilt is None: + raise + return rebuilt + + +def _cast_child(array, target_type): + if array.type == target_type: + return array + try: + return array.cast(target_type, safe=False) + except pa.ArrowNotImplementedError: + return _cast_nested_container(array, target_type) + + +def _cast_nested_container(array, target_type): + """Rebuild a list or map after casting its children.""" + if isinstance(array, pa.ChunkedArray): + chunks = [_cast_nested_container(chunk, target_type) for chunk in array.chunks] + if any(chunk is None for chunk in chunks): + return None + return pa.chunked_array(chunks, type=target_type) + if pa.types.is_list(array.type) and pa.types.is_list(target_type): + values = _cast_child(array.values, target_type.value_type) + if values is None: + return None + return pa.ListArray.from_arrays(_offsets_with_nulls(array), values) + if pa.types.is_large_list(array.type) and pa.types.is_large_list(target_type): + values = _cast_child(array.values, target_type.value_type) + if values is None: + return None + return pa.LargeListArray.from_arrays(_offsets_with_nulls(array), values) + if pa.types.is_map(array.type) and pa.types.is_map(target_type): + keys = _cast_child(array.keys, target_type.key_type) + items = _cast_child(array.items, target_type.item_type) + if keys is None or items is None: + return None + # PyArrow 6 MapArray.from_arrays has no mask argument. A null offset + # entry is how that version represents a null map. + return pa.MapArray.from_arrays(_offsets_with_nulls(array), keys, items) + return None + + +def _value_offsets(array): + """Offsets for a list or map, including a sliced array. + + PyArrow 6's ``MapArray`` has no ``offsets`` attribute. The offset buffer is + still ``buffers()[1]``, and a slice starts at ``array.offset``. Copy that + range into a new array: ``from_arrays`` rejects a sliced offset buffer when + the map also has a null bitmap. The values still index the child arrays. + """ + return _offsets_from_buffer(array) + + +def _offsets_from_buffer(array): + offset_type = pa.int64() if pa.types.is_large_list(array.type) else pa.int32() + width = 8 if offset_type == pa.int64() else 4 + start = array.offset + count = len(array) + 1 + raw = array.buffers()[1] + sliced = None if raw is None else raw.slice(start * width, count * width) + return pa.Array.from_buffers(offset_type, count, [None, sliced]) + + +def _offsets_with_nulls(array): + offsets = _value_offsets(array) + if array.null_count == 0: + return offsets + # Keep the numeric offset under each null so neighboring rows still slice + # the child arrays at the right index. PyArrow 6 reads that null bit as a + # null list or map. + nulls = array.is_null().to_pylist() + nulls.append(False) + return pa.array(offsets.to_pylist(), type=offsets.type, mask=pa.array(nulls)) diff --git a/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py b/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py index 6ad5df55f094..03d78b3b3577 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/auth_masking_reader.py @@ -27,23 +27,19 @@ from pypaimon.read.reader.iface.record_iterator import RecordIterator from pypaimon.read.reader.iface.record_reader import RecordReader from pypaimon.table.row.offset_row import OffsetRow -from pypaimon.utils.arrow_utils import zero_column_batch +from pypaimon.utils.arrow_utils import row_kind_column_index, row_kind_field, with_no_row_kind, zero_column_batch class RecordReaderToBatchAdapter(RecordBatchReader): def __init__(self, inner, schema: pa.Schema, chunk_size: int = 65536, include_row_kind: bool = False): self._inner = inner - self._schema = schema + self._schema = schema if include_row_kind else with_no_row_kind(schema) self._chunk_size = chunk_size self._exhausted = False self._pending_iterator = None self._include_row_kind = include_row_kind - self.file_io = getattr(inner, 'file_io', None) - self.blob_field_indices = getattr(inner, 'blob_field_indices', None) - self.descriptor_field_indices = getattr(inner, 'descriptor_field_indices', None) - self.blob_view_lookup = getattr(inner, 'blob_view_lookup', None) - self.vector_field_indices = getattr(inner, 'vector_field_indices', None) + self._adopt_blob_metadata(inner) def read_arrow_batch(self) -> Optional[pa.RecordBatch]: if self._exhausted: @@ -84,8 +80,8 @@ def _flush(self, row_tuples, row_kinds=None): else zero_column_batch(len(row_tuples), self._schema.metadata)) if row_kinds: row_kind_array = pa.array(row_kinds, type=pa.string()) - row_kind_field = pa.field("_row_kind", pa.string()) - new_schema = pa.schema([row_kind_field] + list(batch.schema)) + new_schema = pa.schema([row_kind_field(batch.schema)] + list(batch.schema), + metadata=batch.schema.metadata) columns = [row_kind_array] + [batch.column(i) for i in range(batch.num_columns)] batch = pa.RecordBatch.from_arrays(columns, schema=new_schema) return batch @@ -99,11 +95,7 @@ class BatchToRecordReaderAdapter(RecordReader): def __init__(self, inner: RecordBatchReader): self._inner = inner - self.file_io = getattr(inner, 'file_io', None) - self.blob_field_indices = getattr(inner, 'blob_field_indices', None) - self.descriptor_field_indices = getattr(inner, 'descriptor_field_indices', None) - self.blob_view_lookup = getattr(inner, 'blob_view_lookup', None) - self.vector_field_indices = getattr(inner, 'vector_field_indices', None) + self._adopt_blob_metadata(inner) def read_batch(self): batch = self._inner.read_arrow_batch() @@ -117,6 +109,8 @@ def read_batch(self): descriptor_field_indices=self.descriptor_field_indices, blob_view_lookup=self.blob_view_lookup, vector_field_indices=self.vector_field_indices, + resolved_blob_field_indices=self.resolved_blob_field_indices, + blob_base_path=self.blob_base_path, ) def close(self): @@ -128,12 +122,13 @@ class _ArrowBatchIterator(RecordIterator): def __init__(self, batch: pa.RecordBatch, file_io=None, blob_field_indices=None, descriptor_field_indices=None, blob_view_lookup=None, - vector_field_indices=None): + vector_field_indices=None, resolved_blob_field_indices=None, + blob_base_path=None): self._batch = batch self._idx = 0 - self._has_rk = "_row_kind" in batch.schema.names + self._rk_idx = row_kind_column_index(batch.schema) + self._has_rk = self._rk_idx >= 0 if self._has_rk: - self._rk_idx = batch.schema.get_field_index("_row_kind") self._data_cols = [j for j in range(batch.num_columns) if j != self._rk_idx] else: self._rk_idx = -1 @@ -143,6 +138,8 @@ def __init__(self, batch: pa.RecordBatch, self._descriptor_field_indices = descriptor_field_indices self._blob_view_lookup = blob_view_lookup self._vector_field_indices = vector_field_indices + self._resolved_blob_field_indices = resolved_blob_field_indices + self._blob_base_path = blob_base_path def next(self): if self._idx >= self._batch.num_rows: @@ -157,7 +154,9 @@ def next(self): blob_field_indices=self._blob_field_indices, descriptor_field_indices=self._descriptor_field_indices, blob_view_lookup=self._blob_view_lookup, - vector_field_indices=self._vector_field_indices) + vector_field_indices=self._vector_field_indices, + resolved_blob_field_indices=self._resolved_blob_field_indices, + blob_base_path=self._blob_base_path) if self._has_rk: from pypaimon.table.row.row_kind import RowKind kind_str = self._batch.column(self._rk_idx)[self._idx].as_py() @@ -249,10 +248,18 @@ def close(self): class ColumnProjectReader(RecordBatchReader): - def __init__(self, inner_reader: RecordBatchReader, columns: List[str]): + def __init__(self, inner_reader: RecordBatchReader, columns: List[str], source_fields=None): self._inner = inner_reader self._columns = columns self._adopt_metadata(inner_reader) + if source_fields is not None: + source_indices = {field.name: i for i, field in enumerate(source_fields)} + for attr in ('blob_field_indices', 'descriptor_field_indices', + 'resolved_blob_field_indices', 'vector_field_indices'): + indices = getattr(self, attr) + if indices is not None: + setattr(self, attr, {i for i, name in enumerate(columns) + if source_indices[name] in indices}) def read_arrow_batch(self) -> Optional[pa.RecordBatch]: batch = self._inner.read_arrow_batch() @@ -260,8 +267,11 @@ def read_arrow_batch(self) -> Optional[pa.RecordBatch]: return None self._refresh_blob_view_lookup(self._inner) columns = self._columns - if "_row_kind" in batch.schema.names and "_row_kind" not in columns: - columns = ["_row_kind"] + list(columns) + row_kind_index = row_kind_column_index(batch.schema) + if row_kind_index >= 0: + name = batch.schema.field(row_kind_index).name + if name not in columns: + columns = [name] + list(columns) if not columns: return zero_column_batch(batch.num_rows, batch.schema.metadata) return pa.RecordBatch.from_arrays( diff --git a/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py b/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py index 188119abd31e..88b044b2d6ce 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/blob_descriptor_convert_reader.py @@ -23,7 +23,7 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.table.row.blob import Blob, BlobViewStruct -from pypaimon.utils.arrow_utils import as_blob_data, clear_blob_data, is_blob_data +from pypaimon.utils.arrow_utils import as_blob_data, clear_blob_data, is_blob_data, row_kind_column_index class BlobDataBatchReader(RecordBatchReader): @@ -68,7 +68,7 @@ class BlobInlineConvertReader(RecordBatchReader): def __init__(self, inner: RecordBatchReader, table, prescan_reader_factory: Optional[Callable[[Set[str]], RecordBatchReader]] = None, - blob_parallelism: int = 1): + blob_parallelism: int = 1, field_names=None): """ Args: inner: The main data reader (reads all columns). @@ -78,6 +78,8 @@ def __init__(self, inner: RecordBatchReader, table, prescan to collect BlobViewStructs without reading all columns. Signature: (field_names: Set[str]) -> RecordBatchReader blob_parallelism: number of threads for concurrent blob reads. + field_names: logical reader fields, so wrappers can adopt payload + provenance before the first batch is read. """ self._inner = inner self._table = table @@ -91,13 +93,22 @@ def __init__(self, inner: RecordBatchReader, table, self._view_fields = CoreOptions.blob_view_fields(table.options) if resolve_enabled else set() self._descriptor_fields = CoreOptions.blob_descriptor_fields(table.options) self._blob_as_descriptor = CoreOptions.blob_as_descriptor(table.options) + self._field_names = field_names if not self._blob_as_descriptor: # Stage 2 materializes descriptor/view fields to payload bytes. # Row-level descriptor routing must not re-parse that content. self.descriptor_field_indices = set() + if field_names is not None: + self._mark_resolved_fields(field_names) self._prescan_done = False self._blob_view_lookup = None + def _mark_resolved_fields(self, field_names): + resolved_names = self._descriptor_fields | self._view_fields + self.resolved_blob_field_indices = set(self.resolved_blob_field_indices or ()) | { + index for index, name in enumerate(field_names) + if name in resolved_names} + def read_arrow_batch(self) -> Optional[RecordBatch]: # Align with Java: only enter blob view resolution when catalog_loader is available. if self._view_fields and not self._prescan_done: @@ -106,6 +117,10 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: batch = self._inner.read_arrow_batch() if batch is None: return None + if not self._blob_as_descriptor and self._field_names is None: + self._mark_resolved_fields([ + name for index, name in enumerate(batch.schema.names) + if index != row_kind_column_index(batch.schema)]) # Resolve view fields using the preloaded lookup. When # blob-as-descriptor=true, leave BlobViewStruct bytes in place so # get_blob() keeps the originating table's FileIO. Arrow paths diff --git a/paimon-python/pypaimon/execution/python/read/reader/blob_view_read_support.py b/paimon-python/pypaimon/execution/python/read/reader/blob_view_read_support.py index 54055b592d1c..59e36395a951 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/blob_view_read_support.py +++ b/paimon-python/pypaimon/execution/python/read/reader/blob_view_read_support.py @@ -64,5 +64,6 @@ def wrap_record_reader_with_blob_inline_convert( split_read.table, prescan_reader_factory=lambda names: split_read._create_blob_view_prescan_reader(names), blob_parallelism=split_read._blob_parallelism, + field_names=[field.name for field in read_fields], ) return BatchToRecordReaderAdapter(batch_reader) diff --git a/paimon-python/pypaimon/execution/python/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/execution/python/read/reader/filter_record_batch_reader.py index 2dccc1c995bc..56de48b046bb 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/filter_record_batch_reader.py @@ -31,10 +31,27 @@ _polars_arrow_table, ) from pypaimon.schema.data_types import DataField, PyarrowFieldParser +from pypaimon.utils.arrow_utils import row_kind_column_index logger = logging.getLogger(__name__) +def _record_batch_from_table(result): + """Return every filtered row, including when the scanner emits several batches.""" + batches = result.to_batches() + if not batches: + return None + if len(batches) == 1: + return batches[0] + concat_batches = getattr(pa, "concat_batches", None) + if concat_batches is not None: + return concat_batches(batches) + return pa.RecordBatch.from_arrays( + [result.column(i).combine_chunks() for i in range(result.num_columns)], + schema=result.schema, + ) + + def _comparison_fields(predicate): if predicate.method in ('and', 'or'): return set().union(*(_comparison_fields(child) for child in predicate.literals)) @@ -142,18 +159,7 @@ def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: ).to_table() if result.num_rows == 0: return None - batches = result.to_batches() - if not batches: - return None - if len(batches) == 1: - return batches[0] - concat_batches = getattr(pa, "concat_batches", None) - if concat_batches is not None: - return concat_batches(batches) - return pa.RecordBatch.from_arrays( - [result.column(i) for i in range(result.num_columns)], - schema=result.schema, - ) + return _record_batch_from_table(result) def _coerce_comparison_types(self, view, batch): # Cast columns a mask retyped, for the filter only; the output is unchanged. @@ -177,7 +183,8 @@ def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Opti # may prepend _row_kind, which is metadata rather than an operand. columns = (self.field_names if self.field_names is not None else [field.name for field in self.schema_fields] if self.schema_fields is not None else - [name for name in filter_batch.schema.names if name != "_row_kind"]) + [name for index, name in enumerate(filter_batch.schema.names) + if index != row_kind_column_index(filter_batch.schema)]) df = polars.from_arrow(_polars_arrow_table(filter_batch).select(columns)) iterator = InternalRowWrapperIterator( self._iter_df_rows(df), @@ -187,6 +194,7 @@ def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Opti self.vector_field_indices, self.descriptor_field_indices, self.blob_view_lookup, + resolved_blob_field_indices=self.resolved_blob_field_indices, ) selected = [] pos = 0 diff --git a/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py b/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py index fa4d71091787..edca4c7f3df3 100644 --- a/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py +++ b/paimon-python/pypaimon/execution/python/read/reader/outer_projection_record_reader.py @@ -70,6 +70,8 @@ def __init__( self.file_io = self._file_io self.blob_field_indices = self._blob_field_indices self.descriptor_field_indices = self._descriptor_field_indices + self.resolved_blob_field_indices = _project_top_level_field_indices( + getattr(inner, 'resolved_blob_field_indices', None), self._specs) self.vector_field_indices = self._vector_field_indices self.blob_view_lookup = getattr(inner, 'blob_view_lookup', None) @@ -82,7 +84,8 @@ def read_batch(self) -> Optional[RecordIterator[InternalRow]]: inner_batch, self._specs, self._flat_arity, self._file_io, self._blob_field_indices, self._vector_field_indices, self._descriptor_field_indices, - blob_view_lookup=self.blob_view_lookup) + blob_view_lookup=self.blob_view_lookup, + resolved_blob_field_indices=self.resolved_blob_field_indices) def close(self) -> None: self._inner.close() @@ -101,6 +104,7 @@ def __init__( vector_field_indices=None, descriptor_field_indices=None, blob_view_lookup=None, + resolved_blob_field_indices=None, ): self._inner = inner self._specs = specs @@ -110,7 +114,8 @@ def __init__( blob_field_indices=blob_field_indices, vector_field_indices=vector_field_indices, descriptor_field_indices=descriptor_field_indices, - blob_view_lookup=blob_view_lookup) + blob_view_lookup=blob_view_lookup, + resolved_blob_field_indices=resolved_blob_field_indices) def next(self) -> Optional[InternalRow]: inner_row = self._inner.next() diff --git a/paimon-python/pypaimon/execution/python/read/split_read.py b/paimon-python/pypaimon/execution/python/read/split_read.py index 3d4b1be412fb..d86122f9a214 100644 --- a/paimon-python/pypaimon/execution/python/read/split_read.py +++ b/paimon-python/pypaimon/execution/python/read/split_read.py @@ -40,6 +40,7 @@ from pypaimon.execution.python.read.interval_partition import IntervalPartition, SortedRun from pypaimon.read.partition_info import PartitionInfo from pypaimon.read.push_down_utils import ( + exclude_predicate_with_fields, predicate_field_names, predicate_for_physical_fields, predicate_matches_physical_names, @@ -166,18 +167,37 @@ def __init__( # len(read_type) and whose coordinate space is read_type — that is # the space FilterRecordReader actually evaluates against. read_type_names = {f.name for f in read_type} + reader_predicate = self._predicate_without_managed_blobs(self.predicate) if ( - self.predicate is not None - and predicate_field_names(self.predicate).issubset(read_type_names) - and predicate_matches_physical_names(self.predicate) + reader_predicate is not None + and predicate_field_names(reader_predicate).issubset(read_type_names) + and predicate_matches_physical_names(reader_predicate) ): self.predicate_for_reader = rewrite_predicate_indices( - self.predicate, read_type + reader_predicate, read_type ) else: self.predicate_for_reader = None self._blob_view_prescan = False + def _predicate_without_managed_blobs(self, predicate): + """Drop predicates on primary-key managed BLOB columns. + + Those columns store descriptors. Comparing them with payload bytes + drops rows that would match after the descriptor is resolved. + """ + if predicate is None or not self.table.is_primary_key_table: + return predicate + if CoreOptions.blob_as_descriptor(self.table.options): + return predicate + from pypaimon.schema.data_types import managed_blob_field_names + inline = (CoreOptions.blob_descriptor_fields(self.table.options) + | CoreOptions.blob_view_fields(self.table.options)) + managed = managed_blob_field_names(self.table.fields, inline) + if not managed: + return predicate + return exclude_predicate_with_fields(predicate, managed) + def _needs_blob_inline_convert(self) -> bool: return needs_blob_inline_convert(self.table) @@ -190,6 +210,7 @@ def _wrap_batch_reader_with_blob_inline_convert( self.table, prescan_reader_factory=lambda names: self._create_blob_view_prescan_reader(names), blob_parallelism=self._blob_parallelism, + field_names=[field.name for field in self.read_fields], ) def _blob_view_prescan_limit(self) -> Optional[int]: @@ -1053,6 +1074,7 @@ def create_reader(self) -> RecordReader: descriptor_field_indices=descriptor_field_indices_for_table( self.table, self.read_fields), vector_field_indices=vector_field_indices(self.read_fields)) + concat_reader.blob_base_path = str(self.table.table_path) reader = concat_reader if (self.predicate_for_reader and (self.table.is_primary_key_table @@ -1452,6 +1474,7 @@ def _create_raw_reader(self) -> RecordReader: descriptor_field_indices=descriptor_field_indices_for_table( self.table, self.read_fields), vector_field_indices=vector_field_indices(self.read_fields)) + merge_reader.blob_base_path = str(self.table.table_path) if self.predicate_for_reader is not None: reader = FilterRecordBatchReader( merge_reader, diff --git a/paimon-python/pypaimon/execution/python/read/table_read.py b/paimon-python/pypaimon/execution/python/read/table_read.py index 0802587eb5a2..ac8b496a5210 100644 --- a/paimon-python/pypaimon/execution/python/read/table_read.py +++ b/paimon-python/pypaimon/execution/python/read/table_read.py @@ -38,6 +38,7 @@ from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.execution.python.read.reader.limited_record_reader import LimitedRecordBatchReader from pypaimon.execution.python.read.reader.filter_record_batch_reader import FilterRecordBatchReader +from pypaimon.read.reader.managed_blob_record_batch_reader import ManagedBlobRecordBatchReader from pypaimon.read.split import Split from pypaimon.read.read_type import adapter_output_paths from pypaimon.read.read_support import close_read_source @@ -691,6 +692,56 @@ def _map_has_aggregator(self, field_name: str) -> bool: or "fields.default-aggregate-function" in options ) + def _managed_blob_specs(self, fields): + from pypaimon.blob.managed_blob_payload import managed_blob_output_specs + read = self._read + return managed_blob_output_specs( + read.table, fields, read._blob_as_descriptor) + + def _resolve_managed_blob_reader(self, reader, fields): + specs = self._managed_blob_specs(fields) + if not specs: + return reader + read = self._read + if not isinstance(reader, RecordBatchReader): + reader = RecordReaderToBatchAdapter( + reader, read_arrow_schema(fields), include_row_kind=True) + from pypaimon.read.reader.field_indices import ( + blob_field_indices, descriptor_field_indices_for_table, vector_field_indices) + if reader.blob_field_indices is None: + reader.blob_field_indices = blob_field_indices(fields) + # An empty set means the inner reader already resolved descriptors. + if reader.descriptor_field_indices is None: + reader.descriptor_field_indices = descriptor_field_indices_for_table( + read.table, fields) + if reader.vector_field_indices is None: + reader.vector_field_indices = vector_field_indices(fields) + return ManagedBlobRecordBatchReader( + reader, specs, read.table.file_io, read.table.table_path) + + def _managed_blob_reader(self, split, blob_parallelism, limit): + read = self._read + fields = read._scan_read_type + managed_fields = {name for _index, name, _kind in self._managed_blob_specs(fields)} + post_filter = ( + bool(predicate_field_names(read.predicate) & managed_fields) + if read.predicate else False) + reader = self._create_split_read( + split, blob_parallelism=blob_parallelism, read_type=fields, + limit=limit, push_down_limit=not post_filter, + predicate=exclude_predicate_with_fields(read.predicate, managed_fields), + ).create_reader() + needs_convert_back = not isinstance(reader, RecordBatchReader) + reader = self._resolve_managed_blob_reader(reader, fields) + if post_filter: + reader = FilterRecordBatchReader( + reader, rewrite_predicate_indices(read.predicate, fields), schema_fields=fields) + if read._predicate_extra_fields: + reader = ColumnProjectReader(reader, read._output_column_names, fields) + if needs_convert_back or not read._adapter_read_type: + reader = BatchToRecordReaderAdapter(reader) + return reader + def _create_reader_for_split( self, split, blob_parallelism=1, limit: Optional[int] = None): # Raw-convertible splits skip MergeFileSplitRead, so validate before @@ -705,6 +756,8 @@ def _create_reader_for_split( if auth_result is not None: return self._authed_reader( split, auth_result, blob_parallelism, limit) + if self._managed_blob_specs(read._scan_read_type): + return self._managed_blob_reader(split, blob_parallelism, limit) if limit is None: return self._create_split_read( split, blob_parallelism=blob_parallelism).create_reader() @@ -741,6 +794,8 @@ def _authed_reader( effective_read_type = filter_read_fields + extra_fields filter_fn = auth_result.extract_row_filter() + managed_fields = { + name for _index, name, _kind in self._managed_blob_specs(effective_read_type)} effective_limit = read.limit if limit is None else limit auth_fields = ( self._auth_filter_field_names(auth_result, effective_read_type) @@ -776,8 +831,8 @@ def _authed_reader( post_merge_filter=filter_fn if embed_filter else None, eager_blob_fields=eager_blob_fields, post_filter_after_inline=post_filter_after_inline, - # push down conjuncts without masked columns - predicate=exclude_predicate_with_fields(post_filter, masked_fields), + # Managed BLOB predicates need resolved payloads, and masks apply after the read. + predicate=exclude_predicate_with_fields(post_filter, masked_fields | managed_fields), defer_blob_resolution=post_filter is not None, ) reader = split_read.create_reader() @@ -797,6 +852,7 @@ def _authed_reader( reader.vector_field_indices = vector_field_indices(effective_read_type) needs_convert_back = True + reader = self._resolve_managed_blob_reader(reader, effective_read_type) if filter_fn and not embed_filter: reader = AuthFilterReader(reader, filter_fn) @@ -823,7 +879,7 @@ def _authed_reader( if len(effective_read_type) > len(read_fields): original_columns = [f.name for f in read_fields] - reader = ColumnProjectReader(reader, original_columns) + reader = ColumnProjectReader(reader, original_columns, effective_read_type) if needs_convert_back: reader = BatchToRecordReaderAdapter(reader) diff --git a/paimon-python/pypaimon/execution/python/write/file_store_commit.py b/paimon-python/pypaimon/execution/python/write/file_store_commit.py index ebebf63d1bce..05d7c71630f2 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_commit.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_commit.py @@ -21,6 +21,7 @@ import uuid from typing import Dict, List, Optional +from pypaimon.blob.managed_blob_cleanup import managed_blob_deletion_paths from pypaimon.build_info import full_version as build_full_version from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.predicate_builder import PredicateBuilder @@ -980,6 +981,7 @@ def abort(self, commit_messages: List[CommitMessage]): Never call after a commit whose outcome is unknown. Internal error paths must preserve prepared files instead of calling this method. """ + collected = [] for message in commit_messages: for file in (list(message.new_files) + list(message.changelog_files) + list(message.compact_after) @@ -1006,7 +1008,7 @@ def abort(self, commit_messages: List[CommitMessage]): error, ) continue - _delete_abort_paths(self.table, paths) + collected.append(paths) for entry in message.index_adds + message.compact_index_adds: file_name = None try: @@ -1029,6 +1031,8 @@ def abort(self, commit_messages: List[CommitMessage]): file_name, error, ) + deleting = [path for paths in collected for path in paths] + _delete_abort_paths(self.table, managed_blob_deletion_paths(self.table, deleting)) def close(self): """Close the FileStoreCommit and release resources.""" diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index af5bb5a127d4..a0714ec8bcf2 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -127,9 +127,12 @@ def write_row( require_columns(values_by_name, column_names, 'write_row') writer = self.data_writers.get(key) data = None - if writer is None and self._has_blob_columns(): + if (writer is None and self._has_blob_columns() + and not self.table.is_primary_key_table): # Dedicated writers consume live payloads themselves. Inline fields # instead encode to Arrow before the writer captures its layout. + # Primary-key Blob tables use KeyValueDataWriter, which consumes + # an Arrow batch like every other key-value writer. descriptor_fields = self.options.blob_descriptor_fields() view_fields = self.options.blob_view_fields() inline_columns = [name for name in column_names if name in descriptor_fields or name in view_fields] @@ -220,13 +223,10 @@ def max_seq_number(): default = -1 if self.table.is_primary_key_table else 1 return self._seq_number_stats(partition).get(bucket, default) - # Dedicated Blob files are an append-table layout. PK tables require - # managed packs and references attached to their key-value Parquet files. - if self._has_blob_columns(): - if self.table.is_primary_key_table: - raise NotImplementedError( - 'Primary-key Blob writes require the native writer; ' - 'enable write.native.enabled and write Arrow batches') + # Dedicated Blob files are an append-table layout. Primary-key tables + # store descriptors in the key-value file and list packs in a .blobref + # sidecar written by KeyValueDataWriter. + if self._has_blob_columns() and not self.table.is_primary_key_table: return DedicatedFormatWriter( table=self.table, partition=partition, @@ -238,7 +238,8 @@ def max_seq_number(): changelog_producer=self.changelog_producer, blob_uri_reader_factory=self.blob_uri_reader_factory, ) - elif self._has_vector_columns() and options.with_vector_format(): + elif (self._has_vector_columns() and options.with_vector_format() + and not (self.table.is_primary_key_table and self._has_blob_columns())): return DataVectorWriter( table=self.table, partition=partition, @@ -257,7 +258,8 @@ def max_seq_number(): max_seq_number=(0 if bucket == BucketMode.POSTPONE_BUCKET.value else max_seq_number()), options=options, merge_function=self._build_pk_merge_function(), - changelog_producer=self.changelog_producer) + changelog_producer=self.changelog_producer, + blob_uri_reader_factory=self.blob_uri_reader_factory) else: seq_number = 0 if self.table.bucket_mode() == BucketMode.BUCKET_UNAWARE else max_seq_number() return AppendOnlyDataWriter( diff --git a/paimon-python/pypaimon/execution/python/write/writer/data_writer.py b/paimon-python/pypaimon/execution/python/write/writer/data_writer.py index edabf75404d0..c7e5e6c2a468 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/data_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/data_writer.py @@ -288,15 +288,16 @@ def _write_data_to_file(self, data: pa.Table): logical_data = data extra_files = [] row_sidecar_path = None + managed_reference_path = None changelog_meta = None shared_shredding_stats = {} if self._variant_shredding: data = self._apply_variant_shredding(data) - # One data file means up to three files on disk -- the data file, its row - # sidecar and its changelog -- and none of them is committed until all of - # them have landed. A caller that retries the flush still holds these - # rows in its buffer, so publishing the data file before the changelog + # One data file means the data file plus its row sidecar, managed BLOB + # reference and changelog. None of them is committed until all of them + # have landed. A caller that retries the flush still holds these rows + # in its buffer, so publishing the data file before the changelog # exists would make the retry write a second copy of rows the first meta # already covers. try: @@ -313,6 +314,12 @@ def _write_data_to_file(self, data: pa.Table): zstd_level=self.zstd_level) extra_files.append(row_sidecar_name) + additional_extra_files = self._additional_extra_files(file_path, logical_data) + if additional_extra_files: + extra_files.extend(additional_extra_files) + managed_reference_path = self._managed_reference_path( + file_path, additional_extra_files) + statistics = self._data_file_statistics(logical_data) data_meta = self._create_data_file_meta( file_name=file_name, @@ -336,6 +343,8 @@ def _write_data_to_file(self, data: pa.Table): self.file_io.delete_quietly(file_path) if row_sidecar_path is not None: self.file_io.delete_quietly(row_sidecar_path) + if managed_reference_path is not None: + self.file_io.delete_quietly(managed_reference_path) raise self._finish_data_file(data_meta, changelog_meta, shared_shredding_stats) @@ -524,6 +533,18 @@ def _generate_file_path(self, file_name: str) -> str: # Python writer path. Never introduce another Python-specific layout. return canonical_data_file_path(self.table, self.partition, self.bucket, file_name) + def _additional_extra_files(self, file_path: str, data: pa.Table) -> List[str]: + """Sidecars written beside one data file. Empty for writers that have none.""" + return [] + + @staticmethod + def _managed_reference_path(file_path: str, extra_files: List[str]): + from pypaimon.blob.managed_blob_reference_file import ManagedBlobReferenceFile + if any(name.endswith(ManagedBlobReferenceFile.REFERENCE_FILE_SUFFIX) + for name in extra_files): + return ManagedBlobReferenceFile.sidecar_path(file_path) + return None + def _should_write_row_sidecar(self) -> bool: return ( self.options.data_evolution_enabled(False) diff --git a/paimon-python/pypaimon/execution/python/write/writer/dedicated_format_writer.py b/paimon-python/pypaimon/execution/python/write/writer/dedicated_format_writer.py index 15fcd3691109..428d631e74aa 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/dedicated_format_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/dedicated_format_writer.py @@ -38,6 +38,7 @@ require_columns, row_to_named_values, row_values_to_arrow_table, + validate_inline_blob_fields, ) from pypaimon.execution.python.write.writer.composite_data_writer import CompositeDataWriter from pypaimon.execution.python.write.writer.video_group import VideoGroupRollingPolicy @@ -351,65 +352,8 @@ def _project_columns(data: pa.RecordBatch, column_names: List[str]) -> pa.Record return project_arrow_batch(data, column_names) def _validate_inline_stored_fields_input(self, data: pa.RecordBatch): - if not self.blob_inline_fields: - return - - from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct - - for field_name in self.blob_descriptor_fields: - if field_name not in data.schema.names: - continue - values = data.column(data.schema.get_field_index(field_name)).to_pylist() - for value in values: - if value is None: - continue - if hasattr(value, 'as_py'): - value = value.as_py() - if isinstance(value, str): - value = value.encode('utf-8') - if not isinstance(value, (bytes, bytearray)): - raise ValueError( - "blob-descriptor-field requires blob field value to be a serialized " - "BlobDescriptor." - ) - descriptor_bytes = bytes(value) - # Like Java, schema-declared descriptors use the prefix parser. - # Exact wire length is only needed when detecting descriptors - # among arbitrary payload bytes. - try: - BlobDescriptor.deserialize(descriptor_bytes) - except Exception as e: - raise ValueError( - "blob-descriptor-field requires blob field value to be a serialized " - "BlobDescriptor." - ) from e - - for field_name in self.blob_view_fields: - if field_name not in data.schema.names: - continue - values = data.column(data.schema.get_field_index(field_name)).to_pylist() - for value in values: - if value is None: - continue - if hasattr(value, 'as_py'): - value = value.as_py() - if isinstance(value, str): - value = value.encode('utf-8') - if not isinstance(value, (bytes, bytearray)): - raise ValueError( - "blob-view-field requires blob field value to be a serialized " - "BlobViewStruct." - ) - try: - view_bytes = bytes(value) - view_struct = BlobViewStruct.deserialize(view_bytes) - if view_struct.serialize() != view_bytes: - raise ValueError("BlobViewStruct payload contains trailing bytes.") - except Exception as e: - raise ValueError( - "blob-view-field requires blob field value to be a serialized " - "BlobViewStruct." - ) from e + if self.blob_inline_fields: + validate_inline_blob_fields(data, self.blob_descriptor_fields, self.blob_view_fields) @staticmethod def _process_normal_data(data: pa.RecordBatch) -> Optional[pa.Table]: diff --git a/paimon-python/pypaimon/execution/python/write/writer/key_value_data_writer.py b/paimon-python/pypaimon/execution/python/write/writer/key_value_data_writer.py index 5817abe62256..5cd65c2e8947 100644 --- a/paimon-python/pypaimon/execution/python/write/writer/key_value_data_writer.py +++ b/paimon-python/pypaimon/execution/python/write/writer/key_value_data_writer.py @@ -15,19 +15,41 @@ # specific language governing permissions and limitations # under the License. +import os +import re from typing import List, Union +import uuid + import numpy as np import pyarrow as pa import pyarrow.compute as pc +from pypaimon.blob.managed_blob_cleanup import managed_blob_deletion_paths +from pypaimon.blob.managed_blob_reference_collector import ( + ManagedBlobReferenceCollector, +) +from pypaimon.blob.managed_blob_reference_file import ManagedBlobReferenceFile +from pypaimon.blob.primary_key_blob_externalizer import PrimaryKeyBlobExternalizer from pypaimon.manifest.schema.data_file_meta import DataFileMeta +from pypaimon.common.options.core_options import ChangelogProducer, CoreOptions from pypaimon.execution.python.read.reader.deduplicate_merge_function import (DeduplicateMergeFunction) -from pypaimon.common.options.core_options import ChangelogProducer +from pypaimon.schema.data_types import managed_blob_field_names from pypaimon.table.row.key_value import KeyValue from pypaimon.table.special_fields import SpecialFields from pypaimon.execution.python.write.writer.data_writer import DataWriter from pypaimon.execution.python.write.writer.write_buffer import WriteBuffer +from pypaimon.write.row_utils import validate_inline_blob_fields + + +def _qualified_local_path(path): + normalized = str(path).replace('\\', '/') + if (normalized.startswith('/') + or normalized.startswith('file:') + or '://' in normalized + or re.match(r'^[A-Za-z]:', normalized)): + return path + return os.path.abspath(path).replace('\\', '/') class KeyValueDataWriter(DataWriter): @@ -44,7 +66,7 @@ class KeyValueDataWriter(DataWriter): def __init__(self, table, partition, bucket, max_seq_number, options=None, write_cols=None, merge_function=None, - changelog_producer=ChangelogProducer.NONE): + changelog_producer=ChangelogProducer.NONE, blob_uri_reader_factory=None): super().__init__(table, partition, bucket, max_seq_number, options, write_cols, changelog_producer) # Defaults to deduplicate so direct callers (tests / future code @@ -58,10 +80,48 @@ def __init__(self, table, partition, bucket, max_seq_number, WriteBuffer(self._merge_data) if changelog_producer == ChangelogProducer.INPUT and self._writes_changelog_before_merge() else None) + inline = (self.options.blob_descriptor_fields() + | self.options.blob_view_fields()) + self._managed_blob_fields = managed_blob_field_names(self.table.fields, inline) + self._blob_externalizer = None + if self._managed_blob_fields: + self._blob_externalizer = PrimaryKeyBlobExternalizer( + self.file_io, + self._new_managed_blob_path, + self.table.fields, + self._managed_blob_fields, + self._managed_blob_target_size(), + self.options.blob_copy_buffer_size(), + uri_reader_factory=(blob_uri_reader_factory if blob_uri_reader_factory is not None + else getattr(self.file_io, "uri_reader_factory", None)), + ) def _writes_changelog_before_merge(self): return True + def write(self, data: pa.RecordBatch): + # Inline fields are persisted as-is, so raw payloads would commit + # rows that no reader can resolve. + validate_inline_blob_fields( + data, self.options.blob_descriptor_fields(), self.options.blob_view_fields()) + super().write(data) + + def _managed_blob_target_size(self) -> int: + if self.options.options.contains(CoreOptions.BLOB_TARGET_FILE_SIZE): + return self.options.blob_target_file_size() + return CoreOptions.BLOB_TARGET_FILE_SIZE.default_value().get_bytes() + + def _new_managed_blob_path(self) -> str: + file_name = "%s%s-0%s" % ( + CoreOptions.data_file_prefix(self.options), + uuid.uuid4(), + ManagedBlobReferenceFile.MANAGED_BLOB_SUFFIX, + ) + # A relative warehouse such as ``./wh`` would otherwise be stored as a + # relative pack URI. Readers then join it onto the table path again. + # Java opens the published URI directly, so local packs must be absolute. + return _qualified_local_path(self._generate_file_path(file_name)) + def _process_data(self, data: pa.RecordBatch) -> pa.Table: # No sort here: sorting once at flush is strictly cheaper than # per-batch sort + a final global sort. The buffer ends up as a @@ -81,6 +141,14 @@ def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: def prepare_commit(self) -> List[DataFileMeta]: if self._buffer.num_rows > 0 or self._has_pending_input_changelog(): self._flush_all() + if self._blob_externalizer is not None: + try: + self._blob_externalizer.prepare_commit() + except Exception: + # Closing the pack already deleted every uncommitted pack, so + # the data files that point at them cannot be committed. + self.abort() + raise # ``_flush_all`` empties the buffer, so super's prepare_commit just # returns ``committed_files``. return super().prepare_commit() @@ -104,6 +172,8 @@ def close(self): try: if self._buffer.num_rows > 0 or self._has_pending_input_changelog(): self._flush_all() + if self._blob_externalizer is not None: + self._blob_externalizer.prepare_commit() except Exception as e: import logging logger = logging.getLogger(__name__) @@ -115,6 +185,35 @@ def close(self): finally: self._buffer.reset() + def _write_data_to_file(self, data: pa.Table): + if self._blob_externalizer is not None: + data = self._blob_externalizer.externalize_table(data) + try: + super()._write_data_to_file(data) + except Exception: + if self._blob_externalizer is not None: + self._blob_externalizer.discard_pending_hand_off() + raise + if self._blob_externalizer is not None: + self._blob_externalizer.hand_off_pending() + + def _additional_extra_files(self, file_path: str, data: pa.Table) -> List[str]: + if not self._managed_blob_fields: + return [] + collector = ManagedBlobReferenceCollector( + self.file_io, file_path, self.table.fields, self._managed_blob_fields) + try: + collector.collect_table(data) + names = [collector.close()] + if self._blob_externalizer is not None: + # Packs can be shared across data files. Only the sidecar is + # exclusively owned and safe for every engine's extra-file cleanup. + self._blob_externalizer.take_referenced_pack_names() + return names + except Exception: + collector.abort() + raise + def _flush_all(self) -> None: """Sort + fold the entire buffer, then roll-write as files. @@ -163,10 +262,24 @@ def _flush_input_changelog(self): self._input_changelog_buffer.reset(data if data.num_rows else None) def abort(self): + if self._blob_externalizer is not None: + self._blob_externalizer.abort() super().abort() if self._input_changelog_buffer is not None: self._input_changelog_buffer.reset() + def _delete_committed_files(self, file_metas: List[DataFileMeta]): + if not self._managed_blob_fields: + return super()._delete_committed_files(file_metas) + # Discover shared references before the base writer removes sidecars. + paths = [path for meta in file_metas for path in meta.collect_files()] + deleting = managed_blob_deletion_paths(self.table, paths) + super()._delete_committed_files(file_metas) + owned_paths = set(paths) + for path in deleting: + if path not in owned_paths: + self.file_io.delete_quietly(path) + def _roll_write(self, data: pa.Table) -> None: """Write ``data`` as one or more files, each <= target_file_size. diff --git a/paimon-python/pypaimon/read/reader/iface/record_batch_reader.py b/paimon-python/pypaimon/read/reader/iface/record_batch_reader.py index e9e8f38c6904..ac4d1cbbb009 100644 --- a/paimon-python/pypaimon/read/reader/iface/record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/iface/record_batch_reader.py @@ -83,18 +83,15 @@ class RecordBatchReader(RecordReader): """ file_io = None + blob_base_path = None blob_field_indices = None descriptor_field_indices = None + resolved_blob_field_indices = None blob_view_lookup = None vector_field_indices = None def _adopt_metadata(self, reader: "RecordBatchReader") -> None: - self.file_io = reader.file_io - self.blob_field_indices = reader.blob_field_indices - self.descriptor_field_indices = getattr( - reader, 'descriptor_field_indices', None) - self.blob_view_lookup = getattr(reader, 'blob_view_lookup', None) - self.vector_field_indices = reader.vector_field_indices + self._adopt_blob_metadata(reader) @abstractmethod def read_arrow_batch(self) -> Optional[RecordBatch]: @@ -127,7 +124,8 @@ def read_batch(self) -> Optional[RecordIterator[InternalRow]]: return InternalRowWrapperIterator( self._iter_df_rows(df), df.width, self.file_io, self.blob_field_indices, self.vector_field_indices, - self.descriptor_field_indices, self.blob_view_lookup) + self.descriptor_field_indices, self.blob_view_lookup, + getattr(self, 'blob_base_path', None), self.resolved_blob_field_indices) @staticmethod def _iter_df_rows(df) -> Iterator[tuple]: @@ -143,14 +141,17 @@ def __init__(self, iterator: Iterator[tuple], width: int, file_io=None, blob_field_indices=None, vector_field_indices=None, descriptor_field_indices=None, - blob_view_lookup=None): + blob_view_lookup=None, + blob_base_path=None, resolved_blob_field_indices=None): self._iterator = iterator self._reused_row = OffsetRow(None, 0, width, file_io=file_io, blob_field_indices=blob_field_indices, vector_field_indices=vector_field_indices, descriptor_field_indices=descriptor_field_indices, - blob_view_lookup=blob_view_lookup) + blob_view_lookup=blob_view_lookup, + blob_base_path=blob_base_path, + resolved_blob_field_indices=resolved_blob_field_indices) def next(self) -> Optional[InternalRow]: row_tuple = next(self._iterator, None) diff --git a/paimon-python/pypaimon/read/reader/iface/record_reader.py b/paimon-python/pypaimon/read/reader/iface/record_reader.py index d777994c6d17..2236e3d37957 100644 --- a/paimon-python/pypaimon/read/reader/iface/record_reader.py +++ b/paimon-python/pypaimon/read/reader/iface/record_reader.py @@ -42,8 +42,10 @@ def close(self): def _adopt_blob_metadata(self, reader) -> None: self.file_io = getattr(reader, 'file_io', None) + self.blob_base_path = getattr(reader, 'blob_base_path', None) self.blob_field_indices = getattr(reader, 'blob_field_indices', None) self.descriptor_field_indices = getattr(reader, 'descriptor_field_indices', None) + self.resolved_blob_field_indices = getattr(reader, 'resolved_blob_field_indices', None) self.blob_view_lookup = getattr(reader, 'blob_view_lookup', None) self.vector_field_indices = getattr(reader, 'vector_field_indices', None) diff --git a/paimon-python/pypaimon/read/reader/managed_blob_record_batch_reader.py b/paimon-python/pypaimon/read/reader/managed_blob_record_batch_reader.py new file mode 100644 index 000000000000..ed46b7c45eeb --- /dev/null +++ b/paimon-python/pypaimon/read/reader/managed_blob_record_batch_reader.py @@ -0,0 +1,73 @@ +# 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 pyarrow as pa + +from pypaimon.blob.managed_blob_payload import materialize_managed_blob +from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader + + +def _replace_batch_column(batch, index, field, array): + """Replace one column without ``RecordBatch.set_column``. + + That method arrived in PyArrow 13. Python 3.6 and 3.7 stay on older releases. + """ + columns = list(batch.columns) + fields = [batch.schema.field(i) for i in range(batch.num_columns)] + columns[index] = array + fields[index] = field + return pa.RecordBatch.from_arrays( + columns, schema=pa.schema(fields, metadata=batch.schema.metadata)) + + +class ManagedBlobRecordBatchReader(RecordBatchReader): + """Resolve managed payloads before filtering, limiting or projecting rows.""" + + def __init__(self, reader, specs, file_io, table_path): + self.reader = reader + self.specs = specs + self.table_path = table_path + self._adopt_metadata(reader) + self.file_io = file_io + # Selected MAP values become scalar BLOB columns after projection. + self.blob_field_indices = set(self.blob_field_indices or ()) | { + index for index, _name, kind in specs if kind == 'scalar'} + # Resolved scalar BLOBs must be exposed as payloads by get_blob(). + resolved = {index for index, _name, _kind in specs} + self.descriptor_field_indices = ( + set(self.descriptor_field_indices or ()) - resolved) + # Payload bytes are opaque even when they contain reference magic headers. + self.resolved_blob_field_indices = set(self.resolved_blob_field_indices or ()) | resolved + + def read_arrow_batch(self): + batch = self.reader.read_arrow_batch() + if batch is None: + return None + self._refresh_blob_view_lookup(self.reader) + for _index, name, kind in self.specs: + column_index = batch.schema.get_field_index(name) + field = batch.schema.field(column_index) + values = [ + materialize_managed_blob(value, kind, self.file_io, self.table_path) + for value in batch.column(column_index).to_pylist() + ] + batch = _replace_batch_column( + batch, column_index, field, pa.array(values, type=field.type)) + return batch + + def close(self): + self.reader.close() diff --git a/paimon-python/pypaimon/schema/data_types.py b/paimon-python/pypaimon/schema/data_types.py index 893a5f2341ac..2e561a29c02f 100755 --- a/paimon-python/pypaimon/schema/data_types.py +++ b/paimon-python/pypaimon/schema/data_types.py @@ -300,6 +300,16 @@ def is_blob_file_field(field: 'DataField') -> bool: return is_blob_file_type(field.type) +def managed_blob_field_names(fields, inline_fields=None): + """Blob-file fields stored in managed packs rather than inline descriptors.""" + inline = inline_fields or set() + return { + field.name + for field in fields + if is_blob_file_field(field) and field.name not in inline + } + + @dataclass class DataField: FIELD_ID = "id" diff --git a/paimon-python/pypaimon/schema/schema_manager.py b/paimon-python/pypaimon/schema/schema_manager.py index 0d375d49a2bb..32d4fdf328c2 100644 --- a/paimon-python/pypaimon/schema/schema_manager.py +++ b/paimon-python/pypaimon/schema/schema_manager.py @@ -31,7 +31,8 @@ DataType, MapType, MultisetType, RowType, is_array_blob_type, is_blob_file_field, is_blob_file_type, is_blob_type, - is_map_blob_type, reassign_field_id) + is_map_blob_type, managed_blob_field_names, + reassign_field_id) from pypaimon.schema.schema import Schema from pypaimon.schema.schema_change import (AddColumn, DropColumn, RemoveOption, RenameColumn, SchemaChange, @@ -349,6 +350,105 @@ def _assert_not_renaming_blob_column( ) +def _effective_bucket_keys(options, primary_keys, partition_keys): + configured = (options or {}).get(CoreOptions.BUCKET_KEY.key()) + if configured and str(configured).strip(): + return [key.strip() for key in str(configured).split(',') if key.strip()] + return [key for key in primary_keys if key not in partition_keys] + + +def _sequence_group_options(options): + return [(key, value) for key, value in options.items() + if key.startswith('fields.') and key.endswith('.sequence-group')] + + +def _validate_blob_sequence_group_ordering(fields, core_options): + from pypaimon.common.options.core_options import MergeEngine + + if core_options.merge_engine() != MergeEngine.PARTIAL_UPDATE: + return + fields_by_name = {field.name: field for field in fields} + for key, _value in _sequence_group_options(core_options.options.to_map()): + for name in key[len('fields.'):-len('.sequence-group')].split(','): + field = fields_by_name.get(name) + if field is not None and _contains_blob_type(field.type): + raise ValueError( + "Field '%s' with type %s cannot be used as a sequence-group " + "ordering field in option '%s'." % (name, field.type, key)) + + +def _validate_managed_blob_retracts(managed, core_options): + from pypaimon.common.options.core_options import MergeEngine + + if core_options.merge_engine() != MergeEngine.PARTIAL_UPDATE or core_options.ignore_delete(): + return + options = core_options.options.to_map() + protected = {name for _key, value in _sequence_group_options(options) + for name in value.split(',')} + for name in sorted(managed & protected): + aggregate = options.get('fields.%s.aggregate-function' % name) + if aggregate is None: + aggregate = options.get('fields.default-aggregate-function') + if aggregate is None or aggregate == 'last_value': + continue + ignore_retract = str(options.get('fields.%s.ignore-retract' % name, 'false')).lower() + if ignore_retract not in ('true', 'false'): + raise ValueError( + "Unrecognized boolean value for 'fields.%s.ignore-retract': %s. " + "Expected true or false (case insensitive)." % (name, ignore_retract)) + if ignore_retract != 'true': + raise ValueError( + "Managed BLOB field '%s' cannot use aggregate function '%s' because " + "managed BLOB payloads are not retained in retract messages. " + "Set 'fields.%s.ignore-retract' to true to ignore retract messages." + % (name, aggregate, name)) + + +def _validate_primary_key_managed_blobs(fields, options, primary_keys, partition_keys): + """Match Java SchemaValidation for primary-key managed BLOB fields.""" + from pypaimon.common.options.core_options import ChangelogProducer, MergeEngine + + core_options = CoreOptions(Options(options or {})) + inline = core_options.blob_descriptor_fields() | core_options.blob_view_fields() + managed = managed_blob_field_names(fields, inline) + if not managed: + return + + engine = core_options.merge_engine() + if engine not in ( + MergeEngine.DEDUPLICATE, MergeEngine.PARTIAL_UPDATE, MergeEngine.FIRST_ROW): + raise ValueError( + "Primary-key managed BLOB tables only support the deduplicate, " + "partial-update or first-row merge engine.") + if core_options.changelog_producer() != ChangelogProducer.NONE: + raise ValueError( + "Primary-key managed BLOB tables only support changelog-producer 'none'.") + if core_options.data_file_external_paths(): + raise ValueError( + "Primary-key managed BLOB tables do not support '%s'." + % CoreOptions.DATA_FILE_EXTERNAL_PATHS.key()) + if str((options or {}).get('pk-clustering-override', '')).lower() == 'true': + raise ValueError( + "Primary-key managed BLOB tables do not support 'pk-clustering-override'.") + + primary_key_blobs = sorted(name for name in primary_keys if name in managed) + if primary_key_blobs: + raise ValueError( + "Managed BLOB fields cannot be primary keys: %s." % primary_key_blobs) + bucket_blobs = sorted( + name for name in _effective_bucket_keys(options, primary_keys, partition_keys) + if name in managed) + if bucket_blobs: + raise ValueError( + "Managed BLOB fields cannot be bucket keys: %s." % bucket_blobs) + sequence_blobs = sorted( + name for name in core_options.sequence_field() if name in managed) + if sequence_blobs: + raise ValueError( + "Managed BLOB fields cannot be sequence fields: %s." % sequence_blobs) + _validate_managed_blob_retracts(managed, core_options) + + def _validate_blob_fields( fields: List[DataField], options: dict, @@ -381,6 +481,8 @@ def _validate_blob_fields( core_options = CoreOptions(Options(options)) + _validate_blob_sequence_group_ordering(fields, core_options) + configured_blob_fields = core_options.blob_field() for field in configured_blob_fields: if field not in blob_field_names: @@ -394,6 +496,8 @@ def _validate_blob_fields( descriptor_fields = core_options.blob_descriptor_fields() view_fields = core_options.blob_view_fields() video_fields = core_options.video_frame_fields() + if primary_keys and video_fields: + raise ValueError("'video-frame-field' only supports append-only tables.") non_scalar_video_fields = video_fields.difference(scalar_blob_field_names) if non_scalar_video_fields: raise ValueError( @@ -438,35 +542,34 @@ def _validate_blob_fields( ) if blob_field_names: - required_options = { - CoreOptions.ROW_TRACKING_ENABLED.key(): 'true', - CoreOptions.DATA_EVOLUTION_ENABLED.key(): 'true' - } - - missing_options = [] - for key, expected_value in required_options.items(): - if key not in options or options[key] != expected_value: - missing_options.append(f"{key}='{expected_value}'") - - if missing_options: - raise ValueError( - f"Schema contains BLOB, ARRAY or MAP type but is " - f"missing required options: " - f"{', '.join(missing_options)}. " - f"Please add these options to the schema." - ) - - if primary_keys: - raise ValueError( - "BLOB, ARRAY or MAP type is not supported with primary key." - ) - if blob_field_names.intersection(partition_keys): raise ValueError( "The BLOB, ARRAY or MAP type column can not be " "part of partition keys." ) + if primary_keys: + _validate_primary_key_managed_blobs( + fields, options, primary_keys, partition_keys) + else: + required_options = { + CoreOptions.ROW_TRACKING_ENABLED.key(): 'true', + CoreOptions.DATA_EVOLUTION_ENABLED.key(): 'true' + } + + missing_options = [] + for key, expected_value in required_options.items(): + if key not in options or options[key] != expected_value: + missing_options.append(f"{key}='{expected_value}'") + + if missing_options: + raise ValueError( + f"Schema contains BLOB, ARRAY or MAP type but is " + f"missing required options: " + f"{', '.join(missing_options)}. " + f"Please add these options to the schema." + ) + def _validate_options(options: dict): core_options = CoreOptions(Options(options)) diff --git a/paimon-python/pypaimon/table/row/blob.py b/paimon-python/pypaimon/table/row/blob.py index 769fdfb7171a..8fc91d6d9cc5 100644 --- a/paimon-python/pypaimon/table/row/blob.py +++ b/paimon-python/pypaimon/table/row/blob.py @@ -693,7 +693,13 @@ def from_descriptor(uri_reader: UriReader, descriptor: BlobDescriptor) -> 'Blob' @staticmethod def _blob_ref_from_descriptor( descriptor: 'BlobDescriptor', file_io=None, uri_reader_factory=None, + table_path=None, ) -> 'BlobRef': + if table_path: + from pypaimon.blob.managed_blob_payload import resolve_table_blob_uri + uri = resolve_table_blob_uri(table_path, descriptor.uri) + if uri != descriptor.uri: + descriptor = BlobDescriptor(uri, descriptor.offset, descriptor.length) if uri_reader_factory is None: if file_io is None: raise ValueError("file_io is required to resolve BlobDescriptor bytes") @@ -705,6 +711,7 @@ def _blob_ref_from_descriptor( @staticmethod def from_descriptor_bytes( data: Optional[bytes], file_io=None, uri_reader_factory=None, + table_path=None, ) -> Optional['Blob']: """Build a Blob from bytes known to contain a descriptor. @@ -731,7 +738,8 @@ def from_descriptor_bytes( raise ValueError( "Expected BlobDescriptor bytes, got raw bytes") from exc return Blob._blob_ref_from_descriptor( - descriptor, file_io=file_io, uri_reader_factory=uri_reader_factory) + descriptor, file_io=file_io, uri_reader_factory=uri_reader_factory, + table_path=table_path) @staticmethod def from_view(*args, view_struct: Optional[BlobViewStruct] = None) -> 'BlobView': @@ -760,6 +768,7 @@ def from_view(*args, view_struct: Optional[BlobViewStruct] = None) -> 'BlobView' @staticmethod def from_bytes( data: Optional[bytes], file_io=None, uri_reader_factory=None, allow_blob_data: bool = True, + table_path=None, ) -> Optional['Blob']: if data is None: return None @@ -777,7 +786,8 @@ def from_bytes( + ("" if allow_blob_data else " (allow_blob_data=False)") ) from exc return Blob._blob_ref_from_descriptor( - descriptor, file_io=file_io, uri_reader_factory=uri_reader_factory) + descriptor, file_io=file_io, uri_reader_factory=uri_reader_factory, + table_path=table_path) return BlobData(data) diff --git a/paimon-python/pypaimon/table/row/offset_row.py b/paimon-python/pypaimon/table/row/offset_row.py index 1592b94ae472..45bc3139c26b 100644 --- a/paimon-python/pypaimon/table/row/offset_row.py +++ b/paimon-python/pypaimon/table/row/offset_row.py @@ -27,12 +27,16 @@ def __init__(self, row_tuple: Optional[tuple], offset: int, arity: int, file_io=None, blob_field_indices: Optional[Iterable[int]] = None, vector_field_indices: Optional[Iterable[int]] = None, descriptor_field_indices: Optional[Iterable[int]] = None, - blob_view_lookup=None): + blob_view_lookup=None, + blob_base_path=None, + resolved_blob_field_indices: Optional[Iterable[int]] = None): self.row_tuple = row_tuple self.offset = offset self.arity = arity self.row_kind_byte: int = RowKind.INSERT.value self._file_io = file_io + self._blob_base_path = blob_base_path + self._resolved_blob_field_indices = frozenset(resolved_blob_field_indices or ()) self._blob_field_indices: FrozenSet[int] = ( frozenset(blob_field_indices) if blob_field_indices is not None else frozenset() ) @@ -105,12 +109,15 @@ def get_blob(self, pos: int): if value is None: return None raw = self._normalize_blob_bytes(value) + if pos in self._resolved_blob_field_indices: + return Blob.from_data(raw) if raw is not None and BlobViewStruct.is_blob_view_struct(raw): return self._resolve_blob_view_struct(BlobViewStruct.deserialize(raw)) if pos in self._descriptor_field_indices: return self._blob_from_descriptor_field_bytes(raw) return Blob.from_bytes( - raw, self._file_io, uri_reader_factory=self._uri_reader_factory()) + raw, self._file_io, uri_reader_factory=self._uri_reader_factory(), + table_path=self._blob_base_path) def get_vector(self, pos: int): from pypaimon.table.row.vector import Vector diff --git a/paimon-python/pypaimon/tests/auth_masking_reader_test.py b/paimon-python/pypaimon/tests/auth_masking_reader_test.py index 342c2548ad0b..9557caf2c6f4 100644 --- a/paimon-python/pypaimon/tests/auth_masking_reader_test.py +++ b/paimon-python/pypaimon/tests/auth_masking_reader_test.py @@ -881,9 +881,76 @@ def test_returns_none_at_end(self): reader = ColumnProjectReader(_FakeBatchReader([]), ["a"]) self.assertIsNone(reader.read_arrow_batch()) + def test_projected_payload_stays_opaque_through_row_adapters(self): + from pypaimon.execution.python.read.reader.auth_masking_reader import ( + BatchToRecordReaderAdapter, RecordReaderToBatchAdapter) + from pypaimon.table.row.blob import BlobViewStruct + + payload = BlobViewStruct('db.payload', 7, 42).serialize() + source = _FakeBatchReader([pa.RecordBatch.from_pydict({'id': [1], 'payload': [payload]})]) + source.blob_field_indices = {1} + source.resolved_blob_field_indices = {1} + projected = ColumnProjectReader(source, ['payload', 'id'], + [_FakeField('id'), _FakeField('payload')]) + rows = BatchToRecordReaderAdapter(projected) + batches = RecordReaderToBatchAdapter(rows, pa.schema([('payload', pa.binary()), ('id', pa.int64())])) + reader = BatchToRecordReaderAdapter(batches) + try: + row = reader.read_batch().next() + self.assertEqual(row.get_blob(0).to_data(), payload) + self.assertEqual(row.get_field(1), 1) + finally: + reader.close() + class TestColumnProjectReaderRowKind(unittest.TestCase): + def test_adapter_without_row_kind_keeps_user_column(self): + from pypaimon.execution.python.read.reader.auth_masking_reader import ( + BatchToRecordReaderAdapter, RecordReaderToBatchAdapter) + from pypaimon.utils.arrow_utils import row_kind_field + + fields = pa.schema([('_row_kind', pa.string()), ('value', pa.int64())]) + batch = pa.RecordBatch.from_arrays( + [pa.array(['+I']), pa.array(['hello']), pa.array([1])], + schema=pa.schema([row_kind_field(fields)] + list(fields))) + for columns in (['_row_kind', 'value'], ['value']): + rows = BatchToRecordReaderAdapter(_FakeBatchReader([batch])) + batches = RecordReaderToBatchAdapter(rows, fields) + projected = ColumnProjectReader(batches, columns, list(fields)) + result = BatchToRecordReaderAdapter(projected) + try: + row = result.read_batch().next() + self.assertEqual(row.arity, len(columns)) + for index, name in enumerate(columns): + self.assertEqual(row.get_field(index), 'hello' if name == '_row_kind' else 1) + finally: + result.close() + + def test_round_trip_preserves_row_kind_with_colliding_user_columns(self): + from pypaimon.execution.python.read.reader.auth_masking_reader import ( + BatchToRecordReaderAdapter, RecordReaderToBatchAdapter) + from pypaimon.table.row.row_kind import RowKind + from pypaimon.utils.arrow_utils import row_kind_field + + fields = pa.schema([('_row_kind', pa.string()), ('__row_kind', pa.string())]) + batch = pa.RecordBatch.from_arrays( + [pa.array(['-D', '+U']), pa.array(['hello', 'world']), pa.array(['one', 'two'])], + schema=pa.schema([row_kind_field(fields)] + list(fields))) + rows = BatchToRecordReaderAdapter(_FakeBatchReader([batch])) + batches = RecordReaderToBatchAdapter(rows, fields, include_row_kind=True) + projected = ColumnProjectReader(batches, ['__row_kind', '_row_kind'], list(fields)) + result = BatchToRecordReaderAdapter(projected) + try: + iterator = result.read_batch() + for kind, values in ((RowKind.DELETE, ('one', 'hello')), (RowKind.UPDATE_AFTER, ('two', 'world'))): + row = iterator.next() + self.assertEqual(row.get_row_kind(), kind) + self.assertEqual(tuple(row.get_field(index) for index in range(2)), values) + self.assertIsNone(iterator.next()) + finally: + result.close() + def test_preserves_row_kind_column(self): batch = pa.RecordBatch.from_pydict({ "_row_kind": ["+I", "-D"], diff --git a/paimon-python/pypaimon/tests/blob_test.py b/paimon-python/pypaimon/tests/blob_test.py index 9d0f5a660b5e..9012315cdcfc 100644 --- a/paimon-python/pypaimon/tests/blob_test.py +++ b/paimon-python/pypaimon/tests/blob_test.py @@ -2200,6 +2200,67 @@ class _Table: self.assertEqual(blob.to_data(), v1_shaped_inline) reader.close() + def test_inline_materialized_payloads_preserve_reference_headers(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.execution.python.read.reader.auth_masking_reader import ( + BatchToRecordReaderAdapter, ColumnProjectReader) + from pypaimon.execution.python.read.reader.blob_descriptor_convert_reader import ( + BlobInlineConvertReader) + from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader + + with tempfile.TemporaryDirectory() as directory: + file_io = LocalFileIO.create() + target = Path(directory) / 'target' + target.write_bytes(b'wrong payload') + payloads = [BlobViewStruct('db.payload', 7, 42).serialize(), + BlobDescriptor(str(target), 0, target.stat().st_size).serialize()] + fields = [DataField(0, 'id', AtomicType('INT')), DataField(1, 'payload', AtomicType('BLOB'))] + for kind in ('descriptor', 'view'): + for payload in payloads: + with self.subTest(kind=kind, payload=payload): + source = Path(directory) / 'source' + source.write_bytes(payload) + descriptor = BlobDescriptor(str(source), 0, len(payload)) + value = (descriptor.serialize() if kind == 'descriptor' + else BlobViewStruct('db.source', 1, 2).serialize()) + batch = pa.RecordBatch.from_arrays( + [pa.array(['-D']), pa.array([1]), + pa.array([value], type=pa.large_binary())], names=['_row_kind', 'id', 'payload']) + + class InnerReader(RecordBatchReader): + def __init__(self): + self.file_io = file_io + self.blob_field_indices = {1} + self.descriptor_field_indices = {1} + self.batch = batch + + def read_arrow_batch(self): + result, self.batch = self.batch, None + return result + + def close(self): + pass + + table = MagicMock() + table.file_io = file_io + table.options = CoreOptions(Options({ + 'blob-as-descriptor': 'false', 'blob-%s-field' % kind: 'payload'})) + reader = BlobInlineConvertReader( + InnerReader(), table, field_names=[field.name for field in fields]) + if kind == 'view': + lookup = MagicMock() + lookup.resolve_to_null.return_value = False + lookup.resolve_blob.return_value = Blob.from_bytes(descriptor.serialize(), file_io) + reader._blob_view_lookup = lookup + reader._prescan_done = True + projected = ColumnProjectReader(reader, ['payload', 'id'], fields) + rows = BatchToRecordReaderAdapter(projected) + row = rows.read_batch().next() + self.assertEqual(row.get_row_kind(), RowKind.DELETE) + self.assertIsInstance(row.get_blob(0), BlobData) + self.assertEqual(row.get_blob(0).to_data(), payload) + rows.close() + def test_wrap_record_reader_propagates_blob_metadata_for_get_blob(self): from pypaimon.common.options import Options from pypaimon.common.options.core_options import CoreOptions diff --git a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py index 9bbb6e351415..546e86e2e734 100644 --- a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py +++ b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py @@ -55,6 +55,20 @@ def boolean_query(queries): return json.dumps({"boolean": {"queries": queries}}, separators=(",", ":")) +def arrow_table_rows(table): + """Row dicts with the same values as ``Table.to_pylist``. + + ``Table.to_pylist`` is absent on PyArrow 6, the version used by Python 3.6. + ``Table.to_pydict`` is present there and returns each column via ``as_py``. + """ + columns = table.to_pydict() + names = list(columns) + return [ + {name: columns[name][index] for name in names} + for index in range(table.num_rows) + ] + + class JavaPyReadWriteTest(unittest.TestCase): @classmethod def setUpClass(cls): @@ -138,6 +152,48 @@ def test_py_write_dynamic_bucket_hash_index(self): writer.close() commit.close() + def test_py_write_managed_blob(self): + table_name = 'default.managed_blob_python_to_java' + self.catalog.drop_table(table_name, True) + arrow_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + ('payload', pa.large_binary()), + ('items', pa.list_(pa.large_binary())), + ('mapping', pa.map_(pa.string(), pa.large_binary())), + ]) + self.catalog.create_table(table_name, Schema.from_pyarrow_schema( + arrow_schema, primary_keys=['id'], options={ + 'bucket': '1', 'file.format': 'parquet', + 'write.native.enabled': 'false', 'commit.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + }), False) + table = self.catalog.get_table(table_name) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'payload': [b'python-payload'], + 'items': [[b'python-item']], 'mapping': [[('k', b'python-map')]], + }, schema=arrow_schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + self._check_python_managed_blob() + + def test_read_python_managed_blob_after_java_gc(self): + self._check_python_managed_blob() + + def _check_python_managed_blob(self): + table = self.catalog.get_table('default.managed_blob_python_to_java') + builder = table.new_read_builder() + rows = arrow_table_rows(builder.new_read().to_arrow( + builder.new_scan().plan().splits())) + self.assertEqual(rows, [{ + 'id': 1, 'payload': b'python-payload', + 'items': [b'python-item'], 'mapping': [('k', b'python-map')], + }]) + @parameterized.expand([ (type_name, order, grouping) for type_name in ('float', 'double') diff --git a/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py b/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py index 1b49fa19c431..4f4bdad3aa36 100644 --- a/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py +++ b/paimon-python/pypaimon/tests/managed_blob_lifecycle_test.py @@ -15,12 +15,18 @@ # specific language governing permissions and limitations # under the License. +import os +import tempfile import unittest -from unittest.mock import Mock, call +from pathlib import Path +from unittest.mock import Mock, call, patch +from pypaimon.blob.managed_blob_reference_file import ManagedBlobReferenceFile, Reference +from pypaimon.blob.managed_blob_cleanup import managed_blob_deletion_paths from pypaimon.manifest.schema.data_file_meta import DataFileMeta +from pypaimon.filesystem.local_file_io import LocalFileIO from pypaimon.write.commit_message import CommitMessage -from pypaimon.write.file_store_commit import FileStoreCommit +from pypaimon.execution.python.write.file_store_commit import FileStoreCommit def _file(path, extra_files): @@ -43,8 +49,68 @@ def _file(path, extra_files): class ManagedBlobLifecycleTest(unittest.TestCase): + def test_cleanup_compares_qualified_local_references(self): + parent = os.path.abspath('warehouse %?#/bucket-0') + pack = parent + '/pack.managed.blob' + sidecar = 'warehouse %?#/bucket-0/data.blobref' + file_io = Mock() + file_io.to_filesystem_path.side_effect = LocalFileIO.create().to_filesystem_path + file_io.list_status.return_value = [Mock(base_name='data.blobref')] + table = Mock(file_io=file_io, table_path=os.path.dirname(parent)) + reference = Reference(Path(parent).as_uri(), 'pack.managed.blob') + with patch.object(ManagedBlobReferenceFile, 'read', return_value=[reference]): + paths = managed_blob_deletion_paths(table, [sidecar]) + self.assertIn(pack, paths) + file_io.list_status.assert_called_once_with(parent) + + def test_local_uri_cleanup_preserves_surviving_references(self): + file_io = LocalFileIO.create() + with tempfile.TemporaryDirectory(prefix='pack cleanup ') as directory: + root = Path(directory) + for name in ('space dir', 'literal%20', 'literal%?#'): + parent = root / name + parent.mkdir() + pack = parent / 'pack.managed.blob' + pack.write_bytes(b'payload') + aborted, live = parent / 'aborted.blobref', parent / 'live.blobref' + table = Mock(file_io=file_io, table_path=str(root)) + roots = (str(parent), parent.as_uri(), 'file:' + str(parent), os.path.relpath(str(parent))) + for aborted_root in roots: + for live_root in roots: + with self.subTest(name=name, aborted_root=aborted_root, live_root=live_root): + reference = Reference(aborted_root, pack.name) + live_reference = Reference(live_root, pack.name) + ManagedBlobReferenceFile.write(file_io, str(aborted), [reference]) + ManagedBlobReferenceFile.write(file_io, str(live), [live_reference]) + paths = managed_blob_deletion_paths(table, [aborted.as_uri()]) + self.assertEqual(paths, [aborted.as_uri()]) + for path in paths: + file_io.delete_quietly(path) + self.assertTrue(pack.exists()) + paths = managed_blob_deletion_paths(table, [str(live)]) + self.assertEqual(set(paths), {str(live), str(pack)}) + + def test_remote_cleanup_uses_full_reference_identity(self): + root = 's3://bucket/table %?#/bucket-0' + pack = root + '/pack.managed.blob' + sidecar = root + '/aborted.blobref' + for surviving_root in (root, 's3://other/table %?#/bucket-0'): + with self.subTest(surviving_root=surviving_root): + file_io = Mock() + file_io.list_status.return_value = [Mock(base_name='live.blobref')] + table = Mock(file_io=file_io, table_path='s3://bucket/table %?#') + + def references(_file_io, path): + return [Reference(root if path == sidecar else surviving_root, 'pack.managed.blob')] + + with patch.object(ManagedBlobReferenceFile, 'read', side_effect=references): + paths = managed_blob_deletion_paths(table, [sidecar]) + self.assertEqual(pack in paths, surviving_root != root) + file_io.list_status.assert_called_once_with(root) + def test_commit_abort_deletes_sidecars_and_unpreserved_packs(self): file_io = Mock() + file_io.list_status.return_value = [] commit = FileStoreCommit.__new__(FileStoreCommit) commit.table = Mock(file_io=file_io) @@ -69,6 +135,7 @@ def test_commit_abort_deletes_sidecars_and_unpreserved_packs(self): def test_commit_abort_deletes_blob_packs(self): file_io = Mock() + file_io.list_status.return_value = [] commit = FileStoreCommit.__new__(FileStoreCommit) commit.table = Mock(file_io=file_io) @@ -92,6 +159,38 @@ def test_commit_abort_deletes_blob_packs(self): call("/warehouse/table/bucket-0/payload.blob"), ], any_order=True) + def test_commit_abort_discovers_packs_from_sidecars(self): + for listing_fails in (False, True): + with self.subTest(listing_fails=listing_fails): + file_io = Mock() + file_io.list_status.return_value = [] + if listing_fails: + file_io.list_status.side_effect = OSError('injected') + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io, table_path='/warehouse/table') + data_file = _file( + '/warehouse/table/bucket-0/data.avro', ['data.avro.blobref']) + reference = Reference('bucket-0', 'pack.managed.blob') + with patch.object(ManagedBlobReferenceFile, 'read', return_value=[reference]): + commit.abort([CommitMessage((), 0, [data_file])]) + deleted = {args[0] for args, _kwargs in file_io.delete_quietly.call_args_list} + self.assertIn('/warehouse/table/bucket-0/data.avro', deleted) + self.assertIn('/warehouse/table/bucket-0/data.avro.blobref', deleted) + self.assertEqual( + '/warehouse/table/bucket-0/pack.managed.blob' in deleted, not listing_fails) + + def test_commit_abort_keeps_pack_when_surviving_sidecar_is_unreadable(self): + file_io = Mock() + file_io.list_status.return_value = [Mock(base_name='committed.avro.blobref')] + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io, table_path='/warehouse/table') + data_file = _file('/warehouse/table/bucket-0/data.avro', ['data.avro.blobref']) + with patch.object(ManagedBlobReferenceFile, 'read', side_effect=[ + [Reference('bucket-0', 'pack.managed.blob')], OSError('injected')]): + commit.abort([CommitMessage((), 0, [data_file])]) + deleted = {args[0] for args, _kwargs in file_io.delete_quietly.call_args_list} + 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) @@ -115,6 +214,28 @@ def collect_files(self, bucket_path=None): file_io.delete_quietly.assert_called_once_with( "/warehouse/table/bucket-0/data.avro") + def test_commit_abort_keeps_packs_when_listing_fails(self): + file_io = Mock() + file_io.list_status.side_effect = OSError("injected") + commit = FileStoreCommit.__new__(FileStoreCommit) + commit.table = Mock(file_io=file_io) + data_file = _file( + "/warehouse/table/bucket-0/data.avro", + ["data.avro.blobref", "pack.managed.blob"], + ) + message = CommitMessage( + partition=(), + bucket=0, + new_files=[data_file], + ) + + commit.abort([message]) + + deleted = [args[0] for args, _kwargs in file_io.delete_quietly.call_args_list] + self.assertIn("/warehouse/table/bucket-0/data.avro", deleted) + self.assertIn("/warehouse/table/bucket-0/data.avro.blobref", deleted) + self.assertNotIn("/warehouse/table/bucket-0/pack.managed.blob", deleted) + if __name__ == "__main__": unittest.main() diff --git a/paimon-python/pypaimon/tests/managed_blob_pk_write_test.py b/paimon-python/pypaimon/tests/managed_blob_pk_write_test.py new file mode 100644 index 000000000000..ec9c2559c5ef --- /dev/null +++ b/paimon-python/pypaimon/tests/managed_blob_pk_write_test.py @@ -0,0 +1,987 @@ +# 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 os +import shutil +import tempfile +import unittest +from unittest.mock import Mock, patch + +import pyarrow as pa +import pyarrow.parquet as pq + +from pypaimon import CatalogFactory, Schema +from pypaimon.blob.managed_blob_payload import ( + materialize_managed_blob, resolve_table_blob_uri) +from pypaimon.blob.managed_blob_reference_file import ManagedBlobReferenceFile +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, MapType, PyarrowFieldParser +from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct +from pypaimon.write.commit_message import CommitMessage + + +class ManagedBlobPrimaryKeyWriteTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.catalog = CatalogFactory.create({'warehouse': self.temp_dir}) + self.catalog.create_database('default', False) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_pk_write_uses_custom_blob_uri_reader_factory(self): + for streaming in (False, True): + with self.subTest(streaming=streaming): + table = self._create_table('custom_uri_%s' % streaming, { + 'bucket': '1', 'write.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false'}) + builder = (table.new_stream_write_builder() if streaming else table.new_batch_write_builder()) + writer, commit = builder.new_write(), builder.new_commit() + factory = Mock() + factory.create.return_value.new_input_stream.side_effect = ( + lambda uri: io.BytesIO(b'prefixPAYLOADsuffix')) + writer.with_blob_uri_reader_factory(factory) + descriptor = BlobDescriptor('custom://source', 6, 7).serialize() + try: + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'payload': [descriptor], 'items': [None], 'mapping': [None], + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields))) + messages = writer.prepare_commit(1) if streaming else writer.prepare_commit() + factory.create.assert_called_once_with('custom://source') + raw = self._parquet_rows(table, messages[0].new_files[0])[0]['payload'] + self.assertEqual(self._written_payload(table, raw), b'PAYLOAD') + commit.abort(messages) + self.assertFalse(any(name.endswith('.managed.blob') for name in self._names_under(table))) + finally: + writer.close() + commit.close() + + def test_merge_read_preserves_user_row_kind_columns(self): + fields = [DataField(0, 'id', AtomicType('INT', False)), + DataField(1, '_row_kind', AtomicType('STRING')), + DataField(2, '__row_kind', AtomicType('STRING')), + DataField(3, 'payload', AtomicType('BLOB'))] + self.catalog.create_table('default.row_kind_columns', Schema( + fields=fields, primary_keys=['id'], options={ + 'bucket': '1', 'write.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false'}), False) + table = self.catalog.get_table('default.row_kind_columns') + data = {'id': [1, 2], '_row_kind': ['hello', 'world'], + '__row_kind': ['first', 'second'], 'payload': [b'visible', b'other']} + expected = [dict(zip(data, values)) for values in zip(*data.values())] + for generation in range(2): + self._write_and_commit(table, pa.Table.from_pydict( + data, schema=PyarrowFieldParser.from_paimon_schema(fields))) + for projection in (list(data), ['payload', '_row_kind'], ['__row_kind', '_row_kind']): + for filtered in (False, True): + with self.subTest(generation=generation, projection=projection, filtered=filtered): + builder = table.new_read_builder() + if filtered: + builder.with_filter(builder.new_predicate_builder().equal('payload', b'other')) + builder.with_limit(1) + builder.with_projection(projection) + splits = builder.new_scan().plan().splits() + self.assertTrue(all(split.raw_convertible == (generation == 0) for split in splits)) + wanted = [{name: row[name] for name in projection} + for row in (expected[1:] if filtered else expected)] + read = builder.new_read() + self.assertEqual(read.to_arrow(splits).to_pylist(), wanted) + batches = read.to_arrow_batch_reader(splits) + self.assertEqual([row for batch in batches for row in batch.to_pylist()], wanted) + actual = [] + for row in read.to_iterator(splits): + actual.append({name: row.get_blob(i).to_data() if name == 'payload' else row.get_field(i) + for i, name in enumerate(projection)}) + self.assertEqual(actual, wanted) + if not filtered: + authorized = [QueryAuthSplit(split, TableQueryAuthResult( + None, {'payload': '{"name":"NULL"}'})) for split in splits] + masked = [{name: None if name == 'payload' else value for name, value in row.items()} + for row in wanted] + self.assertEqual(read.to_arrow(authorized).to_pylist(), masked) + actual = [] + for row in read.to_iterator(authorized): + actual.append({name: row.get_field(i) for i, name in enumerate(projection)}) + self.assertEqual(actual, masked) + + def test_python_pk_write_registers_blobref_and_reads_payload(self): + table = self._create_table('pk_blob', { + 'bucket': '1', + 'write.native.enabled': 'false', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1, 1, 2], + 'payload': [b'old', b'new', b'other'], + 'items': [[b'a', None], [b'a', b'b'], []], + 'mapping': [[('k', b'x')], [('k', b'y'), ('m', None)], []], + }, schema=arrow_schema) + + writer, commit = self._write(table) + try: + writer.write_arrow(data) + messages = writer.prepare_commit() + files = [meta for message in messages for meta in message.new_files] + self.assertTrue(files) + self.assertTrue(all( + any(name.endswith('.blobref') for name in meta.extra_files) + and not any(name.endswith('.managed.blob') for name in meta.extra_files) + for meta in files)) + + payloads = [] + referenced = [] + for meta in files: + payloads.extend(self._parquet_rows(table, meta)) + for name in meta.extra_files: + if not name.endswith('.blobref'): + continue + sidecar = meta.aligned_file_path(name) + references = ManagedBlobReferenceFile.read(table.file_io, sidecar) + referenced.extend(references) + for reference in references: + pack = self._pack_path(table, reference) + self.assertTrue(table.file_io.exists(pack)) + self.assertTrue(pack.endswith('.managed.blob')) + self.assertNotIn(reference.relative_path, meta.extra_files) + + kept = next(row for row in payloads if row['id'] == 1) + descriptor = BlobDescriptor.deserialize(kept['payload']) + self.assertTrue(descriptor.uri.startswith(str(table.table_path) + '/')) + self.assertTrue(table.file_io.exists(descriptor.uri)) + self.assertIn(ManagedBlobReferenceFile.from_descriptor_uri(descriptor.uri), referenced) + self.assertEqual(self._written_payload(table, kept['payload']), b'new') + self.assertEqual(self._written_payload(table, kept['items'][0]), b'a') + self.assertEqual(self._written_payload(table, kept['items'][1]), b'b') + map_values = dict(kept['mapping']) + self.assertEqual(self._written_payload(table, map_values['k']), b'y') + self.assertIsNone(map_values['m']) + self.assertTrue(referenced) + self.assertNotIn(b'old', [row['payload'] for row in payloads]) + + commit.abort(messages) + for meta in files: + self.assertFalse(table.file_io.exists(meta.physical_path())) + for name in meta.extra_files: + self.assertFalse(table.file_io.exists(meta.aligned_file_path(name))) + for reference in referenced: + self.assertFalse(table.file_io.exists(self._pack_path(table, reference))) + finally: + writer.close() + commit.close() + + def test_retract_does_not_record_managed_payload(self): + table = self._create_table('pk_blob_retract', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'rowkind.field': 'op', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1, 2], + 'payload': [b'keep', b'drop'], + 'items': [[b'keep-item'], [b'drop-item']], + 'mapping': [[('k', b'keep-map')], [('k', b'drop-map')]], + 'op': ['+I', '-D'], + }, schema=arrow_schema) + + writer, commit = self._write(table) + try: + writer.write_arrow(data) + messages = writer.prepare_commit() + files = [meta for message in messages for meta in message.new_files] + rows = [] + blobrefs = [] + for meta in files: + rows.extend(self._parquet_rows(table, meta)) + for name in meta.extra_files: + if name.endswith('.blobref'): + blobrefs.extend(ManagedBlobReferenceFile.read( + table.file_io, meta.aligned_file_path(name))) + deleted = next(row for row in rows if row['id'] == 2) + self.assertIsNone(deleted['payload']) + self.assertIsNone(deleted['items']) + self.assertIsNone(deleted['mapping']) + kept = next(row for row in rows if row['id'] == 1) + self.assertEqual(self._written_payload(table, kept['payload']), b'keep') + pack_names = {reference.relative_path for reference in blobrefs} + self.assertTrue(pack_names) + for reference in blobrefs: + pack = self._pack_path(table, reference) + with table.file_io.new_input_stream(pack) as stream: + contents = stream.read() + self.assertNotIn(b'drop', contents) + finally: + writer.close() + commit.close() + + def test_writer_abort_deletes_uncommitted_pack_and_sidecar(self): + table = self._create_table('pk_blob_abort', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'target-file-size': '1 B', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1], + 'payload': [b'payload'], + 'items': [[b'item']], + 'mapping': [[('k', b'map')]], + }, schema=arrow_schema) + + writer = table.new_batch_write_builder().new_write() + try: + writer.write_arrow(data) + written = self._names_under(table) + self.assertTrue(any(name.endswith('.parquet') for name in written)) + self.assertTrue(any(name.endswith('.blobref') for name in written)) + self.assertTrue(any(name.endswith('.managed.blob') for name in written)) + writer.abort() + self.assertFalse(any( + name.endswith(('.parquet', '.blobref', '.managed.blob')) + for name in self._names_under(table))) + finally: + writer.close() + + def test_partial_abort_keeps_pack_used_by_committed_files(self): + self._check_shared_pack_cleanup(False) + + def test_relative_warehouse_abort_cleans_shared_packs(self): + previous = os.getcwd() + os.chdir(self.temp_dir) + try: + catalog = CatalogFactory.create({'warehouse': './wh %?#'}) + catalog.create_database('default', False) + with patch.object(self, 'catalog', catalog): + for action in ('writer', 'commit', 'shared'): + with self.subTest(action=action): + table = self._create_table('abort_' + action, { + 'bucket': '1', 'write.native.enabled': 'false', + 'target-file-size': '1 B', 'blob.target-file-size': '1 GB', + }) + writer, commit = self._write(table) + try: + for row_id in range(2): + writer.write_arrow(pa.Table.from_pydict({ + 'id': [row_id], 'payload': [b'payload'], + 'items': [[b'item']], 'mapping': [[('k', b'map')]], + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields))) + packs = {name for name in self._names_under(table) + if name.endswith('.managed.blob')} + self.assertEqual(len(packs), 3) + if action == 'writer': + writer.abort() + else: + messages = writer.prepare_commit() + if action == 'shared': + files = [meta for message in messages for meta in message.new_files] + self.assertEqual(len(files), 2) + commit.abort([CommitMessage((), 0, new_files=files[:1])]) + self.assertTrue(packs.issubset(set(self._names_under(table)))) + remaining = self._parquet_rows(table, files[1])[0] + self.assertEqual(self._written_payload(table, remaining['payload']), b'payload') + commit.abort([CommitMessage((), 0, new_files=files[1:])]) + else: + commit.abort(messages) + self.assertFalse(any(name.endswith(('.managed.blob', '.parquet', '.blobref')) + for name in self._names_under(table))) + finally: + writer.close() + commit.close() + self._check_shared_pack_cleanup(False) + finally: + os.chdir(previous) + + def test_iterator_preserves_payloads_with_reference_headers(self): + target = os.path.join(self.temp_dir, 'reference-target') + with open(target, 'wb') as stream: + stream.write(b'redirected') + payloads = [BlobViewStruct('db.user_payload', 7, 42).serialize(), + BlobDescriptor(target, 0, 10).serialize(), b'plain', None] + inputs = [] + for index, payload in enumerate(payloads): + if payload is None: + inputs.append(None) + continue + # A source file makes even descriptor-shaped input an opaque payload. + source = os.path.join(self.temp_dir, 'payload-%d' % index) + with open(source, 'wb') as stream: + stream.write(payload) + inputs.append(BlobDescriptor(source, 0, len(payload)).serialize()) + table = self._create_table('opaque_blobs', { + 'bucket': '1', 'write.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + }) + for merged in (False, True): + self._write_and_commit(table, pa.Table.from_pydict({ + 'id': list(range(4)), 'payload': inputs, + 'items': [[value] for value in inputs], + 'mapping': [[('k', value)] for value in inputs], + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields))) + for projection in (['id', 'payload'], ['payload', 'id'], ['payload'], + ["mapping['k']"], ['payload', "mapping['k']"], + {'alias': 'payload'}, {'alias': "mapping['k']"}): + with self.subTest(merged=merged, projection=projection): + builder = table.new_read_builder().with_projection(projection) + splits = builder.new_scan().plan().splits() + self.assertTrue(all(split.raw_convertible != merged for split in splits)) + read = builder.new_read() + result = read.to_arrow(splits) + blob_indices = [i for i, name in enumerate(result.schema.names) if name != 'id'] + for index in blob_indices: + self.assertEqual(result.column(index).to_pylist(), payloads) + with read.to_arrow_batch_reader(splits) as batches: + self.assertEqual(batches.read_all().to_pydict(), result.to_pydict()) + if isinstance(projection, dict): + continue # Named expressions are only supported by Arrow read APIs. + rows = [] + for row in read.to_iterator(splits): + rows.append([None if row.get_blob(i) is None else row.get_blob(i).to_data() + for i in blob_indices]) + self.assertEqual(rows, [[payload] * len(blob_indices) for payload in payloads]) + + builder = table.new_read_builder() + builder.with_filter(builder.new_predicate_builder().equal('payload', payloads[0])) + builder.with_projection(['payload']).with_limit(1) + splits = builder.new_scan().plan().splits() + for auth in (None, TableQueryAuthResult(None, {})): + selected = splits if auth is None else [QueryAuthSplit(split, auth) for split in splits] + self.assertEqual([row.get_blob(0).to_data() for row in builder.new_read().to_iterator(selected)], + [payloads[0]]) + + builder = table.copy({'blob-as-descriptor': 'true'}).new_read_builder() + builder.with_projection(['payload']) + splits = builder.new_scan().plan().splits() + self.assertEqual([None if row.get_blob(0) is None else row.get_blob(0).to_data() + for row in builder.new_read().to_iterator(splits)], payloads) + + def test_exclusive_file_cleanup_keeps_pack_used_by_committed_files(self): + self._check_shared_pack_cleanup(True) + + def _check_shared_pack_cleanup(self, exclusive_cleanup): + table = self._create_table('pk_blob_partial_abort', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + 'target-file-size': '1 B', + 'blob.target-file-size': '1 GB', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + writer, commit = self._write(table) + try: + for row_id in range(6): + writer.write_arrow(pa.Table.from_pydict({ + 'id': [row_id], + 'payload': [b'p%d' % row_id], + 'items': [[b'i%d' % row_id]], + 'mapping': [[('k', b'm%d' % row_id)]], + }, schema=arrow_schema)) + messages = writer.prepare_commit() + files = [meta for message in messages for meta in message.new_files] + self.assertGreaterEqual(len(files), 2) + pack_lists = [ + reference.relative_path + for meta in files for reference in ManagedBlobReferenceFile.read( + table.file_io, meta.aligned_file_path(meta.file_name + '.blobref')) + ] + self.assertGreater(len(pack_lists), len(set(pack_lists))) + + message = messages[0] + dropped = files[-1] + dropped_ids = {row['id'] for row in self._parquet_rows(table, dropped)} + commit.commit([ + CommitMessage(message.partition, message.bucket, files[:-1])]) + if exclusive_cleanup: + # Java abort deletes collectFiles() without reference checks. + # The metadata must be safe for that exclusive-ownership contract. + for path in dropped.collect_files(): + table.file_io.delete_quietly(path) + else: + commit.abort([ + CommitMessage(message.partition, message.bucket, [dropped])]) + + for name in set(pack_lists): + self.assertTrue(table.file_io.exists(dropped.aligned_file_path(name))) + self.assertFalse(table.file_io.exists(dropped.physical_path())) + + builder = table.new_read_builder() + rows = builder.new_read().to_arrow( + builder.new_scan().plan().splits()).to_pylist() + self.assertEqual( + {row['id'] for row in rows}, + {row_id for row_id in range(6) if row_id not in dropped_ids}) + for row in rows: + self.assertEqual(row['payload'], b'p%d' % row['id']) + finally: + writer.close() + commit.close() + + def test_read_builder_returns_managed_blob_payload(self): + table = self._create_table('pk_blob_read', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1], + 'payload': [b'visible'], + 'items': [[b'item']], + 'mapping': [[('k', b'map')]], + }, schema=arrow_schema) + + writer, commit = self._write(table) + try: + writer.write_arrow(data) + commit.commit(writer.prepare_commit()) + builder = table.new_read_builder() + rows = builder.new_read().to_arrow( + builder.new_scan().plan().splits()).to_pylist() + self.assertEqual(len(rows), 1) + row = rows[0] + self.assertEqual(row['payload'], b'visible') + self.assertEqual(row['items'], [b'item']) + self.assertEqual(dict(row['mapping']), {'k': b'map'}) + finally: + writer.close() + commit.close() + + def test_payload_predicate_matches_resolved_bytes(self): + self._check_payload_predicate(False) + + def test_map_key_projection_resolves_payload(self): + self._check_map_key_projection(False) + + def test_merged_map_key_projection_resolves_payload(self): + self._check_map_key_projection(True) + + def _check_map_key_projection(self, merged): + table = self._create_table('pk_blob_map_projection', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }) + for payload in ([b'old', b'new'] if merged else [b'new']): + self._write_and_commit(table, pa.Table.from_pydict({ + 'id': [1, 2, 3], + 'payload': [b'visible', None, None], + 'items': [None, None, None], + 'mapping': [[('k', payload), ('n', None), + ('k.a', b'dot'), ('k_a', b'underscore')], [], None], + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields))) + + projections = [ + (["mapping['k']"], {'mapping_k': [b'new', None, None]}), + ({'alias': "mapping['k']"}, {'alias': [b'new', None, None]}), + (["mapping['k']", 'payload'], { + 'mapping_k': [b'new', None, None], 'payload': [b'visible', None, None]}), + (["mapping['k.a']", "mapping['k_a']", "mapping['n']", "mapping['missing']"], { + 'mapping_k_a': [b'dot', None, None], + 'mapping_k_a__0': [b'underscore', None, None], + 'mapping_n': [None, None, None], 'mapping_missing': [None, None, None]}), + ] + for projection, expected in projections: + with self.subTest(projection=projection): + builder = table.new_read_builder().with_projection(projection) + splits = builder.new_scan().plan().splits() + self.assertTrue(splits) + self.assertTrue(all(split.raw_convertible != merged for split in splits)) + read = builder.new_read() + self.assertEqual(read.to_arrow(splits).to_pydict(), expected) + with read.to_arrow_batch_reader(splits) as reader: + self.assertEqual(reader.read_all().to_pydict(), expected) + if isinstance(projection, list): + rows = [tuple(row.get_field(i) for i in range(len(expected))) + for row in read.to_iterator(splits)] + self.assertEqual(rows, list(zip(*expected.values()))) + + builder = table.new_read_builder() + builder.with_filter(builder.new_predicate_builder().equal('payload', b'visible')) + builder.with_projection(["mapping['k']"]).with_limit(1) + splits = builder.new_scan().plan().splits() + self.assertEqual(builder.new_read().to_arrow(splits).to_pydict(), {'mapping_k': [b'new']}) + + builder = table.copy({'blob-as-descriptor': 'true'}).new_read_builder() + builder.with_projection(["mapping['k']"]) + rows = builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist() + self.assertEqual(self._written_payload(table, rows[0]['mapping_k']), b'new') + self.assertEqual(rows[1:], [{'mapping_k': None}, {'mapping_k': None}]) + + def test_merged_payload_predicate_matches_resolved_bytes(self): + self._check_payload_predicate(True) + + def _check_payload_predicate(self, merged): + table = self._create_table('pk_blob_filter', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1, 2], + 'payload': [b'visible', b'other'], + 'items': [[b'a'], [b'b']], + 'mapping': [[('k', b'a')], [('k', b'b')]], + }, schema=arrow_schema) + self._write_and_commit(table, data) + + if merged: + self._write_and_commit(table, data) + + for projection in (None, ['id'], ['payload', 'id'], ['id', 'payload'], []): + for limit in (None, 1): + for payload, expected_id in ((b'visible', 1), (b'other', 2), (b'missing', None)): + with self.subTest(projection=projection, limit=limit, payload=payload): + builder = table.new_read_builder() + predicate = builder.new_predicate_builder().equal('payload', payload) + builder.with_filter(predicate) + if projection is not None: + builder.with_projection(projection) + if limit is not None: + builder.with_limit(limit) + splits = builder.new_scan().plan().splits() + if expected_id is not None: + self.assertTrue(splits) + self.assertTrue(all(split.raw_convertible != merged for split in splits)) + read = builder.new_read() + expected_count = 0 if expected_id is None else 1 + result = read.to_arrow(splits) + with read.to_arrow_batch_reader(splits) as batch_reader: + self.assertEqual(batch_reader.read_all().num_rows, expected_count) + self.assertEqual(result.num_rows, expected_count) + if expected_count and 'id' in result.schema.names: + self.assertEqual(result.column('id').to_pylist(), [expected_id]) + if expected_count and 'payload' in result.schema.names: + self.assertEqual(result.column('payload').to_pylist(), [payload]) + names = projection if projection is not None else [f.name for f in table.fields] + rows = [tuple(row.get_field(i) for i in range(len(names))) + for row in read.to_iterator(splits)] + self.assertEqual(len(rows), expected_count) + if expected_count and 'id' in names: + self.assertEqual(rows[0][names.index('id')], expected_id) + if expected_count and 'payload' in names: + self.assertEqual(rows[0][names.index('payload')], payload) + + # Authorization masks must be applied before the user predicate, + # including when its payload column is absent from the output. + builder = table.new_read_builder() + splits = builder.new_scan().plan().splits() + builder.with_filter(builder.new_predicate_builder().is_null('payload')) + builder.with_projection(['id']).with_limit(1) + auth = TableQueryAuthResult(None, {'payload': '{"name":"NULL"}'}) + splits = [QueryAuthSplit(split, auth) for split in splits] + read = builder.new_read() + self.assertEqual(read.to_arrow(splits).num_rows, 1) + self.assertEqual(len(list(read.to_iterator(splits))), 1) + + # Authorized reads must also defer unmasked managed BLOB predicates + # until descriptors have been resolved, before applying LIMIT. + builder = table.new_read_builder() + builder.with_filter(builder.new_predicate_builder().equal('payload', b'other')) + builder.with_projection(['id']).with_limit(1) + auth = TableQueryAuthResult(None, {}) + splits = [QueryAuthSplit(split, auth) for split in builder.new_scan().plan().splits()] + read = builder.new_read() + self.assertEqual(read.to_arrow(splits).to_pylist(), [{'id': 2}]) + self.assertEqual([row.get_field(0) for row in read.to_iterator(splits)], [2]) + + def test_parallel_payload_filter_applies_limit_to_matches(self): + table = self._create_table('pk_blob_parallel_filter', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + 'target-file-size': '1 B', + 'source.split.target-size': '1 B', + }) + writer, commit = self._write(table) + try: + for row_id, payload in enumerate([b'other', b'visible', b'other', b'visible']): + writer.write_arrow(pa.Table.from_pydict({ + 'id': [row_id], 'payload': [payload], 'items': [None], 'mapping': [None], + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields))) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + for projection in (['id'], ['payload', 'id']): + builder = table.new_read_builder() + builder.with_filter(builder.new_predicate_builder().equal('payload', b'visible')) + builder.with_projection(projection).with_limit(1) + splits = builder.new_scan().plan().splits() + self.assertGreater(len(splits), 1) + result = builder.new_read().to_arrow(splits, parallelism=2) + self.assertEqual(result.num_rows, 1) + self.assertIn(result.column('id')[0].as_py(), [1, 3]) + + def test_payload_filter_preserves_large_batch(self): + table = self._create_table('pk_blob_large_filter', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + 'read.batch-size': '200000', + }) + count = 140000 + data = pa.Table.from_pydict({ + 'id': list(range(count)), + 'payload': [b'x'] * count, + 'items': [None] * count, + 'mapping': [None] * count, + }, schema=PyarrowFieldParser.from_paimon_schema(table.fields)) + self._write_and_commit(table, data) + builder = table.new_read_builder() + builder.with_filter(builder.new_predicate_builder().equal('payload', b'x')) + result = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + self.assertEqual(result.num_rows, count) + self.assertEqual(sorted(result.column('id').to_pylist()), list(range(count))) + + def test_row_iterator_opens_table_relative_blob(self): + table = self._create_table('pk_blob_iterator', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data = pa.Table.from_pydict({ + 'id': [1], + 'payload': [b'visible'], + 'items': [[b'item']], + 'mapping': [[('k', b'map')]], + }, schema=arrow_schema) + self._write_and_commit(table, data) + + builder = table.new_read_builder() + read = builder.new_read() + rows = list(read.to_iterator(builder.new_scan().plan().splits())) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0].get_blob(1).to_data(), b'visible') + + def test_mixed_blob_iterator_preserves_resolved_descriptor_fields(self): + self._check_mixed_blob_iterator(os.path.join(self.temp_dir, 'external.bin')) + + def test_mixed_blob_iterator_preserves_relative_external_uri(self): + previous = os.getcwd() + try: + os.chdir(self.temp_dir) + self._check_mixed_blob_iterator('external.bin') + finally: + os.chdir(previous) + + def _check_mixed_blob_iterator(self, external_path): + fields = [ + DataField(0, 'id', AtomicType('INT', False)), + DataField(1, 'managed', AtomicType('BLOB')), + DataField(2, 'external', AtomicType('BLOB')), + ] + schema = Schema(fields=fields, primary_keys=['id'], options={ + 'bucket': '1', + 'blob-descriptor-field': 'external', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }) + self.catalog.create_table('default.mixed_blobs', schema, False) + table = self.catalog.get_table('default.mixed_blobs') + external_payload = b'external payload' + with open(external_path, 'wb') as stream: + stream.write(external_payload) + descriptor = BlobDescriptor(external_path, 0, len(external_payload)).serialize() + for generation in range(2): + managed_payload = b'managed-%d' % generation + self._write_and_commit(table, pa.Table.from_pydict({ + 'id': [1], 'managed': [managed_payload], 'external': [descriptor], + }, schema=PyarrowFieldParser.from_paimon_schema(fields))) + for projection in (['id', 'managed', 'external'], ['external', 'managed'], ['id', 'external']): + with self.subTest(generation=generation, projection=projection): + builder = table.new_read_builder().with_projection(projection) + splits = builder.new_scan().plan().splits() + self.assertTrue(splits) + self.assertTrue(all(split.raw_convertible == (generation == 0) for split in splits)) + read = builder.new_read() + expected = {'id': 1, 'managed': managed_payload, 'external': external_payload} + self.assertEqual(read.to_arrow(splits).to_pylist(), [ + {name: expected[name] for name in projection}]) + rows = list(read.to_iterator(splits)) + self.assertEqual(len(rows), 1) + for index, name in enumerate(projection): + if name != 'id': + self.assertEqual(rows[0].get_blob(index).to_data(), expected[name]) + descriptor_builder = table.copy({'blob-as-descriptor': 'true'}).new_read_builder() + descriptor_builder.with_projection(projection) + descriptor_splits = descriptor_builder.new_scan().plan().splits() + rows = list(descriptor_builder.new_read().to_iterator(descriptor_splits)) + self.assertEqual(len(rows), 1) + for index, name in enumerate(projection): + if name != 'id': + self.assertEqual(rows[0].get_blob(index).to_data(), expected[name]) + + def test_inline_blob_fields_reject_raw_payloads(self): + fields = [ + DataField(0, 'id', AtomicType('INT', False)), + DataField(1, 'payload', AtomicType('BLOB')), + ] + for option, struct_name in (('blob-descriptor-field', 'BlobDescriptor'), + ('blob-view-field', 'BlobViewStruct')): + with self.subTest(option=option): + name = 'raw_inline_' + option.replace('-', '_') + self.catalog.create_table('default.' + name, Schema( + fields=fields, primary_keys=['id'], options={ + 'bucket': '1', option: 'payload', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }), False) + table = self.catalog.get_table('default.' + name) + writer, commit = self._write(table) + try: + with self.assertRaisesRegex(ValueError, '%s requires .* serialized %s' % (option, struct_name)): + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'payload': [b'not-a-reference'], + }, schema=PyarrowFieldParser.from_paimon_schema(fields))) + self.assertEqual(writer.prepare_commit(), []) + finally: + writer.close() + commit.close() + self.assertIsNone(table.snapshot_manager().get_latest_snapshot()) + self.assertFalse([name for name in self._names_under(table) + if not name.startswith(('schema-', '.'))]) + + def test_failed_prepare_retry_cleans_unreferenced_packs(self): + self._check_failed_prepare_finalization('prepare') + + def test_failed_prepare_close_cleans_unreferenced_packs(self): + self._check_failed_prepare_finalization('close') + + def test_failed_prepare_abort_cleans_unreferenced_packs(self): + self._check_failed_prepare_finalization('abort') + + def _check_failed_prepare_finalization(self, action): + for stage in ('_data_file_statistics', '_create_data_file_meta'): + for pack_size in ('1 B', '1 GB'): + for shared in (False, True): + with self.subTest(stage=stage, pack_size=pack_size, shared=shared): + name = 'failed_%s_%s_%s_%s' % (action, stage, pack_size.replace(' ', ''), shared) + table = self._create_table(name, { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + 'target-file-size': '1 B' if shared else '1 GB', + 'blob.target-file-size': pack_size, + }) + schema = PyarrowFieldParser.from_paimon_schema(table.fields) + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + if shared: + writer.write_arrow(pa.Table.from_pydict({ + 'id': [0], 'payload': [b'first'], 'items': [None], 'mapping': [None], + }, schema=schema)) + data_writer = next(iter(writer._execution.file_store_write.data_writers.values())) + data_writer.target_file_size = 1024 ** 3 + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'payload': [b'retried'], 'items': [None], 'mapping': [None], + }, schema=schema)) + data_writer = next(iter(writer._execution.file_store_write.data_writers.values())) + with patch.object(data_writer, stage, side_effect=OSError('injected')): + with self.assertRaisesRegex(OSError, 'injected'): + writer.prepare_commit(1) + if shared: + first = data_writer.committed_files[0] + references = ManagedBlobReferenceFile.read( + table.file_io, first.aligned_file_path(first.file_name + '.blobref')) + for reference in references: + self.assertTrue(table.file_io.exists(self._pack_path(table, reference))) + if pack_size == '1 GB': + self.assertEqual(sum(name.endswith('.managed.blob') + for name in self._names_under(table)), 1) + if action == 'abort': + writer.abort() + self.assertFalse(any( + name.endswith(('.parquet', '.blobref', '.managed.blob')) + for name in self._names_under(table))) + continue + if action == 'prepare': + messages = writer.prepare_commit(1) + files = [meta for message in messages for meta in message.new_files] + commit.commit(messages, 1) + writer.close() + else: + writer.close() + files = list(data_writer.committed_files) + references = set() + rows = [] + for meta in files: + rows.extend(self._parquet_rows(table, meta)) + for name in meta.extra_files: + if name.endswith('.blobref'): + references.update(reference.relative_path for reference in + ManagedBlobReferenceFile.read( + table.file_io, meta.aligned_file_path(name))) + packs = {name for name in self._names_under(table) + if name.endswith('.managed.blob')} + self.assertTrue(references) + self.assertEqual(packs, references) + expected = {0: b'first', 1: b'retried'} if shared else {1: b'retried'} + self.assertEqual({row['id']: self._written_payload(table, row['payload']) + for row in rows}, expected) + finally: + writer.close() + commit.close() + + def test_later_flush_failure_keeps_completed_pack(self): + from pypaimon.blob.primary_key_blob_externalizer import PrimaryKeyBlobExternalizer + table = self._create_table('pk_blob_flush_fail', { + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + 'target-file-size': '1 B', + 'blob.target-file-size': '1 GB', + }) + arrow_schema = PyarrowFieldParser.from_paimon_schema(table.fields) + writer, commit = self._write(table) + original_write = PrimaryKeyBlobExternalizer._write_descriptor + original_take = PrimaryKeyBlobExternalizer.take_referenced_pack_names + + def take(self): + names = original_take(self) + self._files_done = getattr(self, '_files_done', 0) + 1 + return names + + def fail_later(self, field_name, raw): + if getattr(self, '_files_done', 0) >= 1 and field_name == 'payload': + raise OSError('injected') + return original_write(self, field_name, raw) + + try: + PrimaryKeyBlobExternalizer.take_referenced_pack_names = take + PrimaryKeyBlobExternalizer._write_descriptor = fail_later + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], + 'payload': [b'first'], + 'items': [[b'a']], + 'mapping': [[('k', b'a')]], + }, schema=arrow_schema)) + packs = [name for name in self._names_under(table) if name.endswith('.managed.blob')] + self.assertTrue(packs) + data_writer = next(iter(writer._execution.file_store_write.data_writers.values())) + data_writer.target_file_size = 1024 * 1024 * 1024 + writer.write_arrow(pa.Table.from_pydict({ + 'id': [2], + 'payload': [b'second'], + 'items': [[b'b']], + 'mapping': [[('k', b'b')]], + }, schema=arrow_schema)) + with self.assertRaises(OSError): + writer.prepare_commit() + for name in packs: + self.assertTrue(any(found == name for found in self._names_under(table))) + finally: + PrimaryKeyBlobExternalizer.take_referenced_pack_names = original_take + PrimaryKeyBlobExternalizer._write_descriptor = original_write + writer.close() + commit.close() + + def test_relative_warehouse_publishes_absolute_pack_uri(self): + previous = os.getcwd() + os.chdir(self.temp_dir) + try: + catalog = CatalogFactory.create({'warehouse': './wh'}) + catalog.create_database('default', False) + fields = [ + DataField(0, 'id', AtomicType('INT', False)), + DataField(1, 'payload', AtomicType('BLOB')), + ] + catalog.create_table('default.rel_blob', Schema( + fields=fields, primary_keys=['id'], options={ + 'bucket': '1', + 'write.native.enabled': 'false', + 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', + }), False) + table = catalog.get_table('default.rel_blob') + messages = self._write_and_commit(table, pa.Table.from_pydict({ + 'id': [1], 'payload': [b'visible'], + }, schema=PyarrowFieldParser.from_paimon_schema(fields))) + builder = table.new_read_builder() + rows = builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist() + self.assertEqual(rows, [{'id': 1, 'payload': b'visible'}]) + data_file = messages[0].new_files[0] + descriptor = BlobDescriptor.deserialize(self._parquet_rows(table, data_file)[0]['payload']) + self.assertTrue(os.path.isabs(descriptor.uri)) + self.assertTrue(descriptor.uri.startswith(os.path.abspath('./wh') + '/')) + self.assertEqual(descriptor.uri.count('default.db'), 1) + self.assertTrue(table.file_io.exists(descriptor.uri)) + finally: + os.chdir(previous) + + def _create_table(self, name, options): + fields = [ + DataField(0, 'id', AtomicType('INT', False)), + DataField(1, 'payload', AtomicType('BLOB')), + DataField(2, 'items', ArrayType(True, AtomicType('BLOB'))), + DataField(3, 'mapping', MapType( + True, AtomicType('STRING', False), AtomicType('BLOB'))), + ] + if options.get('rowkind.field') == 'op': + fields.append(DataField(4, 'op', AtomicType('STRING'))) + schema = Schema(fields=fields, primary_keys=['id'], options=options) + self.catalog.create_table('default.' + name, schema, False) + return self.catalog.get_table('default.' + name) + + def _write(self, table): + builder = table.new_batch_write_builder() + return builder.new_write(), builder.new_commit() + + def _write_and_commit(self, table, data): + writer, commit = self._write(table) + try: + writer.write_arrow(data) + messages = writer.prepare_commit() + commit.commit(messages) + return messages + finally: + writer.close() + commit.close() + + def _parquet_rows(self, table, meta): + with table.file_io.new_input_stream(meta.physical_path()) as stream: + return pq.ParquetFile(stream).read().to_pylist() + + def _names_under(self, table): + names = [] + for dirpath, _, filenames in os.walk(str(table.table_path)): + names.extend(filenames) + return names + + def _written_payload(self, table, raw): + return materialize_managed_blob( + raw, 'scalar', table.file_io, table.table_path) + + def _pack_path(self, table, reference): + return resolve_table_blob_uri( + table.table_path, + "%s/%s" % (reference.storage_root_id, reference.relative_path)) diff --git a/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py b/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py index 4d985feccf6b..40c9dd984313 100644 --- a/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py +++ b/paimon-python/pypaimon/tests/managed_blob_write_ownership_test.py @@ -17,7 +17,7 @@ import unittest -from pypaimon.write.file_store_write import FileStoreWrite +from pypaimon.execution.python.write.file_store_write import FileStoreWrite class _KeepingWriter: @@ -40,9 +40,14 @@ def _release_prepared_files(self): def _file_store(writer, blob_consumer=None): file_store = object.__new__(FileStoreWrite) + file_store.table = type("_Table", (), {"is_primary_key_table": False})() file_store.data_writers = {((), 0): writer} file_store.commit_identifier = 0 + file_store.max_seq_numbers = {} file_store._runtime_total_buckets = {} + file_store._last_modified_commit_identifiers = {} + file_store._written_buckets = set() + file_store._prepared_writer_lifecycle = None file_store.blob_consumer = blob_consumer return file_store diff --git a/paimon-python/pypaimon/tests/py36/filter_record_batch_reader_test.py b/paimon-python/pypaimon/tests/py36/filter_record_batch_reader_test.py index 503ef4d0b1f8..a33e18a6a6a1 100644 --- a/paimon-python/pypaimon/tests/py36/filter_record_batch_reader_test.py +++ b/paimon-python/pypaimon/tests/py36/filter_record_batch_reader_test.py @@ -16,11 +16,14 @@ # under the License. import unittest +from unittest.mock import patch import pyarrow as pa from pypaimon.common.predicate import Predicate -from pypaimon.execution.python.read.reader.filter_record_batch_reader import (FilterRecordBatchReader) +from pypaimon.execution.python.read.reader.filter_record_batch_reader import ( + FilterRecordBatchReader, _record_batch_from_table) +from pypaimon.read.reader.managed_blob_record_batch_reader import _replace_batch_column from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.schema.data_types import AtomicType, DataField @@ -62,3 +65,36 @@ def test_row_filter_uses_logical_field_order_and_preserves_metadata(self): self.assertIsNone(reader.read_arrow_batch()) finally: reader.close() + + def test_large_filter_keeps_every_row(self): + count = 140000 + batch = pa.RecordBatch.from_arrays([pa.array(range(count))], names=['id']) + reader = FilterRecordBatchReader( + _Batches([batch]), Predicate('greaterOrEqual', 0, 'id', [0])) + result = reader.read_arrow_batch() + self.assertEqual(result.num_rows, count) + self.assertEqual(result.column(0).to_pylist(), list(range(count))) + self.assertIsNone(reader.read_arrow_batch()) + + def test_record_batch_from_table_without_concat_batches(self): + count = 140000 + half = count // 2 + table = pa.Table.from_batches([ + pa.RecordBatch.from_arrays([pa.array(range(half))], names=['id']), + pa.RecordBatch.from_arrays([pa.array(range(half, count))], names=['id']), + ]) + self.assertGreater(len(table.to_batches()), 1) + with patch.object(pa, 'concat_batches', None, create=True): + result = _record_batch_from_table(table) + self.assertEqual(result.num_rows, count) + self.assertEqual(result.column(0).to_pylist(), list(range(count))) + + def test_replace_record_batch_column_without_set_column(self): + field = pa.field('payload', pa.large_binary()) + batch = pa.RecordBatch.from_arrays( + [pa.array([1]), pa.array([b'old'], type=pa.large_binary())], + schema=pa.schema([pa.field('id', pa.int64()), field])) + updated = _replace_batch_column( + batch, 1, field, pa.array([b'new'], type=pa.large_binary())) + self.assertEqual(updated.column(0).to_pylist(), [1]) + self.assertEqual(updated.column(1).to_pylist(), [b'new']) diff --git a/paimon-python/pypaimon/tests/schema_evolution_cast_test.py b/paimon-python/pypaimon/tests/schema_evolution_cast_test.py index 992f456b272c..d3f635793887 100644 --- a/paimon-python/pypaimon/tests/schema_evolution_cast_test.py +++ b/paimon-python/pypaimon/tests/schema_evolution_cast_test.py @@ -25,9 +25,11 @@ import pyarrow as pa import pytest +from pypaimon.execution.python.casting.schema_evolution import ( + _cast_nested_container, _offsets_from_buffer, cast_array_for_schema_evolution) from pypaimon.schema.arrow_schema import read_arrow_schema -from pypaimon.execution.python.casting.schema_evolution import cast_array_for_schema_evolution -from pypaimon.schema.data_types import ArrayType, AtomicType, DataField, MapType, PyarrowFieldParser, RowType +from pypaimon.schema.data_types import ( + ArrayType, AtomicType, DataField, MapType, PyarrowFieldParser, RowType) @pytest.mark.parametrize('chunked', [False, True]) @@ -518,3 +520,31 @@ def limited_localtime(seconds=None): # Golden values from real Java CastExecutors, including a local midnight # whose DATE changes if a standard offset replaces the correct DST offset. assert result.to_pylist() == expected + + +def test_map_binary_values_rebuild_when_container_cast_is_missing(): + # PyArrow 6 rejects map -> map. + keys = pa.array(['k', 'n']) + items = pa.array([b'python-map', None], type=pa.binary()) + source = pa.MapArray.from_arrays( + pa.array([0, 2, 2], type=pa.int32()), keys, items, + mask=pa.array([False, True])) + target = pa.map_(pa.string(), pa.large_binary()) + rebuilt = _cast_nested_container(source, target) + assert rebuilt.type == target + assert rebuilt.to_pylist() == [[('k', b'python-map'), ('n', None)], None] + + elements = pa.array([b'python-item', None], type=pa.binary()) + listed = pa.ListArray.from_arrays(pa.array([0, 2], type=pa.int32()), elements) + rebuilt_list = _cast_nested_container(listed, pa.list_(pa.large_binary())) + assert rebuilt_list.to_pylist() == [[b'python-item', None]] + + wider = pa.MapArray.from_arrays( + pa.array([0, 1, 1, 3], type=pa.int32()), + pa.array(['k', 'm', 'n']), + pa.array([b'a', b'b', b'c'], type=pa.binary()), + mask=pa.array([False, True, False])) + sliced = wider.slice(1, 2) + assert _offsets_from_buffer(sliced).to_pylist() == [1, 1, 3] + rebuilt_slice = _cast_nested_container(sliced, target) + assert rebuilt_slice.to_pylist() == [None, [('m', b'b'), ('n', b'c')]] diff --git a/paimon-python/pypaimon/tests/schema_manager_test.py b/paimon-python/pypaimon/tests/schema_manager_test.py index dd73d9142d36..2e9b7889961c 100644 --- a/paimon-python/pypaimon/tests/schema_manager_test.py +++ b/paimon-python/pypaimon/tests/schema_manager_test.py @@ -101,6 +101,19 @@ def tearDown(self): import shutil shutil.rmtree(self.temp_dir, ignore_errors=True) + def test_primary_key_video_frame_field_is_rejected(self): + fields = [DataField(0, 'id', AtomicType('INT', False)), DataField(1, 'payload', AtomicType('BLOB'))] + manager = SchemaManager(self.file_io, self.temp_dir + '/video') + with self.assertRaisesRegex(ValueError, 'only supports append-only tables'): + manager.create_table(Schema(fields=fields, primary_keys=['id'], options={ + 'bucket': '1', 'video-frame-field': 'payload'})) + self.assertIsNone(manager.latest()) + manager.create_table(Schema(fields=fields, primary_keys=['id'], options={'bucket': '1'})) + with self.assertRaisesRegex(RuntimeError, 'only supports append-only tables'): + manager.commit_changes([SetOption('video-frame-field', 'payload')]) + self.assertEqual(manager.latest().id, 0) + self.assertNotIn('video-frame-field', manager.latest().options) + @staticmethod def _schema(blob_type): return Schema( @@ -212,7 +225,7 @@ def test_create_table_rejects_unsupported_nested_blob(self): ): manager.create_table(schema) - def test_create_table_still_rejects_primary_key_map_blob(self): + def test_create_table_accepts_primary_key_map_blob(self): manager = SchemaManager( self.file_io, f"{self.temp_dir}/test_db.db/pk_map_blob", @@ -231,19 +244,129 @@ def test_create_table_still_rejects_primary_key_map_blob(self): ), ], primary_keys=["id"], - options={ - "row-tracking.enabled": "true", - "data-evolution.enabled": "true", - }, + options={"bucket": "1"}, ) - with self.assertRaisesRegex( - ValueError, - "MAP type is not supported with primary key", - ): + manager.create_table(schema) + self.assertEqual(manager.latest().primary_keys, ["id"]) + + def test_create_table_rejects_managed_blob_primary_key(self): + manager = SchemaManager( + self.file_io, + f"{self.temp_dir}/test_db.db/pk_blob_key", + ) + schema = Schema( + fields=[ + DataField(0, "id", AtomicType("INT", False)), + DataField(1, "payload", AtomicType("BLOB")), + ], + primary_keys=["payload", "id"], + options={"bucket": "1"}, + ) + + with self.assertRaisesRegex(ValueError, "cannot be primary keys"): manager.create_table(schema) +class TestPartialUpdateBlobValidation(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.file_io = FileIO.get(self.temp_dir, Options({})) + + def tearDown(self): + import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) + + @staticmethod + def _schema(options, payload_type=None): + return Schema(fields=[ + DataField(0, 'id', AtomicType('INT', False)), + DataField(1, 'payload', payload_type or AtomicType('BLOB')), + DataField(2, 'ts', AtomicType('INT')), + ], primary_keys=['id'], options=dict({'bucket': '1', 'merge-engine': 'partial-update'}, **options)) + + def test_rejects_blob_sequence_group_ordering(self): + types = [AtomicType('BLOB'), ArrayType(True, AtomicType('BLOB')), + MapType(True, AtomicType('STRING', False), AtomicType('BLOB'))] + for index, payload_type in enumerate(types): + for ordering in ('payload', 'ts,payload'): + with self.subTest(payload_type=payload_type, ordering=ordering): + manager = SchemaManager(self.file_io, '%s/ordering_%s_%s' % ( + self.temp_dir, index, ordering)) + options = {'fields.%s.sequence-group' % ordering: 'ts'} + with self.assertRaisesRegex(ValueError, 'sequence-group ordering field'): + manager.create_table(self._schema(options, payload_type)) + self.assertIsNone(manager.latest()) + for inline in ('blob-descriptor-field', 'blob-view-field'): + manager = SchemaManager(self.file_io, self.temp_dir + '/' + inline) + options = {inline: 'payload', 'fields.payload.sequence-group': 'ts'} + with self.assertRaisesRegex(ValueError, 'sequence-group ordering field'): + manager.create_table(self._schema(options)) + self.assertIsNone(manager.latest()) + + def test_managed_blob_aggregate_retract_validation(self): + for key in ('fields.payload.aggregate-function', 'fields.default-aggregate-function'): + for index, extra in enumerate(({}, {'fields.payload.ignore-retract': 'true'}, + {'partial-update.ignore-delete': 'true'}, + {key: 'last_value'}, + {'blob-descriptor-field': 'payload'})): + with self.subTest(key=key, extra=extra): + manager = SchemaManager(self.file_io, '%s/%s_%s' % (self.temp_dir, key, index)) + options = {'fields.ts.sequence-group': 'payload', key: 'last_non_null_value'} + options.update(extra) + if extra: + manager.create_table(self._schema(options)) + self.assertIsNotNone(manager.latest()) + else: + with self.assertRaisesRegex(ValueError, 'fields.payload.ignore-retract'): + manager.create_table(self._schema(options)) + self.assertIsNone(manager.latest()) + + def test_alter_rejects_invalid_blob_options_without_persisting(self): + cases = [ + ({}, {'fields.payload.sequence-group': 'ts'}, 'sequence-group ordering field'), + ({'fields.ts.sequence-group': 'payload'}, + {'fields.payload.aggregate-function': 'last_non_null_value'}, 'ignore-retract'), + ({'fields.ts.sequence-group': 'payload', 'fields.default-aggregate-function': 'last_non_null_value', + 'fields.payload.ignore-retract': 'true'}, + {'fields.payload.ignore-retract': 'false'}, 'ignore-retract'), + ] + for index, (options, changes, message) in enumerate(cases): + with self.subTest(options=options, changes=changes): + manager = SchemaManager(self.file_io, '%s/alter_%s' % (self.temp_dir, index)) + manager.create_table(self._schema(options)) + original = manager.latest().options.copy() + with self.assertRaisesRegex(RuntimeError, message): + manager.commit_changes([SetOption(key, value) for key, value in changes.items()]) + self.assertEqual(manager.latest().id, 0) + self.assertEqual(manager.latest().options, original) + + def test_managed_blob_ignore_retract_requires_java_boolean(self): + for index, value in enumerate(('yes', '1', ' true ', 'on', '0', 'no', 'false ', 'TRUE', 'TrUe')): + with self.subTest(value=value): + manager = SchemaManager(self.file_io, '%s/boolean_%s' % (self.temp_dir, index)) + options = {'fields.ts.sequence-group': 'payload', + 'fields.payload.aggregate-function': 'last_non_null_value', + 'fields.payload.ignore-retract': value} + if value.lower() == 'true': + manager.create_table(self._schema(options)) + self.assertIsNotNone(manager.latest()) + else: + with self.assertRaisesRegex(ValueError, 'fields.payload.ignore-retract'): + manager.create_table(self._schema(options)) + self.assertIsNone(manager.latest()) + + def test_accepts_unprotected_and_inactive_sequence_group_options(self): + for index, options in enumerate(( + {'fields.payload.aggregate-function': 'last_non_null_value'}, + {'fields.ts.sequence-group': 'payload'}, + {'merge-engine': 'deduplicate', 'fields.payload.sequence-group': 'ts'})): + manager = SchemaManager(self.file_io, '%s/valid_%s' % (self.temp_dir, index)) + manager.create_table(self._schema(options)) + self.assertIsNotNone(manager.latest()) + + class TestOptionValidation(unittest.TestCase): def setUp(self): diff --git a/paimon-python/pypaimon/utils/arrow_utils.py b/paimon-python/pypaimon/utils/arrow_utils.py index 77e1064813e7..a3d4afe94ea9 100644 --- a/paimon-python/pypaimon/utils/arrow_utils.py +++ b/paimon-python/pypaimon/utils/arrow_utils.py @@ -18,6 +18,34 @@ import pyarrow as pa +_ROW_KIND = b'paimon.row-kind' + + +def row_kind_field(schema): + """Give internal row kinds a marked column that cannot shadow user data.""" + name = '_row_kind' + while name in schema.names: + name = '_' + name + return pa.field(name, pa.string(), metadata={_ROW_KIND: b'true'}) + + +def row_kind_column_index(schema): + for index, field in enumerate(schema): + if _ROW_KIND in (field.metadata or {}): + return index + if _ROW_KIND in (schema.metadata or {}): + return -1 + # Older batch readers use the conventional transport column without a marker. + return schema.get_field_index('_row_kind') + + +def with_no_row_kind(schema): + """Mark all columns as data, including a user column named _row_kind.""" + metadata = dict(schema.metadata or {}) + metadata[_ROW_KIND] = b'none' + return schema.with_metadata(metadata) + + def zero_column_batch(num_rows: int, metadata=None) -> pa.RecordBatch: """Build a zero-column batch without losing its logical row count.""" empty_struct = pa.Array.from_buffers( diff --git a/paimon-python/pypaimon/utils/path.py b/paimon-python/pypaimon/utils/path.py index 6987353bb093..2769df0ef242 100644 --- a/paimon-python/pypaimon/utils/path.py +++ b/paimon-python/pypaimon/utils/path.py @@ -115,3 +115,20 @@ def to_file_io_path(path): if _WINDOWS and _has_drive(local_path): local_path = local_path.lstrip('/') return local_path + + +def path_identity(path, file_io): + """Compare local paths independently of qualification, preserving remote URIs.""" + path = str(path) + scheme, authority, local_path = _parts(path) + if scheme not in ('', 'file:') or authority not in ('', '//localhost'): + return path + if scheme == 'file:': + # Match the FileIO that will read/delete the path. Only URI inputs are + # decoded; percent characters in plain local names remain literal. + local_path = file_io.to_filesystem_path(path) + if local_path.startswith('file:'): + raise ValueError('FileIO cannot resolve local file URI: %s' % path) + if _WINDOWS and _has_drive(local_path): + local_path = local_path.lstrip('/') + return os.path.abspath(local_path).replace(os.sep, '/') diff --git a/paimon-python/pypaimon/write/row_utils.py b/paimon-python/pypaimon/write/row_utils.py index 6127543caf6c..d03d3b66431c 100644 --- a/paimon-python/pypaimon/write/row_utils.py +++ b/paimon-python/pypaimon/write/row_utils.py @@ -90,6 +90,58 @@ def inline_blob_value(value: Any, descriptor: bool, view: bool) -> Any: return value +def _inline_reference_bytes(value: Any, option: str, struct_name: str) -> bytes: + if hasattr(value, 'as_py'): + value = value.as_py() + if isinstance(value, str): + value = value.encode('utf-8') + if not isinstance(value, (bytes, bytearray)): + raise ValueError( + "%s requires blob field value to be a serialized %s." % (option, struct_name)) + return bytes(value) + + +def validate_inline_blob_fields(data: pa.RecordBatch, descriptor_fields, view_fields): + """Reject raw payloads in fields that store serialized BLOB references inline.""" + from pypaimon.table.row.blob import BlobDescriptor, BlobViewStruct + + for field_name in descriptor_fields: + if field_name not in data.schema.names: + continue + for value in data.column(data.schema.get_field_index(field_name)).to_pylist(): + if value is None: + continue + descriptor_bytes = _inline_reference_bytes( + value, 'blob-descriptor-field', 'BlobDescriptor') + # Like Java, schema-declared descriptors use the prefix parser. + # Exact wire length is only needed when detecting descriptors + # among arbitrary payload bytes. + try: + BlobDescriptor.deserialize(descriptor_bytes) + except Exception as e: + raise ValueError( + "blob-descriptor-field requires blob field value to be a serialized " + "BlobDescriptor." + ) from e + + for field_name in view_fields: + if field_name not in data.schema.names: + continue + for value in data.column(data.schema.get_field_index(field_name)).to_pylist(): + if value is None: + continue + view_bytes = _inline_reference_bytes(value, 'blob-view-field', 'BlobViewStruct') + try: + view_struct = BlobViewStruct.deserialize(view_bytes) + if view_struct.serialize() != view_bytes: + raise ValueError("BlobViewStruct payload contains trailing bytes.") + except Exception as e: + raise ValueError( + "blob-view-field requires blob field value to be a serialized " + "BlobViewStruct." + ) from e + + def value_for_arrow(value: Any, field: DataField) -> Any: if isinstance(value, Vector): return value.to_list()