diff --git a/paimon-python/pypaimon/manifest/manifest_file_manager.py b/paimon-python/pypaimon/manifest/manifest_file_manager.py index 71a5a149e623..af2dec5388c7 100644 --- a/paimon-python/pypaimon/manifest/manifest_file_manager.py +++ b/paimon-python/pypaimon/manifest/manifest_file_manager.py @@ -21,6 +21,14 @@ import fastavro +try: + from fastavro._read import (read_bytes as _read_bytes, + read_long as _read_long, + read_record as _read_record, + skip_record as _skip_record) +except ImportError: # pragma: no cover - supported fastavro versions provide these + _read_bytes = _read_long = _read_record = _skip_record = None + from datetime import datetime from pypaimon.manifest.schema.data_file_meta import DataFileMeta @@ -34,6 +42,68 @@ from pypaimon.table.row.binary_row import BinaryRow +_MANIFEST_FIELD_NAMES = ( + '_VERSION', '_KIND', '_PARTITION', '_BUCKET', '_TOTAL_BUCKETS', '_FILE') +_MANIFEST_PREFIX_TYPES = ('int', 'int', 'bytes', 'int', 'int') + + +def _read_manifest_records(buffer, early_entry_filter, partition_filter, + partition_fields): + """Yield manifest records, skipping ``_FILE`` before decoding when possible.""" + if ((early_entry_filter is None and partition_filter is None) + or _read_record is None): + for record in fastavro.reader(buffer): + yield record, None, False + return + + blocks = fastavro.block_reader(buffer) + fields = blocks.writer_schema.get('fields', []) + if (tuple(field.get('name') for field in fields) != _MANIFEST_FIELD_NAMES + or tuple(field.get('type') for field in fields[:5]) + != _MANIFEST_PREFIX_TYPES + or not isinstance(fields[5].get('type'), dict) + or fields[5]['type'].get('type') != 'record'): + # Avro values follow writer field order. Keep the generic reader for + # historical manifests whose top-level fields were reordered. + buffer.seek(0) + for record in fastavro.reader(buffer): + yield record, None, False + return + + file_schema = fields[5]['type'] + named_schemas = getattr(blocks, '_named_schemas', {}) + + for block in blocks: + stream = block.bytes_ + for _ in range(block.num_records): + version = _read_long(stream) + kind = _read_long(stream) + partition_bytes = _read_bytes(stream) + bucket = _read_long(stream) + total_buckets = _read_long(stream) + if (early_entry_filter is not None + and not early_entry_filter(bucket, total_buckets)): + _skip_record(stream, file_schema, named_schemas) + continue + + partition = None + if partition_filter is not None: + partition = GenericRowDeserializer.from_bytes( + partition_bytes, partition_fields) + if not partition_filter.test(partition): + _skip_record(stream, file_schema, named_schemas) + continue + + yield { + '_VERSION': version, + '_KIND': kind, + '_PARTITION': partition_bytes, + '_BUCKET': bucket, + '_TOTAL_BUCKETS': total_buckets, + '_FILE': _read_record(stream, file_schema, named_schemas), + }, partition, True + + class ManifestFileManager: """Writer for manifest files in Avro format using unified FileIO.""" @@ -78,14 +148,23 @@ def _entry_identifier(e: ManifestEntry) -> tuple: deleted_entry_keys = set() added_entries = [] - with ThreadPoolExecutor(max_workers=max_workers) as executor: - future_results = executor.map(_process_single_manifest, manifest_files) - for entries in future_results: - for entry in entries: - if entry.kind == 0: # ADD - added_entries.append(entry) - else: # DELETE - deleted_entry_keys.add(_entry_identifier(entry)) + + def _collect(entries): + for entry in entries: + if entry.kind == 0: # ADD + added_entries.append(entry) + else: # DELETE + deleted_entry_keys.add(_entry_identifier(entry)) + + if len(manifest_files) == 1: + # Avoid executor overhead and keep native block decoding on the + # caller thread for the common single-manifest case. + _collect(_process_single_manifest(manifest_files[0])) + else: + with ThreadPoolExecutor(max_workers=max_workers) as executor: + for entries in executor.map( + _process_single_manifest, manifest_files): + _collect(entries) final_entries = [ entry for entry in added_entries @@ -113,10 +192,12 @@ def read(self, manifest_file_name: str, manifest_entry_filter=None, drop_stats=T with self.file_io.new_input_stream(manifest_file_path) as input_stream: avro_bytes = input_stream.read() buffer = BytesIO(avro_bytes) - reader = fastavro.reader(buffer) + records = _read_manifest_records( + buffer, early_entry_filter, partition_filter, + self.partition_keys_fields) - for record in reader: - if early_entry_filter is not None: + for record, partition, prefix_filtered in records: + if not prefix_filtered and early_entry_filter is not None: try: bucket = record['_BUCKET'] total_buckets = record['_TOTAL_BUCKETS'] @@ -127,8 +208,7 @@ def read(self, manifest_file_name: str, manifest_entry_filter=None, drop_stats=T continue if early_record_filter is not None and not early_record_filter(record): continue - partition = None - if partition_filter is not None: + if partition_filter is not None and partition is None: partition = GenericRowDeserializer.from_bytes( record['_PARTITION'], self.partition_keys_fields) if not partition_filter.test(partition): diff --git a/paimon-python/pypaimon/tests/manifest/manifest_manager_test.py b/paimon-python/pypaimon/tests/manifest/manifest_manager_test.py index f1def7a19387..c415f20a92b9 100644 --- a/paimon-python/pypaimon/tests/manifest/manifest_manager_test.py +++ b/paimon-python/pypaimon/tests/manifest/manifest_manager_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import copy import os import shutil import subprocess @@ -24,6 +25,7 @@ import unittest from dataclasses import replace from io import BytesIO +from unittest import mock import fastavro import pyarrow as pa @@ -36,7 +38,8 @@ from pypaimon.manifest.manifest_file_manager import ManifestFileManager from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.manifest.schema.data_file_meta import DataFileMeta -from pypaimon.manifest.schema.manifest_entry import ManifestEntry +from pypaimon.manifest.schema.manifest_entry import (MANIFEST_ENTRY_SCHEMA, + ManifestEntry) from pypaimon.manifest.schema.manifest_file_meta import ManifestFileMeta from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.schema.data_types import AtomicType, DataField @@ -273,6 +276,118 @@ def _create_manifest_entry(self, file_name, bucket=0): ) return entry + def _partitioned_manifest(self): + table_id = 'default.selective_manifest' + identifier = Identifier.from_string(table_id) + path = self.catalog.get_table_path(identifier) + if self.catalog.file_io.exists(path): + self.catalog.file_io.delete(path, recursive=True) + schema = Schema.from_pyarrow_schema( + pa.schema([('pt', pa.string()), ('id', pa.int32())]), + partition_keys=['pt']) + self.catalog.create_table(table_id, schema, False) + table = self.catalog.get_table(table_id) + manager = ManifestFileManager(table) + entries = [] + for name, partition, bucket in [ + ('selected.parquet', 'keep', 1), + ('partition-pruned.parquet', 'drop', 1), + ('bucket-pruned.parquet', 'keep', 2)]: + entry = self._create_manifest_entry(name, bucket) + entry.total_buckets = 4 + entry.partition = GenericRow( + [partition], table.partition_keys_fields) + entries.append(entry) + return table, manager, entries + + def test_partition_and_bucket_filters_skip_file_decoding(self): + import pypaimon.manifest.manifest_file_manager as manager_module + + _, manager, entries = self._partitioned_manifest() + name = 'selective-manifest.avro' + manager.write(name, entries) + + class KeepPartition: + @staticmethod + def test(row): + return row.get_field(0) == 'keep' + + with mock.patch.object( + manager_module, '_read_record', + wraps=manager_module._read_record) as read_record, \ + mock.patch.object( + manager_module, '_skip_record', + wraps=manager_module._skip_record) as skip_record: + actual = manager.read( + name, + early_entry_filter=lambda bucket, _: bucket == 1, + partition_filter=KeepPartition()) + + self.assertEqual([entry.file.file_name for entry in actual], + ['selected.parquet']) + self.assertEqual(read_record.call_count, 1) + self.assertEqual(skip_record.call_count, 2) + + def test_single_manifest_read_does_not_create_executor(self): + import pypaimon.manifest.manifest_file_manager as manager_module + + _, manager, entries = self._partitioned_manifest() + manifests = manager.rolling_write( + entries, 1024 * 1024, 'single-manifest') + + with mock.patch.object( + manager_module, 'ThreadPoolExecutor') as executor: + actual = manager.read_entries_parallel( + manifests, + early_entry_filter=lambda bucket, _: bucket == 1) + + executor.assert_not_called() + self.assertEqual( + [entry.file.file_name for entry in actual], + ['selected.parquet', 'partition-pruned.parquet']) + + def test_reordered_manifest_fields_use_compatible_reader(self): + import pypaimon.manifest.manifest_file_manager as manager_module + + table, manager, entries = self._partitioned_manifest() + schema = copy.deepcopy(MANIFEST_ENTRY_SCHEMA) + fields = schema['fields'] + schema['fields'] = [fields[1], fields[0]] + fields[2:] + buffer = BytesIO() + fastavro.writer(buffer, schema, manager._to_avro_records(entries)) + name = 'reordered-manifest.avro' + path = '{}/{}'.format(manager.manifest_path, name) + with table.file_io.new_output_stream(path) as output: + output.write(buffer.getvalue()) + + with mock.patch.object( + manager_module, '_read_record', + wraps=manager_module._read_record) as read_record: + actual = manager.read( + name, early_entry_filter=lambda bucket, _: bucket == 1) + + self.assertEqual([entry.file.file_name for entry in actual], + ['selected.parquet', 'partition-pruned.parquet']) + self.assertEqual(read_record.call_count, 0) + + def test_file_schema_named_type_reference(self): + table, manager, entries = self._partitioned_manifest() + schema = copy.deepcopy(MANIFEST_ENTRY_SCHEMA) + file_fields = schema['fields'][5]['type']['fields'] + file_fields[6]['type'] = 'record_KEY_STATS' + buffer = BytesIO() + fastavro.writer(buffer, schema, manager._to_avro_records(entries)) + name = 'named-type-manifest.avro' + path = '{}/{}'.format(manager.manifest_path, name) + with table.file_io.new_output_stream(path) as output: + output.write(buffer.getvalue()) + + actual = manager.read( + name, early_entry_filter=lambda bucket, _: bucket == 1) + + self.assertEqual([entry.file.file_name for entry in actual], + ['selected.parquet', 'partition-pruned.parquet']) + def test_manifest_bucket_and_level_stats(self): manager = self._make_manager() entries = [self._create_manifest_entry('a', bucket=2),