This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 9da018c6b9 [python] Support structured-data local cache selection
(#10225)
9da018c6b9 is described below
commit 9da018c6b95413ce5c5c2dbf1a287709ef70a16f
Author: XiaoHongbo <[email protected]>
AuthorDate: Tue Oct 27 13:14:40 2026 +0800
[python] Support structured-data local cache selection (#10225)
---
docs/docs/program-api/file-cache.mdx | 6 +-
.../pypaimon/common/options/core_options.py | 4 +-
.../python/read/reader/format_avro_reader.py | 8 +-
.../python/read/reader/format_pyarrow_reader.py | 31 ++-
.../pypaimon/filesystem/caching_file_io.py | 31 ++-
.../pypaimon/filesystem/caching_file_system.py | 78 ++++++++
.../pypaimon/tests/caching_file_io_test.py | 71 +++++++
paimon-python/pypaimon/tests/file_type_test.py | 32 ++-
.../pypaimon/tests/parquet_local_cache_test.py | 217 +++++++++++++++++++++
paimon-python/pypaimon/utils/file_type.py | 22 ++-
10 files changed, 468 insertions(+), 32 deletions(-)
diff --git a/docs/docs/program-api/file-cache.mdx
b/docs/docs/program-api/file-cache.mdx
index ad6926ee4a..fb187a768b 100644
--- a/docs/docs/program-api/file-cache.mdx
+++ b/docs/docs/program-api/file-cache.mdx
@@ -50,6 +50,7 @@ The cache classifies files by path. Set
`local-cache.whitelist` to select the fi
| GLOBAL_INDEX | global-index | BTree and full-text index files | Yes |
| BUCKET_INDEX | bucket-index | Hash, deletion vector index files | No |
| DATA | data | Data files (ORC, Parquet, etc.) | No |
+| STRUCTURED_DATA | structured-data | Parquet, ORC, Avro, ROW, Mosaic
(PyPaimon/Rust) | No |
| FILE_INDEX | file-index | Data-file level bloom filter, bitmap | No |
The default whitelist is `meta,global-index`, and `*` selects every file type.
Files rewritten in
@@ -128,12 +129,15 @@ The Java snippet is a method-body fragment; declare or
handle the catalog except
| `local-cache.dir` | String | (none) | Directory for storing cached blocks on
disk. If not configured, memory cache is used. |
| `local-cache.max-size` | MemorySize | See below | Maximum cached block bytes
per cache manager. Least recently used blocks are evicted when the limit is
exceeded. |
| `local-cache.block-size` | MemorySize | 1 mb | Block size for caching. Files
are logically divided into fixed-size blocks and cached independently. |
-| `local-cache.whitelist` | String | meta,global-index | Comma-separated list
of file types to cache. Supported values: `meta`, `global-index`,
`bucket-index`, `data`, `file-index`, or `*` for all of them. |
+| `local-cache.whitelist` | String | meta,global-index | File types to cache
(see [Cached File Types](#cached-file-types)); `*` selects all. |
When `local-cache.max-size` is omitted, Java has no configured size limit.
PyPaimon uses a 256 MiB
memory limit or a 10 GiB disk limit. Set the option explicitly when you want
the same limit across
clients. The limit accounts for cached block bytes, not all reader buffers or
process memory.
+Use `local-cache.whitelist=meta,global-index,structured-data` to cache
structured files without
+BLOB or video bodies. PyPaimon's Vortex and Lance readers bypass this cache.
+
## How It Works
1. The reader requests bytes from an eligible file. The cache maps the byte
range to fixed-size
diff --git a/paimon-python/pypaimon/common/options/core_options.py
b/paimon-python/pypaimon/common/options/core_options.py
index d462943b93..93f04872c0 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -1150,8 +1150,8 @@ class CoreOptions:
.default_value("meta,global-index")
.with_description(
"Comma-separated list of file types to cache. "
- "Supported values: meta, global-index, bucket-index, data,
file-index, "
- "or * for all of them."
+ "Supported values: meta, global-index, bucket-index, data, "
+ "structured-data, file-index, or * for all of them."
)
)
diff --git
a/paimon-python/pypaimon/execution/python/read/reader/format_avro_reader.py
b/paimon-python/pypaimon/execution/python/read/reader/format_avro_reader.py
index 57309fab1d..ec20f9d8e3 100644
--- a/paimon-python/pypaimon/execution/python/read/reader/format_avro_reader.py
+++ b/paimon-python/pypaimon/execution/python/read/reader/format_avro_reader.py
@@ -36,8 +36,12 @@ class FormatAvroReader(RecordBatchReader):
def __init__(self, file_io: FileIO, file_path: str, read_fields:
List[str], full_fields: List[DataField],
push_down_predicate: Any, batch_size: int = 1024,
nested_name_paths: Optional[List[List[str]]] = None):
- file_path_for_io = file_io.to_filesystem_path(file_path)
- self._file = file_io.filesystem.open_input_file(file_path_for_io)
+ from pypaimon.filesystem.caching_file_io import CachingFileIO
+ if isinstance(file_io, CachingFileIO) and
file_io._is_cacheable(file_path):
+ self._file = file_io.new_input_stream(file_path)
+ else:
+ file_path_for_io = file_io.to_filesystem_path(file_path)
+ self._file = file_io.filesystem.open_input_file(file_path_for_io)
self._avro_reader = fastavro.reader(self._file)
self._batch_size = batch_size
self._push_down_predicate = push_down_predicate
diff --git
a/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py
b/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py
index be7eee5368..f303b358f7 100644
---
a/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py
+++
b/paimon-python/pypaimon/execution/python/read/reader/format_pyarrow_reader.py
@@ -275,6 +275,18 @@ def _file_format_metadata_cache_max_size(file_io: FileIO)
-> int:
CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE).get_bytes()
+def _uses_block_cache(file_io: FileIO, file_path: str) -> bool:
+ from pypaimon.filesystem.caching_file_io import CachingFileIO
+ return isinstance(file_io, CachingFileIO) and
file_io._is_cacheable(file_path)
+
+
+def _open_pyarrow_input_file(file_io: FileIO, file_path: str, file_size=None):
+ if _uses_block_cache(file_io, file_path):
+ return pa.PythonFile(
+ file_io._new_input_stream_with_known_size(file_path, file_size),
mode='r')
+ return
file_io.filesystem.open_input_file(file_io.to_filesystem_path(file_path))
+
+
def _file_format_dataset(file_io: FileIO, file_format: str, file_path: str,
cache_max_size: int,
file_size: Optional[int] = None):
@@ -291,6 +303,11 @@ def _file_format_dataset(file_io: FileIO, file_format:
str, file_path: str,
if known_size is not None and register_file_size is not None:
register_file_size(file_path_for_pyarrow, known_size)
+ if _uses_block_cache(file_io, file_path):
+ from pypaimon.filesystem.caching_file_system import
CachedFileSystemHandler
+ import pyarrow.fs as pafs
+ filesystem = pafs.PyFileSystem(CachedFileSystemHandler(file_io,
file_path, known_size))
+
def load():
if file_format == 'parquet':
parquet_format = ds.ParquetFileFormat()
@@ -315,7 +332,8 @@ def _file_format_dataset(file_io: FileIO, file_format: str,
file_path: str,
file_path_for_pyarrow, format=file_format, filesystem=filesystem)
key = (
- _FilesystemIdentity(filesystem), file_format, file_path_for_pyarrow,
+ _FilesystemIdentity(file_io if _uses_block_cache(file_io, file_path)
else filesystem),
+ file_format, file_path_for_pyarrow,
known_size)
if cache_max_size <= 0:
_reset_file_format_dataset_cache()
@@ -335,8 +353,7 @@ def _orc_schema_with_field_metadata(file_io: FileIO,
file_path: str,
return fallback
import pyarrow.orc as orc
- source = file_io.filesystem.open_input_file(
- file_io.to_filesystem_path(file_path))
+ source = _open_pyarrow_input_file(file_io, file_path)
try:
metadata = orc.ORCFile(source).metadata
arrow_schema = metadata.get(b"ARROW:schema")
@@ -383,7 +400,7 @@ class FormatPyArrowReader(RecordBatchReader):
file_path_for_pyarrow = file_io.to_filesystem_path(file_path)
self._row_group_cache = row_group_cache
self._row_group_cache_filesystem = _FilesystemIdentity(
- file_io.filesystem)
+ file_io if _uses_block_cache(file_io, file_path) else
file_io.filesystem)
self._row_group_cache_path = file_path_for_pyarrow
cache_max_size = _file_format_metadata_cache_max_size(file_io)
self.dataset = _file_format_dataset(
@@ -561,8 +578,7 @@ class FormatPyArrowReader(RecordBatchReader):
and self._selected_shared_map_paths)):
import pyarrow.parquet as pq
# ParquetFile(filesystem=...) is unavailable in PyArrow 6.
- self._parquet_source = file_io.filesystem.open_input_file(
- file_path_for_pyarrow)
+ self._parquet_source = _open_pyarrow_input_file(file_io,
file_path, file_size)
try:
self._parquet_file = pq.ParquetFile(self._parquet_source)
if (self._selected_parquet_row_groups is not None
@@ -581,8 +597,7 @@ class FormatPyArrowReader(RecordBatchReader):
raise
if file_format == 'orc' and self._selected_shared_map_paths:
import pyarrow.orc as orc
- self._orc_source = file_io.filesystem.open_input_file(
- file_path_for_pyarrow)
+ self._orc_source = _open_pyarrow_input_file(file_io, file_path)
self._orc_file = orc.ORCFile(self._orc_source)
if self._exhausted:
self._raw_batches = iter(())
diff --git a/paimon-python/pypaimon/filesystem/caching_file_io.py
b/paimon-python/pypaimon/filesystem/caching_file_io.py
index 38a5b2f2b2..d60d8c2ecd 100644
--- a/paimon-python/pypaimon/filesystem/caching_file_io.py
+++ b/paimon-python/pypaimon/filesystem/caching_file_io.py
@@ -26,6 +26,7 @@ the delegate FileIO.
"""
import hashlib
+import io
import os
import threading
import uuid
@@ -262,7 +263,7 @@ class LocalDiskCacheManager:
self._file_size_cache[file_path] = size
-class CachingInputStream:
+class CachingInputStream(io.RawIOBase):
"""Wraps a remote stream with block-level caching."""
def __init__(self, file_io, file_path: str, cache, file_size=None):
@@ -275,6 +276,18 @@ class CachingInputStream:
self._io_lock = threading.Lock()
self._remote_supports_pread = None
+ def readable(self):
+ return True
+
+ def seekable(self):
+ return True
+
+ def readinto(self, buffer):
+ view = memoryview(buffer).cast('B')
+ data = self.read(len(view))
+ view[:len(data)] = data
+ return len(data)
+
def _get_file_size(self) -> int:
if self._file_size == -1:
cached = self._cache.get_file_size(self._file_path)
@@ -320,14 +333,14 @@ class CachingInputStream:
first_block = offset // block_size
last_block = (end - 1) // block_size
- result = bytearray()
+ parts = []
for bi, block_data in self._read_blocks(first_block, last_block):
block_start = bi * block_size
start_in_block = max(offset - block_start, 0)
end_in_block = min(end - block_start, len(block_data))
- result.extend(block_data[start_in_block:end_in_block])
+ parts.append(block_data[start_in_block:end_in_block])
- return bytes(result)
+ return b''.join(parts)
def _read_blocks(self, first: int, last: int):
block_size = self._cache.block_size
@@ -393,6 +406,7 @@ class CachingInputStream:
if self._stream is not None:
self._stream.close()
self._stream = None
+ super().close()
def __enter__(self):
return self
@@ -471,12 +485,17 @@ class CachingFileIO(FileIO):
def properties(self):
return self._delegate.properties
+ def _is_cacheable(self, path: str):
+ file_type = FileType.classify(path)
+ return (self._cache is not None
+ and file_type in self._whitelist
+ and not FileType.is_mutable(path))
+
def new_input_stream(self, path: str):
return self._new_input_stream_with_known_size(path, None)
def _new_input_stream_with_known_size(self, path: str, file_size):
- file_type = FileType.classify(path)
- if self._cache is None or file_type not in self._whitelist or
FileType.is_mutable(path):
+ if not self._is_cacheable(path):
return self._delegate.new_input_stream(path)
return CachingInputStream(self._delegate, path, self._cache, file_size)
diff --git a/paimon-python/pypaimon/filesystem/caching_file_system.py
b/paimon-python/pypaimon/filesystem/caching_file_system.py
new file mode 100644
index 0000000000..5d530a0bff
--- /dev/null
+++ b/paimon-python/pypaimon/filesystem/caching_file_system.py
@@ -0,0 +1,78 @@
+# 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 pyarrow as pa
+import pyarrow.fs as pafs
+
+
+class CachedFileSystemHandler(pafs.FileSystemHandler):
+ """Route one fragment's input through FileIO while preserving filesystem
operations."""
+
+ def __init__(self, file_io, path, file_size=None):
+ self._file_io = file_io
+ self._path = path
+ self._delegate = file_io.filesystem
+ self._filesystem_path =
self._delegate.normalize_path(file_io.to_filesystem_path(path))
+ self._file_size = file_size
+
+ def get_type_name(self):
+ return 'paimon-cached-file'
+
+ def open_input_file(self, path):
+ if self._delegate.normalize_path(path) != self._filesystem_path:
+ return self._delegate.open_input_file(path)
+ return pa.PythonFile(
+ self._file_io._new_input_stream_with_known_size(self._path,
self._file_size), mode='r')
+
+ def open_input_stream(self, path):
+ return self.open_input_file(path)
+
+ def normalize_path(self, path):
+ return self._delegate.normalize_path(path)
+
+ def get_file_info(self, paths):
+ return self._delegate.get_file_info(paths)
+
+ def get_file_info_selector(self, selector):
+ return self._delegate.get_file_info(selector)
+
+ def create_dir(self, path, recursive):
+ return self._delegate.create_dir(path, recursive=recursive)
+
+ def delete_dir(self, path):
+ return self._delegate.delete_dir(path)
+
+ def delete_dir_contents(self, path, missing_dir_ok):
+ return self._delegate.delete_dir_contents(path,
missing_dir_ok=missing_dir_ok)
+
+ def delete_root_dir_contents(self):
+ return self._delegate.delete_dir_contents("", accept_root_dir=True)
+
+ def delete_file(self, path):
+ return self._delegate.delete_file(path)
+
+ def move(self, src, dest):
+ return self._delegate.move(src, dest)
+
+ def copy_file(self, src, dest):
+ return self._delegate.copy_file(src, dest)
+
+ def open_output_stream(self, path, metadata):
+ return self._delegate.open_output_stream(path, metadata=metadata)
+
+ def open_append_stream(self, path, metadata):
+ return self._delegate.open_append_stream(path, metadata=metadata)
diff --git a/paimon-python/pypaimon/tests/caching_file_io_test.py
b/paimon-python/pypaimon/tests/caching_file_io_test.py
index 25078989c2..ae5ff07044 100644
--- a/paimon-python/pypaimon/tests/caching_file_io_test.py
+++ b/paimon-python/pypaimon/tests/caching_file_io_test.py
@@ -26,6 +26,7 @@ import pickle
import shutil
import tempfile
import threading
+import tracemalloc
import unittest
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import MagicMock, patch
@@ -242,6 +243,25 @@ class CachingInputStreamTest(unittest.TestCase):
result = stream.read(30) # 90..120, spans blocks 0 and 1
self.assertEqual(data[90:120], result)
+ def test_large_cached_read_uses_one_result_buffer(self):
+ block_size = 256 * 1024
+ data = b'x' * (32 * block_size)
+ cache = LocalMemoryCacheManager(len(data), block_size=block_size)
+ for i in range(32):
+ cache.put_block('test', i, data[i * block_size:(i + 1) *
block_size])
+ for positioned in (False, True):
+ with self.subTest(positioned=positioned):
+ stream = CachingInputStream(self._make_file_io(data), 'test',
cache)
+ tracemalloc.start()
+ try:
+ result = (stream.read_at(len(data), 0) if positioned else
stream.read())
+ _, peak = tracemalloc.get_traced_memory()
+ finally:
+ tracemalloc.stop()
+ self.assertEqual(data, result)
+ self.assertEqual(0 if positioned else len(data), stream.tell())
+ self.assertLess(peak, len(data) * 3 // 2)
+
def test_cache_hit_avoids_remote_read(self):
data = b"0123456789abcdef"
cache = LocalDiskCacheManager(self.cache_dir, 2 ** 63 - 1,
block_size=8)
@@ -472,6 +492,57 @@ class CachingFileIOTest(unittest.TestCase):
result = caching_io.new_input_stream("global-index-uuid.index")
self.assertNotIsInstance(result, CachingInputStream)
+ def test_structured_data_cache_preserves_existing_selectors(self):
+ from pypaimon.filesystem.caching_file_io import LocalMemoryCacheManager
+ from pypaimon.utils.file_type import FileType
+
+ for disk in (False, True):
+ for whitelist in ('meta,global-index', 'data', 'structured-data',
+ 'meta,global-index,structured-data', '*'):
+ for name in ('data.parquet', 'data.blob', 'data.orc',
'data.avro',
+ 'data.row', 'data.mosaic', 'data.vortex',
'data.lance',
+ 'data.video', 'data.unknown',
'data.parquet.index',
+ 'snapshot-1', 'snapshot-1.blob'):
+ with self.subTest(disk=disk, whitelist=whitelist,
name=name):
+ with tempfile.TemporaryDirectory() as directory:
+ cache = (LocalDiskCacheManager(directory, 1024,
block_size=4) if disk
+ else LocalMemoryCacheManager(1024,
block_size=4))
+ delegate = self._make_delegate({name: b'abcdefgh'})
+ caching_io = CachingFileIO(
+ delegate, cache,
FileType.parse_whitelist(whitelist))
+ if whitelist == '*':
+ cached = True
+ elif whitelist == 'data':
+ cached = name.startswith('data.') and not
name.endswith('.index')
+ elif whitelist == 'structured-data':
+ cached = name in (
+ 'data.parquet', 'data.orc', 'data.avro',
+ 'data.row', 'data.mosaic')
+ elif whitelist ==
'meta,global-index,structured-data':
+ cached = name in (
+ 'data.parquet', 'data.orc', 'data.avro',
'data.row',
+ 'data.mosaic', 'snapshot-1',
'snapshot-1.blob')
+ else:
+ cached = name.startswith('snapshot-')
+ for _ in range(2):
+ with caching_io.new_input_stream(name) as
stream:
+ stream.seek(1)
+ self.assertEqual(b'bcdefg', stream.read(6))
+ self.assertEqual(1 if cached else 2,
+
delegate.new_input_stream.call_count)
+ if not cached:
+ delegate.get_file_size.assert_not_called()
+
+ def test_data_whitelist_still_caches_blob(self):
+ from pypaimon.utils.file_type import FileType
+ delegate = self._make_delegate({"data.blob": b"abcdefgh"})
+ cache = LocalDiskCacheManager(self.cache_dir, 1024, block_size=4)
+ caching_io = CachingFileIO(delegate, cache, {FileType.DATA})
+ for _ in range(2):
+ with caching_io.new_input_stream("data.blob") as stream:
+ self.assertEqual(b"abcdefgh", stream.read())
+ delegate.new_input_stream.assert_called_once()
+
def test_data_file_not_cached(self):
data = b"data content"
delegate = self._make_delegate({"data-abc.orc": data})
diff --git a/paimon-python/pypaimon/tests/file_type_test.py
b/paimon-python/pypaimon/tests/file_type_test.py
index 6b00d4a069..31f9bdb56a 100644
--- a/paimon-python/pypaimon/tests/file_type_test.py
+++ b/paimon-python/pypaimon/tests/file_type_test.py
@@ -83,10 +83,19 @@ class FileTypeClassifyTest(unittest.TestCase):
self.assertEqual(FileType.BUCKET_INDEX,
FileType.classify("index-uuid-1"))
def test_data_files(self):
- self.assertEqual(FileType.DATA, FileType.classify("data-abc.orc"))
- self.assertEqual(FileType.DATA, FileType.classify("data-abc.parquet"))
+ for name in ("data-abc.parquet", "data-abc.ORC", "data-abc.avro",
+ "data-abc.row", "data-abc.mosaic"):
+ self.assertEqual(FileType.STRUCTURED_DATA, FileType.classify(name))
+ for name in ("data-abc.blob", "data-abc.video", "data-abc.vortex",
+ "data-abc.lance"):
+ self.assertEqual(FileType.DATA, FileType.classify(name))
self.assertEqual(FileType.DATA, FileType.classify("unknown-file"))
+ def test_structured_data_keeps_metadata_and_indexes_separate(self):
+ self.assertEqual(FileType.META, FileType.classify("manifest-1.avro"))
+ self.assertEqual(FileType.META,
FileType.classify("snapshot-1.parquet"))
+ self.assertEqual(FileType.FILE_INDEX,
FileType.classify("data.parquet.index"))
+
def test_temp_file_unwrap(self):
# .{originalName}.{UUID}.tmp -> originalName
uuid = "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
@@ -94,7 +103,7 @@ class FileTypeClassifyTest(unittest.TestCase):
self.assertEqual(FileType.META, FileType.classify(temp_name))
temp_data = f".data-abc.orc.{uuid}.tmp"
- self.assertEqual(FileType.DATA, FileType.classify(temp_data))
+ self.assertEqual(FileType.STRUCTURED_DATA,
FileType.classify(temp_data))
def test_temp_file_not_matching_pattern(self):
self.assertEqual(FileType.DATA, FileType.classify(".short.tmp"))
@@ -106,6 +115,16 @@ class FileTypeClassifyTest(unittest.TestCase):
self.assertTrue(FileType.FILE_INDEX.is_index())
self.assertFalse(FileType.META.is_index())
self.assertFalse(FileType.DATA.is_index())
+ self.assertFalse(FileType.STRUCTURED_DATA.is_index())
+
+ def test_parse_whitelist(self):
+ self.assertEqual(
+ {FileType.STRUCTURED_DATA},
+ FileType.parse_whitelist("structured-data"))
+ self.assertEqual(
+ {FileType.DATA, FileType.STRUCTURED_DATA},
+ FileType.parse_whitelist("data"))
+ self.assertEqual(set(FileType), FileType.parse_whitelist("*"))
class FileTypeMutableTest(unittest.TestCase):
@@ -162,13 +181,14 @@ class FileTypeMutableTest(unittest.TestCase):
class FileTypeParseWhitelistTest(unittest.TestCase):
- ALL = {FileType.META, FileType.DATA, FileType.BUCKET_INDEX,
- FileType.GLOBAL_INDEX, FileType.FILE_INDEX}
+ ALL = {FileType.META, FileType.DATA, FileType.STRUCTURED_DATA,
+ FileType.BUCKET_INDEX, FileType.GLOBAL_INDEX, FileType.FILE_INDEX}
def test_named_values(self):
self.assertEqual({FileType.META, FileType.GLOBAL_INDEX},
FileType.parse_whitelist("meta,global-index"))
- self.assertEqual({FileType.BUCKET_INDEX, FileType.DATA,
FileType.FILE_INDEX},
+ self.assertEqual({FileType.BUCKET_INDEX, FileType.DATA,
+ FileType.STRUCTURED_DATA, FileType.FILE_INDEX},
FileType.parse_whitelist(" bucket-index ,
data,file-index "))
def test_wildcard(self):
diff --git a/paimon-python/pypaimon/tests/parquet_local_cache_test.py
b/paimon-python/pypaimon/tests/parquet_local_cache_test.py
new file mode 100644
index 0000000000..5779941008
--- /dev/null
+++ b/paimon-python/pypaimon/tests/parquet_local_cache_test.py
@@ -0,0 +1,217 @@
+# 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.
+
+
+from unittest.mock import patch
+
+import fastavro
+import pyarrow as pa
+import pyarrow.orc as orc
+import pyarrow.parquet as pq
+import pytest
+
+from pypaimon.common.options import Options
+from pypaimon.filesystem.caching_file_io import CachingFileIO
+from pypaimon.filesystem.local_file_io import LocalFileIO
+from pypaimon.execution.python.read.reader.format_avro_reader import
FormatAvroReader
+from pypaimon.execution.python.read.reader.format_mosaic_reader import
FormatMosaicReader
+from pypaimon.execution.python.read.reader.format_pyarrow_reader import (
+ FormatPyArrowReader, _reset_file_format_dataset_cache,
+)
+from pypaimon.execution.python.read.reader.format_row_reader import
FormatRowReader
+from pypaimon.execution.python.write.format_file_writer import write_row_file
+from pypaimon.schema.data_types import AtomicType, DataField
+
+
+def _read(file_io, path, file_format, rows, file_size=None):
+ reader = FormatPyArrowReader(
+ file_io, file_format, path, [DataField(0, 'v', AtomicType('INT'))],
+ None, row_indices=rows, batch_size=2, file_size=file_size)
+ values = []
+ try:
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ values.extend(batch.column(0).to_pylist())
+ finally:
+ reader.close()
+ return values
+
+
+def _read_custom(reader):
+ 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()
+
+
[email protected]('disk', [False, True])
[email protected]('metadata_cache', ['0 b', '1 mb'])
[email protected]('file_format,rows', [('parquet', None), ('parquet',
[1]), ('orc', None)])
+def test_format_reader_uses_block_cache(tmp_path, disk, metadata_cache,
file_format, rows):
+ path = tmp_path / ('data.' + file_format)
+ table = pa.table({'v': pa.array([1, 2, 3], type=pa.int32())})
+ if file_format == 'parquet':
+ pq.write_table(table, path, row_group_size=1)
+ else:
+ orc.write_table(table, path)
+ opts = Options({'local-cache.enabled': 'true', 'local-cache.block-size':
'64 b',
+ 'local-cache.whitelist': 'structured-data',
+ 'file-format.metadata-cache.max-size': metadata_cache})
+ if disk:
+ opts.data['local-cache.dir'] = str(tmp_path / 'cache')
+ delegate = LocalFileIO(catalog_options=opts)
+ cache = CachingFileIO.create_cache_manager(opts)
+ file_io = CachingFileIO.wrap_with_caching_if_needed(delegate, opts, cache)
+ expected = [2] if rows is not None else [1, 2, 3]
+ _reset_file_format_dataset_cache()
+ try:
+ opened = []
+ original_open = delegate.new_input_stream
+
+ def open_source(location):
+ stream = original_open(location)
+ opened.append(stream)
+ return stream
+
+ with patch.object(delegate, 'new_input_stream',
side_effect=open_source) as opens:
+ assert _read(file_io, path.as_uri(), file_format, rows) == expected
+ assert opens.call_count > 0
+ assert all(stream.closed for stream in opened)
+ if file_format == 'parquet':
+ path.unlink()
+ # Drop the parsed-footer cache, proving that raw cached blocks are
sufficient.
+ _reset_file_format_dataset_cache()
+ with patch.object(delegate, 'new_input_stream',
side_effect=AssertionError('source read')):
+ with patch.object(delegate, 'get_file_size',
side_effect=AssertionError('source stat')):
+ assert _read(file_io, path.as_uri(), file_format, rows) ==
expected
+ finally:
+ _reset_file_format_dataset_cache()
+
+
+def test_known_parquet_size_avoids_source_stat(tmp_path):
+ path = tmp_path / 'data.parquet'
+ pq.write_table(pa.table({'v': pa.array([1, 2, 3], type=pa.int32())}), path)
+ opts = Options({'local-cache.enabled': 'true',
+ 'local-cache.whitelist': 'structured-data',
+ 'file-format.metadata-cache.max-size': '0 b'})
+ delegate = LocalFileIO(catalog_options=opts)
+ file_io = CachingFileIO.wrap_with_caching_if_needed(
+ delegate, opts, CachingFileIO.create_cache_manager(opts))
+ with patch.object(delegate, 'get_file_size',
side_effect=AssertionError('source stat')):
+ assert _read(file_io, str(path), 'parquet', None, path.stat().st_size)
== [1, 2, 3]
+
+
[email protected]('disk', [False, True])
+def test_known_parquet_size_selected_rows_use_warm_blocks(tmp_path, disk):
+ path = tmp_path / 'data.parquet'
+ pq.write_table(pa.table({'v': pa.array([1, 2, 3], type=pa.int32())}),
+ path, row_group_size=1)
+ file_size = path.stat().st_size
+ opts = Options({'local-cache.enabled': 'true', 'local-cache.block-size':
'64 b',
+ 'local-cache.whitelist': 'structured-data'})
+ if disk:
+ opts.data['local-cache.dir'] = str(tmp_path / 'cache')
+ delegate = LocalFileIO(catalog_options=opts)
+ file_io = CachingFileIO.wrap_with_caching_if_needed(
+ delegate, opts, CachingFileIO.create_cache_manager(opts))
+ _reset_file_format_dataset_cache()
+ try:
+ assert _read(file_io, path.as_uri(), 'parquet', None, file_size) ==
[1, 2, 3]
+ path.unlink()
+ _reset_file_format_dataset_cache()
+ with patch.object(delegate, 'new_input_stream',
+ side_effect=AssertionError('source read')):
+ with patch.object(delegate, 'get_file_size',
+ side_effect=AssertionError('source stat')):
+ assert _read(file_io, path.as_uri(), 'parquet', [1],
file_size) == [2]
+ finally:
+ _reset_file_format_dataset_cache()
+
+
[email protected]('disk', [False, True])
+def test_avro_reader_uses_block_cache(tmp_path, disk):
+ path = tmp_path / 'data.avro'
+ schema = {'type': 'record', 'name': 'row',
+ 'fields': [{'name': 'v', 'type': 'int'}]}
+ with path.open('wb') as output:
+ fastavro.writer(output, schema, [{'v': 1}, {'v': 2}, {'v': 3}])
+
+ opts = Options({'local-cache.enabled': 'true', 'local-cache.block-size':
'64 b',
+ 'local-cache.whitelist': 'structured-data'})
+ if disk:
+ opts.data['local-cache.dir'] = str(tmp_path / 'cache')
+ delegate = LocalFileIO(catalog_options=opts)
+ file_io = CachingFileIO.wrap_with_caching_if_needed(
+ delegate, opts, CachingFileIO.create_cache_manager(opts))
+
+ def read():
+ reader = FormatAvroReader(file_io, path.as_uri(), ['v'],
+ [DataField(0, 'v', AtomicType('INT'))],
None, batch_size=2)
+ 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()
+
+ assert read() == [1, 2, 3]
+ path.unlink()
+ with patch.object(delegate, 'new_input_stream',
side_effect=AssertionError('source read')):
+ with patch.object(delegate, 'get_file_size',
side_effect=AssertionError('source stat')):
+ assert read() == [1, 2, 3]
+
+
[email protected]('file_format', ['row', 'mosaic'])
+def test_other_structured_readers_use_block_cache(tmp_path, file_format):
+ if file_format == 'mosaic':
+ pytest.importorskip('mosaic')
+
+ path = tmp_path / ('data.' + file_format)
+ fields = [DataField(0, 'v', AtomicType('INT'))]
+ table = pa.table({'v': pa.array([1, 2, 3], type=pa.int32())})
+ opts = Options({'local-cache.enabled': 'true', 'local-cache.block-size':
'64 b',
+ 'local-cache.whitelist': 'structured-data'})
+ delegate = LocalFileIO(catalog_options=opts)
+ if file_format == 'row':
+ write_row_file(delegate, str(path), table, fields=fields)
+ else:
+ delegate.write_mosaic(str(path), table)
+ file_io = CachingFileIO.wrap_with_caching_if_needed(
+ delegate, opts, CachingFileIO.create_cache_manager(opts))
+
+ def read():
+ if file_format == 'row':
+ reader = FormatRowReader(file_io, str(path), ['v'], fields, None)
+ else:
+ reader = FormatMosaicReader(file_io, str(path), fields, None)
+ return _read_custom(reader)
+
+ with patch.object(delegate, 'new_input_stream',
wraps=delegate.new_input_stream) as opens:
+ assert read() == [1, 2, 3]
+ cold_opens = opens.call_count
+ assert cold_opens > 0
+ assert read() == [1, 2, 3]
+ assert opens.call_count == cold_opens
diff --git a/paimon-python/pypaimon/utils/file_type.py
b/paimon-python/pypaimon/utils/file_type.py
index 7039bf6204..2d3514c04b 100644
--- a/paimon-python/pypaimon/utils/file_type.py
+++ b/paimon-python/pypaimon/utils/file_type.py
@@ -31,12 +31,14 @@ class FileType(Enum):
- META: snapshot, schema, manifest, manifest sidecar, statistics, tag,
changelog metadata,
hint files, _SUCCESS, consumer, service files
- DATA: data files and any unrecognized files (default)
+ - STRUCTURED_DATA: block-cache-readable Parquet, ORC, Avro, ROW and Mosaic
files
- BUCKET_INDEX: bucket level index files (Hash, DV)
- GLOBAL_INDEX: table level global index files (btree, lumina, full-text)
- FILE_INDEX: data-file index files (bloom filter, bitmap, etc.)
"""
META = "META"
DATA = "DATA"
+ STRUCTURED_DATA = "STRUCTURED_DATA"
BUCKET_INDEX = "BUCKET_INDEX"
GLOBAL_INDEX = "GLOBAL_INDEX"
FILE_INDEX = "FILE_INDEX"
@@ -92,16 +94,21 @@ class FileType(Enum):
if parent == "changelog":
return FileType.META
+ stem, _, extension = name.rpartition('.')
+ if stem and extension.lower() in ('parquet', 'orc', 'avro', 'row',
'mosaic'):
+ return FileType.STRUCTURED_DATA
+
return FileType.DATA
@staticmethod
def parse_whitelist(whitelist_str: str) -> set:
mapping = {
- "meta": FileType.META,
- "global-index": FileType.GLOBAL_INDEX,
- "bucket-index": FileType.BUCKET_INDEX,
- "data": FileType.DATA,
- "file-index": FileType.FILE_INDEX,
+ "meta": (FileType.META,),
+ "global-index": (FileType.GLOBAL_INDEX,),
+ "bucket-index": (FileType.BUCKET_INDEX,),
+ "data": (FileType.DATA, FileType.STRUCTURED_DATA),
+ "structured-data": (FileType.STRUCTURED_DATA,),
+ "file-index": (FileType.FILE_INDEX,),
}
result = set()
for name in whitelist_str.split(","):
@@ -109,11 +116,12 @@ class FileType(Enum):
if name == "*":
result.update(FileType)
elif name in mapping:
- result.add(mapping[name])
+ result.update(mapping[name])
elif name:
logger.warning(
"Unknown local-cache.whitelist value '%s'. "
- "Supported values: meta, global-index, bucket-index, data,
file-index, or * for all of them.",
+ "Supported values: meta, global-index, bucket-index, data,
"
+ "structured-data, file-index, or * for all of them.",
name,
)
return result