[python] Cache file-format metadata across reads

docs/docs/pypaimon/pytorch.md

@@ -59,6 +59,26 @@ when it is false, it will read the full amount of data into memory.
 
 **`prefetch_concurrency`** (default: 1): When streaming is true, number of threads used for parallel prefetch within each DataLoader worker. Set to a value greater than 1 to partition splits across threads and increase read throughput. Has no effect when streaming is false.
 
+## File Format Metadata Cache
+
+Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated
+size limit in the catalog options:
+
+```python
+catalog = CatalogFactory.create({
+    "warehouse": "file:///path/to/warehouse",
+    "file-format.metadata-cache.max-size": "50 mb",
+})
+table = catalog.get_table("database.table")
+read_builder = table.new_read_builder()
+```
+
+The default limit is 50 MB; set it to `0 b` to disable and clear the cache. The
+cache is local to each process and benefits workers reused with
+`DataLoader(..., persistent_workers=True)`. The cache uses a conservative
+per-entry memory estimate and an internal entry-count safeguard; actual native
+PyArrow memory may still be higher. The cache assumes immutable Paimon data files.
+
 ## Shuffle
 
 PyPaimon supports streaming shuffle for PyTorch `IterableDataset`. The shuffle

paimon-python/pypaimon/common/options/config.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from pypaimon.common.memory_size import MemorySize
 from pypaimon.common.options.config_options import ConfigOptions
 
 
@@ -78,6 +79,16 @@ class CatalogOptions:
     METASTORE = ConfigOptions.key("metastore").string_type().default_value("filesystem").with_description(
         "Metastore type")
     WAREHOUSE = ConfigOptions.key("warehouse").string_type().no_default_value().with_description("Warehouse path")
+    FILE_FORMAT_METADATA_CACHE_MAX_SIZE = (
+        ConfigOptions.key("file-format.metadata-cache.max-size")
+        .memory_type()
+        .default_value(MemorySize.of_mebi_bytes(50))
+        .with_description(
+            "Maximum estimated size of reusable PyArrow Dataset metadata "
+            "cached in the current process. Set to 0 to disable and clear "
+            "the cache."
+        )
+    )
     TOKEN_PROVIDER = ConfigOptions.key("token.provider").string_type().no_default_value().with_description(
         "Token provider")
     TOKEN = ConfigOptions.key("token").string_type().no_default_value().with_description("Authentication token")

paimon-python/pypaimon/read/reader/format_pyarrow_reader.py

@@ -15,14 +15,19 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from collections import deque
-from typing import Any, Deque, Dict, Iterator, List, Optional, Set, Tuple
+import os
+import sys
+import threading
+from collections import OrderedDict, deque
+from concurrent.futures import Future
+from typing import Any, Callable, Deque, Dict, Iterator, List, Optional, Set, Tuple
 
 import pyarrow as pa
 import pyarrow.dataset as ds
 from pyarrow import RecordBatch
 
 from pypaimon.common.file_io import FileIO
+from pypaimon.common.options.config import CatalogOptions
 from pypaimon.common.options.core_options import CoreOptions
 from pypaimon.data.variant_shredding import (
     VariantSchema,
@@ -43,6 +48,194 @@
 from pypaimon.table.special_fields import SpecialFields
 
 
+_DEFAULT_FILE_FORMAT_METADATA_CACHE_MAX_SIZE = 50 * 1024 * 1024
+_FILE_FORMAT_METADATA_CACHE_MAX_ENTRIES = 4096
+_FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE = 8 * 1024
+_FILE_FORMAT_METADATA_CACHE_CONTAINER_OVERHEAD = 256
+
+
+class _FilesystemIdentity:
+    def __init__(self, filesystem):
+        self.filesystem = filesystem
+
+    def __hash__(self):
+        return id(self.filesystem)
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, _FilesystemIdentity)
+            and self.filesystem is other.filesystem
+        )
+
+
+class _FileFormatDatasetCache:
+    def __init__(
+            self,
+            max_size: int,
+            max_entries: int = _FILE_FORMAT_METADATA_CACHE_MAX_ENTRIES):
+        self.max_size = max_size
+        self.max_entries = max_entries
+        self.estimated_size = 0
+        self._entries = OrderedDict()
+        self._loads = {}
+        self._lock = threading.Lock()
+
+    def get_or_load(self, key: Tuple[Any, str, str], loader: Callable[[], Any],
+                    size_estimator: Callable[[Any], Optional[int]]):
+        with self._lock:
+            entry = self._entries.get(key)
+            if entry is not None:
+                self._entries.move_to_end(key)
+                return entry[0]
+
+            future = self._loads.get(key)
+            if future is None:
+                future = Future()
+                self._loads[key] = future
+                should_load = True
+            else:
+                should_load = False
+
+        if not should_load:
+            return future.result()
+
+        try:
+            dataset = loader()
+            estimated_size = size_estimator(dataset)
+        except BaseException as exception:
+            future.set_exception(exception)
+            with self._lock:
+                self._loads.pop(key, None)
+            raise
+
+        with self._lock:
+            if estimated_size is not None:
+                estimated_size = max(1, estimated_size)
+                self._entries[key] = (dataset, estimated_size)
+                self.estimated_size += estimated_size
+                self._entries.move_to_end(key)
+                self._evict()
+        future.set_result(dataset)
+        with self._lock:
+            self._loads.pop(key, None)
+        return dataset
+
+    def resize(self, max_size: int):
+        with self._lock:
+            self.max_size = max_size
+            self._evict()
+
+    def _evict(self):
+        while (
+                self.estimated_size > self.max_size
+                or len(self._entries) > self.max_entries):
+            _, (_, evicted_size) = self._entries.popitem(last=False)
+            self.estimated_size -= evicted_size
+
+
+_FILE_FORMAT_DATASET_CACHE = None
+_FILE_FORMAT_DATASET_CACHE_LOCK = threading.Lock()
+_FILE_FORMAT_DATASET_CACHE_PID = os.getpid()
+
+
+def _ensure_file_format_dataset_cache_process():
+    global _FILE_FORMAT_DATASET_CACHE
+    global _FILE_FORMAT_DATASET_CACHE_LOCK
+    global _FILE_FORMAT_DATASET_CACHE_PID
+    current_pid = os.getpid()
+    if current_pid != _FILE_FORMAT_DATASET_CACHE_PID:
+        _FILE_FORMAT_DATASET_CACHE = None
+        _FILE_FORMAT_DATASET_CACHE_LOCK = threading.Lock()
+        _FILE_FORMAT_DATASET_CACHE_PID = current_pid
+
+
+def _file_format_dataset_cache(max_size: int) -> _FileFormatDatasetCache:
+    global _FILE_FORMAT_DATASET_CACHE
+    _ensure_file_format_dataset_cache_process()
+    with _FILE_FORMAT_DATASET_CACHE_LOCK:
+        if _FILE_FORMAT_DATASET_CACHE is None:
+            _FILE_FORMAT_DATASET_CACHE = _FileFormatDatasetCache(max_size)
+        else:
+            _FILE_FORMAT_DATASET_CACHE.resize(max_size)
+        return _FILE_FORMAT_DATASET_CACHE
+
+
+def _reset_file_format_dataset_cache():
+    global _FILE_FORMAT_DATASET_CACHE
+    _ensure_file_format_dataset_cache_process()
+    with _FILE_FORMAT_DATASET_CACHE_LOCK:
+        _FILE_FORMAT_DATASET_CACHE = None
+
+
+def _estimate_file_format_dataset_size(dataset, file_format: str) -> Optional[int]:
+    try:
+        if file_format == 'parquet':
+            footer_size = 0
+            for fragment in dataset.get_fragments():
+                metadata = fragment.metadata
+                if metadata is not None:
+                    footer_size += int(metadata.serialized_size)
+            if footer_size > 0:
+                return footer_size
+        return int(dataset.schema.serialize().size)
+    except Exception:
+        return None
+
+
+def _estimate_file_format_cache_entry_size(
+        key: Tuple[Any, str, str],
+        dataset,
+        file_format: str) -> Optional[int]:
+    metadata_size = _estimate_file_format_dataset_size(dataset, file_format)
+    if metadata_size is None:
+        return None
+
+    # PyArrow does not expose the native size retained by Dataset and Fragment
+    # objects. Account for all visible Python objects and apply a conservative
+    # floor so tiny files cannot turn a byte-bounded cache into an effectively
+    # unbounded object cache.
+    visible_size = (
+        metadata_size
+        + sys.getsizeof(key)
+        + sys.getsizeof(key[0])
+        + sys.getsizeof(key[1])
+        + sys.getsizeof(key[2])
+        + sys.getsizeof(dataset)
+        + sys.getsizeof((dataset, metadata_size))
+        + _FILE_FORMAT_METADATA_CACHE_CONTAINER_OVERHEAD
+    )
+    return max(_FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE, visible_size)
+
+
+def _file_format_metadata_cache_max_size(file_io: FileIO) -> int:
+    properties = getattr(file_io, 'properties', None)
+    if properties is None:
+        return _DEFAULT_FILE_FORMAT_METADATA_CACHE_MAX_SIZE
+    return properties.get(
+        CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE).get_bytes()
+
+
+def _file_format_dataset(file_io: FileIO, file_format: str, file_path: str,
+                         cache_max_size: int):
+    file_path_for_pyarrow = file_io.to_filesystem_path(file_path)
+    filesystem = file_io.filesystem
+
+    def load():
+        return ds.dataset(
+            file_path_for_pyarrow, format=file_format, filesystem=filesystem)
+
+    key = (_FilesystemIdentity(filesystem), file_format, file_path_for_pyarrow)
+    if cache_max_size <= 0:
+        _reset_file_format_dataset_cache()
+        return load()
+
+    return _file_format_dataset_cache(cache_max_size).get_or_load(
+        key,
+        load,
+        lambda dataset: _estimate_file_format_cache_entry_size(
+            key, dataset, file_format))
+
+
 class FormatPyArrowReader(RecordBatchReader):
     """
     A Format Reader that reads record batch from a Parquet or ORC file using PyArrow,
@@ -63,7 +256,9 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str,
                  row_ranges: Optional[List[Tuple[int, int]]] = None):
         self._predicate_field_names = predicate_field_names or set()
         file_path_for_pyarrow = file_io.to_filesystem_path(file_path)
-        self.dataset = ds.dataset(file_path_for_pyarrow, format=file_format, filesystem=file_io.filesystem)
+        cache_max_size = _file_format_metadata_cache_max_size(file_io)
+        self.dataset = _file_format_dataset(
+            file_io, file_format, file_path, cache_max_size)
         self._range_slicer = None
         self._selected_parquet_row_groups = None
         self._exhausted = False