From ffdbe6c11597c99fe5130d0cc48d1a819fbdd78a Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Tue, 15 Sep 2026 08:46:40 -0700 Subject: [PATCH] Increase the default PyArrow IO buffer to 8 MiB Use one buffer default across FileIO and Parquet readers while preserving explicit overrides. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyiceberg/io/pyarrow.py | 15 +++++++++------ tests/io/test_pyarrow.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..ac0311bd3d 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -195,8 +195,8 @@ logger = logging.getLogger(__name__) -ONE_MEGABYTE = 1024 * 1024 BUFFER_SIZE = "buffer-size" +_DEFAULT_BUFFER_SIZE = 8 * 1024 * 1024 # 8 MiB ICEBERG_SCHEMA = b"iceberg.schema" # The PARQUET: in front means that it is Parquet specific, in this case the field_id PYARROW_PARQUET_FIELD_ID_KEY = b"PARQUET:field_id" @@ -260,6 +260,9 @@ class PyArrowFile(InputFile, OutputFile): Args: location (str): A URI or a path to a local file. + path (str): The filesystem path. + fs (FileSystem): The filesystem used to access the file. + buffer_size (int): Buffer size in bytes. Defaults to 8 MiB. Attributes: location(str): The URI or path to a local file for a PyArrowFile instance. @@ -281,7 +284,7 @@ class PyArrowFile(InputFile, OutputFile): _path: str _buffer_size: int - def __init__(self, location: str, path: str, fs: FileSystem, buffer_size: int = ONE_MEGABYTE): + def __init__(self, location: str, path: str, fs: FileSystem, buffer_size: int = _DEFAULT_BUFFER_SIZE): self._filesystem = fs self._path = path self._buffer_size = buffer_size @@ -642,7 +645,7 @@ def new_input(self, location: str) -> PyArrowFile: fs=self.fs_by_scheme(scheme, netloc), location=location, path=path, - buffer_size=int(self.properties.get(BUFFER_SIZE, ONE_MEGABYTE)), + buffer_size=int(self.properties.get(BUFFER_SIZE, _DEFAULT_BUFFER_SIZE)), ) @override @@ -660,7 +663,7 @@ def new_output(self, location: str) -> PyArrowFile: fs=self.fs_by_scheme(scheme, netloc), location=location, path=path, - buffer_size=int(self.properties.get(BUFFER_SIZE, ONE_MEGABYTE)), + buffer_size=int(self.properties.get(BUFFER_SIZE, _DEFAULT_BUFFER_SIZE)), ) @override @@ -1131,7 +1134,7 @@ def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray] if data_file.file_format == FileFormat.PARQUET: with io.new_input(data_file.file_path).open() as fi: delete_fragment = _get_file_format( - data_file.file_format, dictionary_columns=("file_path",), pre_buffer=True, buffer_size=ONE_MEGABYTE + data_file.file_format, dictionary_columns=("file_path",), pre_buffer=True, buffer_size=_DEFAULT_BUFFER_SIZE ).make_fragment(fi) table = ds.Scanner.from_fragment(fragment=delete_fragment).to_table() table = table.unify_dictionaries() @@ -1641,7 +1644,7 @@ def _task_to_record_batches( downcast_ns_timestamp_to_us: bool | None = None, dictionary_columns: tuple[str, ...] = (), ) -> Iterator[pa.RecordBatch]: - format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8} + format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": _DEFAULT_BUFFER_SIZE} if dictionary_columns and task.file.file_format == FileFormat.PARQUET: format_kwargs["dictionary_columns"] = dictionary_columns arrow_format = _get_file_format(task.file.file_format, **format_kwargs) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..eb5c7bccc1 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -222,6 +222,21 @@ def test_pyarrow_output_file() -> None: assert len(output_file) == 3 +def test_pyarrow_file_buffer_sizes() -> None: + # Constructor default + assert PyArrowFile("buffered.bin", path="buffered.bin", fs=LocalFileSystem())._buffer_size == 8 * 1024 * 1024 + + # No override + io = PyArrowFileIO() + assert io.new_input("buffered.bin")._buffer_size == 8 * 1024 * 1024 + assert io.new_output("buffered.bin")._buffer_size == 8 * 1024 * 1024 + + # Custom override + io = PyArrowFileIO({"buffer-size": "4096"}) + assert io.new_input("buffered.bin")._buffer_size == 4096 + assert io.new_output("buffered.bin")._buffer_size == 4096 + + def test_pyarrow_invalid_scheme() -> None: """Test that a ValueError is raised if a location is provided with an invalid scheme"""