[python] Add block-based manifest Avro reader

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):

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),