From d970daec5db7ca99c3a32579e508e61efe01e93a Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 17 Sep 2026 20:55:41 -0700 Subject: [PATCH] [python] Reuse Parquet fragment metadata for single-file reads --- .../read/reader/format_pyarrow_reader.py | 7 + .../tests/parquet_metadata_cache_test.py | 130 +++++++++++++++--- 2 files changed, 118 insertions(+), 19 deletions(-) diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index 33999fec75f8..ead8dd3bdf93 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -279,6 +279,13 @@ def _file_format_dataset(file_io: FileIO, file_format: str, file_path: str, filesystem = file_io.filesystem def load(): + if file_format == 'parquet': + parquet_format = ds.ParquetFileFormat() + fragment = parquet_format.make_fragment( + file_path_for_pyarrow, filesystem=filesystem) + # Reuse this fragment's footer for schema discovery and scanning. + return ds.FileSystemDataset( + [fragment], fragment.physical_schema, parquet_format, filesystem) return ds.dataset( file_path_for_pyarrow, format=file_format, filesystem=filesystem) diff --git a/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py b/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py index 257560ef630e..12cead96faeb 100644 --- a/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py +++ b/paimon-python/pypaimon/tests/parquet_metadata_cache_test.py @@ -80,6 +80,22 @@ def reset_counts(self): self.reads = [] +class _CountingFileSystemHandler(pafs.FSSpecHandler): + """Count Arrow calls, not backend HTTP requests.""" + + def __init__(self): + super().__init__(FsspecLocalFileSystem(skip_instance_cache=True)) + self.calls = [] + + def get_file_info(self, paths): + self.calls.append("get_file_info") + return super().get_file_info(paths) + + def open_input_file(self, path): + self.calls.append("open_input_file") + return super().open_input_file(path) + + class FileFormatMetadataCacheTest(unittest.TestCase): def setUp(self): reader_module._reset_file_format_dataset_cache() @@ -129,8 +145,8 @@ def test_enabled_by_default(self): 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: self._read(self.paths[0]) self._read(self.paths[0]) self.assertEqual(1, dataset.call_count) @@ -138,8 +154,8 @@ def test_enabled_by_default(self): 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: self._read(self.paths[0], enabled) self._read(self.paths[0], disabled) self._read(self.paths[0], enabled) @@ -148,8 +164,8 @@ def test_zero_size_bypasses_and_removes_entry(self): 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: self._read(self.paths[0], enabled) self._read(self.paths[1], enabled) self._read(self.paths[0], disabled) @@ -158,8 +174,8 @@ def test_zero_size_clears_other_entries(self): 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: first = self._read(self.paths[0], file_io) second = self._read(self.paths[0], file_io) @@ -198,6 +214,78 @@ def test_repeated_scan_skips_footer_io(self): self.assertLess(counting.opens, uncached_opens) self.assertLess(len(counting.reads), uncached_reads) + def test_single_file_request_counts(self): + for max_size, expected_opens in [ + (DEFAULT_CACHE_SIZE, [2, 1]), (0, [2, 2]), (1, [2, 2])]: + with self.subTest(cache_max_size=max_size): + reader_module._reset_file_format_dataset_cache() + handler = _CountingFileSystemHandler() + self.file_io.filesystem = pafs.PyFileSystem(handler) + results = [] + for opens in expected_opens: + handler.calls.clear() + dataset = reader_module._file_format_dataset( + self.file_io, "parquet", self.paths[0], max_size) + results.append(dataset.to_table().to_pydict()) + self.assertEqual(0, handler.calls.count("get_file_info")) + self.assertEqual(opens, handler.calls.count("open_input_file")) + self.assertEqual({"value": list(range(10))}, results[0]) + self.assertEqual(results[0], results[1]) + + def test_fragment_metadata_is_reused_without_io(self): + handler = _CountingFileSystemHandler() + self.file_io.filesystem = pafs.PyFileSystem(handler) + dataset = reader_module._file_format_dataset( + self.file_io, "parquet", self.paths[0], 0) + self.assertEqual(["open_input_file"], handler.calls) + handler.calls.clear() + for _ in range(2): + fragment = next(dataset.get_fragments()) + self.assertEqual(dataset.schema, fragment.physical_schema) + self.assertEqual(10, fragment.metadata.num_rows) + self.assertEqual(5, len(fragment.split_by_row_group())) + self.assertGreater( + reader_module._estimate_file_format_dataset_size( + dataset, "parquet"), 0) + self.assertEqual([], handler.calls) + self.assertEqual(list(range(10)), dataset.to_table()[0].to_pylist()) + self.assertEqual(["open_input_file"], handler.calls) + + def test_concurrent_scans_load_footer_once(self): + handler = _CountingFileSystemHandler() + self.file_io.filesystem = pafs.PyFileSystem(handler) + barrier = threading.Barrier(8) + + def read(_): + barrier.wait(timeout=10) + return self._read(self.paths[0]) + + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(read, range(8))) + self.assertTrue(all(result == list(range(10)) for result in results)) + self.assertEqual(0, handler.calls.count("get_file_info")) + self.assertEqual(9, handler.calls.count("open_input_file")) + + def test_missing_and_corrupt_files_fail_and_can_retry(self): + for max_size in [0, DEFAULT_CACHE_SIZE]: + for corrupt in [False, True]: + with self.subTest(cache_max_size=max_size, corrupt=corrupt): + reader_module._reset_file_format_dataset_cache() + path = os.path.join(self.temp_dir.name, "invalid.parquet") + if corrupt: + with open(path, "wb") as output: + output.write(b"not a parquet file") + elif os.path.exists(path): + os.remove(path) + with self.assertRaises((OSError, pa.ArrowInvalid)): + reader_module._file_format_dataset( + self.file_io, "parquet", path, max_size).to_table() + pq.write_table(pa.table({"value": [42]}), path) + self.assertEqual( + [42], reader_module._file_format_dataset( + self.file_io, "parquet", path, max_size + ).to_table()[0].to_pylist()) + def test_evicts_least_recently_used_entry_by_estimated_size(self): cache = reader_module._FileFormatDatasetCache(10) first_key = (None, "parquet", "first") @@ -357,8 +445,8 @@ 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", 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) @@ -367,8 +455,8 @@ 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: reader_module._file_format_dataset( self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) reader_module._file_format_dataset( @@ -378,8 +466,8 @@ def test_shares_cache_across_file_io_with_same_filesystem(self): 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: + original = reader_module.ds.FileSystemDataset + with patch.object(reader_module.ds, "FileSystemDataset", wraps=original) as dataset: reader_module._file_format_dataset( self.file_io, "parquet", self.paths[0], DEFAULT_CACHE_SIZE) reader_module._file_format_dataset( @@ -390,8 +478,10 @@ 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: + reader_module.ds, "FileSystemDataset", + return_value=parquet_dataset) as parquet_loader, patch.object( + reader_module.ds, "dataset", + return_value=orc_dataset) as orc_loader: with patch.object( reader_module, "_estimate_file_format_dataset_size", return_value=1): @@ -402,7 +492,9 @@ def test_does_not_share_across_file_formats(self): self.assertIs(parquet_dataset, first) self.assertIs(orc_dataset, second) - self.assertEqual(2, dataset.call_count) + self.assertEqual(1, parquet_loader.call_count) + orc_loader.assert_called_once_with( + self.paths[0], format="orc", filesystem=self.file_io.filesystem) def test_cache_key_retains_filesystem_wrapper(self): root = pafs.LocalFileSystem() @@ -459,14 +551,14 @@ def test_resets_after_process_change(self): self.assertIsNot(parent_cache, child_cache) def test_coalesces_concurrent_loads(self): - original = reader_module.ds.dataset + original = reader_module.ds.FileSystemDataset 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: + reader_module.ds, "FileSystemDataset", side_effect=delayed_dataset) as dataset: with ThreadPoolExecutor(max_workers=8) as executor: results = list(executor.map( lambda _: self._read(self.paths[0]),