From 50bde33502af3c123ee6210099245846a96e4f52 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 21:25:52 +0800 Subject: [PATCH] [python] Cache file-format metadata across reads Cache reusable PyArrow Dataset metadata across repeated reads in the same process. Configure the cache through catalog options and bound it with conservative byte accounting plus an entry-count safeguard. --- docs/docs/pypaimon/pytorch.md | 20 + .../pypaimon/common/options/config.py | 11 + .../read/reader/format_pyarrow_reader.py | 201 +++++++- .../tests/parquet_metadata_cache_test.py | 479 ++++++++++++++++++ 4 files changed, 708 insertions(+), 3 deletions(-) create mode 100644 paimon-python/pypaimon/tests/parquet_metadata_cache_test.py diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 0f7f7bbdef6d..6ab0af5173ca 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/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 diff --git a/paimon-python/pypaimon/common/options/config.py b/paimon-python/pypaimon/common/options/config.py index 604d62911a8f..2573782c424f 100644 --- a/paimon-python/pypaimon/common/options/config.py +++ b/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") diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index 565ec1f5b186..c84a1383add7 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/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 diff --git a/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py b/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py new file mode 100644 index 000000000000..2ca91f577690 --- /dev/null +++ b/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py @@ -0,0 +1,479 @@ +# 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 gc +import os +import tempfile +import threading +import time +import unittest +import weakref +from concurrent.futures import Future, ThreadPoolExecutor +from unittest.mock import patch + +import pyarrow as pa +import pyarrow.fs as pafs +import pyarrow.parquet as pq +from fsspec.implementations.local import LocalFileSystem as FsspecLocalFileSystem + +from pypaimon.common.options import Options +from pypaimon.common.options.config import CatalogOptions +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.filesystem.local_file_io import LocalFileIO +from pypaimon.read.reader import format_pyarrow_reader as reader_module +from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader +from pypaimon.schema.data_types import AtomicType, DataField + + +DEFAULT_CACHE_SIZE = 50 * 1024 * 1024 + + +class _CountingInputFile: + def __init__(self, wrapped, file_system): + self._wrapped = wrapped + self._file_system = file_system + + def read(self, size=-1): + offset = self._wrapped.tell() + data = self._wrapped.read(size) + self._file_system.reads.append((offset, len(data))) + return data + + def readinto(self, buffer): + offset = self._wrapped.tell() + size = self._wrapped.readinto(buffer) + self._file_system.reads.append((offset, size)) + return size + + def __getattr__(self, name): + return getattr(self._wrapped, name) + + +class _CountingLocalFileSystem(FsspecLocalFileSystem): + def __init__(self): + super().__init__() + self.opens = 0 + self.reads = [] + + def _open(self, path, mode="rb", **kwargs): + wrapped = super()._open(path, mode=mode, **kwargs) + if "r" not in mode: + return wrapped + self.opens += 1 + return _CountingInputFile(wrapped, self) + + def reset_counts(self): + self.opens = 0 + self.reads = [] + + +class FileFormatMetadataCacheTest(unittest.TestCase): + def setUp(self): + reader_module._reset_file_format_dataset_cache() + self.temp_dir = tempfile.TemporaryDirectory() + self.file_io = LocalFileIO(self.temp_dir.name, Options({})) + self.paths = [] + for index in range(3): + path = os.path.join(self.temp_dir.name, "data-{}.parquet".format(index)) + pq.write_table( + pa.table({"value": list(range(index * 10, index * 10 + 10))}), + path, + row_group_size=2, + ) + self.paths.append(path) + + def tearDown(self): + reader_module._reset_file_format_dataset_cache() + self.temp_dir.cleanup() + + def _file_io(self, max_size="50 mb"): + return LocalFileIO(self.temp_dir.name, Options({ + "file-format.metadata-cache.max-size": max_size, + })) + + def _read(self, path, file_io=None, options=None): + reader = FormatPyArrowReader( + file_io or self.file_io, + "parquet", + path, + [DataField(0, "value", AtomicType("BIGINT"))], + None, + options=options, + ) + values = [] + try: + while True: + batch = reader.read_arrow_batch() + if batch is None: + return values + values.extend(batch.column(0).to_pylist()) + finally: + reader.close() + + def test_enabled_by_default(self): + self.assertEqual( + DEFAULT_CACHE_SIZE, + self.file_io.properties.get( + CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE).get_bytes()) + + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + self._read(self.paths[0]) + self._read(self.paths[0]) + self.assertEqual(1, dataset.call_count) + + def test_zero_size_bypasses_and_removes_entry(self): + enabled = self._file_io() + disabled = self._file_io("0 b") + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + self._read(self.paths[0], enabled) + self._read(self.paths[0], disabled) + self._read(self.paths[0], enabled) + self.assertEqual(3, dataset.call_count) + + def test_zero_size_clears_other_entries(self): + enabled = self._file_io() + disabled = self._file_io("0 b") + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + self._read(self.paths[0], enabled) + self._read(self.paths[1], enabled) + self._read(self.paths[0], disabled) + self._read(self.paths[1], enabled) + self.assertEqual(4, dataset.call_count) + + def test_reuses_dataset(self): + file_io = self._file_io() + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + first = self._read(self.paths[0], file_io) + second = self._read(self.paths[0], file_io) + + self.assertEqual(list(range(10)), first) + self.assertEqual(first, second) + self.assertEqual(1, dataset.call_count) + + def test_repeated_scan_skips_footer_io(self): + path = os.path.join(self.temp_dir.name, "footer-io.parquet") + pq.write_table( + pa.table({ + "value": list(range(10000)), + "payload": ["x" * 100] * 10000, + }), + path, + row_group_size=100, + compression="none", + ) + counting = _CountingLocalFileSystem() + file_io = self._file_io("0 b") + file_io.filesystem = pafs.PyFileSystem(pafs.FSSpecHandler(counting)) + + uncached = self._read(path, file_io) + uncached_opens = counting.opens + uncached_reads = len(counting.reads) + + counting.reset_counts() + reader_module._reset_file_format_dataset_cache() + file_io.properties.set( + CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE, "50 mb") + self._read(path, file_io) + counting.reset_counts() + cached = self._read(path, file_io) + + self.assertEqual(uncached, cached) + self.assertLess(counting.opens, uncached_opens) + self.assertLess(len(counting.reads), uncached_reads) + + def test_evicts_least_recently_used_entry_by_estimated_size(self): + cache = reader_module._FileFormatDatasetCache(10) + first_key = (None, "parquet", "first") + second_key = (None, "parquet", "second") + third_key = (None, "parquet", "third") + + cache.get_or_load(first_key, lambda: "first", lambda _: 4) + cache.get_or_load(second_key, lambda: "second", lambda _: 4) + cache.get_or_load(first_key, lambda: "unused", lambda _: 4) + cache.get_or_load(third_key, lambda: "third", lambda _: 4) + + self.assertEqual([first_key, third_key], list(cache._entries.keys())) + self.assertEqual(8, cache.estimated_size) + + def test_evicts_least_recently_used_entry_by_entry_count(self): + cache = reader_module._FileFormatDatasetCache( + 1024, max_entries=2) + first_key = (None, "parquet", "first") + second_key = (None, "parquet", "second") + third_key = (None, "parquet", "third") + + cache.get_or_load(first_key, lambda: "first", lambda _: 1) + cache.get_or_load(second_key, lambda: "second", lambda _: 1) + cache.get_or_load(third_key, lambda: "third", lambda _: 1) + + self.assertEqual( + [second_key, third_key], list(cache._entries.keys())) + self.assertEqual(2, cache.estimated_size) + + def test_does_not_retain_entry_larger_than_size_limit(self): + cache = reader_module._FileFormatDatasetCache(5) + loads = [] + key = (None, "parquet", "large") + + def load(): + loads.append(True) + return "large" + + self.assertEqual( + "large", cache.get_or_load(key, load, lambda _: 6)) + self.assertEqual( + "large", cache.get_or_load(key, load, lambda _: 6)) + self.assertEqual(2, len(loads)) + self.assertEqual(0, len(cache._entries)) + self.assertEqual(0, cache.estimated_size) + + def test_does_not_retain_entry_without_size_estimate(self): + cache = reader_module._FileFormatDatasetCache(10) + key = (None, "unknown", "data") + loads = [] + + def load(): + loads.append(True) + return "unknown" + + self.assertEqual( + "unknown", cache.get_or_load(key, load, lambda _: None)) + self.assertEqual( + "unknown", cache.get_or_load(key, load, lambda _: None)) + self.assertEqual(2, len(loads)) + self.assertEqual(0, len(cache._entries)) + + def test_coalesces_load_while_uncached_result_completes(self): + cache = reader_module._FileFormatDatasetCache(10) + key = (None, "unknown", "data") + setting_result = threading.Event() + waiter_started = threading.Event() + release_result = threading.Event() + loads = [] + future_count = [] + + def new_future(): + future = Future() + future_count.append(future) + if len(future_count) == 1: + original_set_result = future.set_result + original_result = future.result + + def delayed_set_result(result): + setting_result.set() + release_result.wait() + original_set_result(result) + + def observed_result(*args, **kwargs): + waiter_started.set() + return original_result(*args, **kwargs) + + future.set_result = delayed_set_result + future.result = observed_result + return future + + def load(): + loads.append(True) + return "unknown" + + with patch.object(reader_module, "Future", side_effect=new_future): + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit( + cache.get_or_load, key, load, lambda _: None) + self.assertTrue(setting_result.wait(1)) + second = executor.submit( + cache.get_or_load, key, load, lambda _: None) + waited = waiter_started.wait(1) + release_result.set() + + self.assertTrue(waited) + self.assertEqual("unknown", first.result()) + self.assertEqual("unknown", second.result()) + self.assertEqual(1, len(loads)) + + def test_estimates_serialized_parquet_footer_size(self): + dataset = reader_module.ds.dataset(self.paths[0], format="parquet") + expected = sum( + fragment.metadata.serialized_size + for fragment in dataset.get_fragments() + ) + self.assertGreater(expected, 0) + self.assertEqual( + expected, + reader_module._estimate_file_format_dataset_size( + dataset, "parquet")) + + def test_cache_entry_estimate_has_conservative_floor(self): + dataset = reader_module.ds.dataset( + self.paths[0], format="parquet") + key = ( + reader_module._FilesystemIdentity(self.file_io.filesystem), + "parquet", + self.paths[0], + ) + + estimated = reader_module._estimate_file_format_cache_entry_size( + key, dataset, "parquet") + + self.assertGreaterEqual( + estimated, + reader_module._FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE) + + def test_process_cache_can_shrink_requested_capacity(self): + cache = reader_module._file_format_dataset_cache(10) + cache.get_or_load(("first", "parquet", "first"), + lambda: "first", lambda _: 4) + cache.get_or_load(("second", "parquet", "second"), + lambda: "second", lambda _: 4) + same_cache = reader_module._file_format_dataset_cache(5) + + self.assertIs(cache, same_cache) + self.assertEqual(5, cache.max_size) + self.assertEqual( + [("second", "parquet", "second")], + list(cache._entries.keys())) + self.assertEqual(4, cache.estimated_size) + + def test_table_option_does_not_configure_process_cache(self): + table_options = CoreOptions(Options({ + "file-format.metadata-cache.max-size": "0 b", + })) + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + self._read(self.paths[0], options=table_options) + self._read(self.paths[0], options=table_options) + self.assertEqual(1, dataset.call_count) + + def test_shares_cache_across_file_io_with_same_filesystem(self): + other_file_io = LocalFileIO(self.temp_dir.name, Options({})) + other_file_io.filesystem = self.file_io.filesystem + + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + reader_module._file_format_dataset( + self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) + reader_module._file_format_dataset( + other_file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) + + self.assertEqual(1, dataset.call_count) + + def test_does_not_share_across_filesystems(self): + other_file_io = LocalFileIO(self.temp_dir.name, Options({})) + original = reader_module.ds.dataset + with patch.object(reader_module.ds, "dataset", wraps=original) as dataset: + reader_module._file_format_dataset( + self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) + reader_module._file_format_dataset( + other_file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) + self.assertEqual(2, dataset.call_count) + + def test_does_not_share_across_file_formats(self): + parquet_dataset = object() + orc_dataset = object() + with patch.object( + reader_module.ds, "dataset", + side_effect=[parquet_dataset, orc_dataset]) as dataset: + with patch.object( + reader_module, "_estimate_file_format_dataset_size", + return_value=1): + first = reader_module._file_format_dataset( + self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) + second = reader_module._file_format_dataset( + self.file_io, "orc", self.paths[0], DEFAULT_CACHE_SIZE) + + self.assertIs(parquet_dataset, first) + self.assertIs(orc_dataset, second) + self.assertEqual(2, dataset.call_count) + + def test_cache_key_retains_filesystem_wrapper(self): + root = pafs.LocalFileSystem() + filesystem = pafs.SubTreeFileSystem(self.temp_dir.name, root) + filesystem_ref = weakref.ref(filesystem) + file_io = LocalFileIO(self.temp_dir.name, Options({})) + file_io.filesystem = filesystem + + reader_module._file_format_dataset( + file_io, "parquet", os.path.basename(self.paths[0]), + DEFAULT_CACHE_SIZE) + file_io.filesystem = root + del filesystem + gc.collect() + + self.assertIsNotNone(filesystem_ref()) + + def test_filesystem_hash_collision_does_not_share_dataset(self): + first_dir = tempfile.TemporaryDirectory() + second_dir = tempfile.TemporaryDirectory() + try: + file_name = "same.parquet" + pq.write_table( + pa.table({"value": [1]}), os.path.join(first_dir.name, file_name)) + pq.write_table( + pa.table({"value": [2]}), os.path.join(second_dir.name, file_name)) + file_io = LocalFileIO(first_dir.name, Options({})) + first_filesystem = pafs.SubTreeFileSystem( + first_dir.name, pafs.LocalFileSystem()) + second_filesystem = pafs.SubTreeFileSystem( + second_dir.name, pafs.LocalFileSystem()) + + with patch.object( + reader_module._FilesystemIdentity, "__hash__", return_value=1): + file_io.filesystem = first_filesystem + first = reader_module._file_format_dataset( + file_io, "parquet", file_name, + DEFAULT_CACHE_SIZE).to_table() + file_io.filesystem = second_filesystem + second = reader_module._file_format_dataset( + file_io, "parquet", file_name, + DEFAULT_CACHE_SIZE).to_table() + + self.assertEqual([1], first.column("value").to_pylist()) + self.assertEqual([2], second.column("value").to_pylist()) + finally: + first_dir.cleanup() + second_dir.cleanup() + + def test_resets_after_process_change(self): + parent_cache = reader_module._file_format_dataset_cache(DEFAULT_CACHE_SIZE) + with patch.object(reader_module.os, "getpid", return_value=os.getpid() + 1): + child_cache = reader_module._file_format_dataset_cache(DEFAULT_CACHE_SIZE) + self.assertIsNot(parent_cache, child_cache) + + def test_coalesces_concurrent_loads(self): + original = reader_module.ds.dataset + + def delayed_dataset(*args, **kwargs): + time.sleep(0.05) + return original(*args, **kwargs) + + with patch.object( + reader_module.ds, "dataset", side_effect=delayed_dataset) as dataset: + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map( + lambda _: self._read(self.paths[0]), + range(8), + )) + + self.assertEqual(1, dataset.call_count) + self.assertTrue(all(value == list(range(10)) for value in results)) + + +if __name__ == "__main__": + unittest.main()