From e25850c869662e5e938b486cbfdf923d3784894f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 23:49:00 +0900 Subject: [PATCH 01/11] Compress pipe_file() values up front for the compression argument pipe_file() passed fsspec's compression argument to PutObject on the single-request path, which botocore rejects, and to open() on the buffered path, which returns a compression wrapper. A failed write on the buffered path then raised AttributeError from _close_without_commit() on the wrapper, hiding the original error, and still uploaded an empty compressed object. Compress the value in memory with the codec that open() uses, including "infer" from the path, and upload the compressed bytes through the usual paths. The same applies to the transaction path of AioS3FileSystem. Closes #1037 Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 48 +++++++++++++-- pyathena/filesystem/s3_async.py | 14 +++-- tests/pyathena/filesystem/test_s3.py | 68 ++++++++++++++++++++++ tests/pyathena/filesystem/test_s3_async.py | 15 +++++ 4 files changed, 137 insertions(+), 8 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 5434a297..3b84e08d 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -25,6 +25,8 @@ from botocore.client import BaseClient, Config from fsspec import AbstractFileSystem from fsspec.callbacks import _DEFAULT_CALLBACK, Callback +from fsspec.compression import compr +from fsspec.core import get_compression from fsspec.implementations.local import trailing_sep from fsspec.spec import AbstractBufferedFile from fsspec.utils import isfilelike, other_paths, tokenize @@ -49,6 +51,36 @@ _logger = logging.getLogger(__name__) +def _compress( + path: str, value: bytes | bytearray | memoryview, compression: str +) -> bytes | bytearray | memoryview: + """Compress a value with the codec that ``open()`` uses for a compression. + + Args: + path: Path of the file, from which ``"infer"`` takes the codec. + value: The bytes to compress. + compression: Name of a codec in ``fsspec.compression.compr``, or + ``"infer"`` to take it from the extension of the path. + + Returns: + The compressed bytes, or the value itself when ``"infer"`` finds no + codec for the path. + + Raises: + ValueError: If the codec is not supported. + """ + compression = get_compression(path, compression) + if compression is None: + return value + if isinstance(value, memoryview) and not value.c_contiguous: + # Codecs cannot compress a non-contiguous memoryview. + value = value.tobytes() + buffer = BytesIO() + with compr[compression](buffer, mode="w") as f: + f.write(value) + return buffer.getvalue() + + class S3FileSystem(AbstractFileSystem): """A filesystem interface for Amazon S3 that implements the fsspec protocol. @@ -1760,16 +1792,24 @@ def pipe_file( (e.g., ContentType, StorageClass) on the single-request path. The ``block_size``, ``max_workers``, and ``s3_additional_kwargs`` parameters of the ``open()`` path - are also accepted. + are also accepted, and so is ``compression``: the codec of + ``open()`` to compress the value with before it is + uploaded, or ``"infer"`` to take it from the extension of + the path. Raises: FileExistsError: If the mode is "create" and the path already exists, or an object is created at it before the write is committed. ValueError: If the path does not contain a key or specifies a - version, or if the data takes more than - ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. - """ + version, if the compression is not supported, or if the data + takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + """ + compression = kwargs.pop("compression", None) + if compression is not None: + # Compressed up front, so that every path uploads the compressed + # bytes, and open() returns the file instead of a wrapper. + value = _compress(path, value, compression) block_size = kwargs.get("block_size") or self.default_block_size # The size in bytes; the length of a memoryview counts its items. self._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index a7d933bc..fe941ae2 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -19,7 +19,7 @@ from fsspec.asyn import AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK -from pyathena.filesystem.s3 import S3File, S3FileSystem +from pyathena.filesystem.s3 import S3File, S3FileSystem, _compress from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import ( S3Metadata, @@ -187,14 +187,20 @@ def _pipe_file_in_transaction( opened in ``xb`` mode: raise FileExistsError when the object already exists, including one created before the transaction is committed, which is not replaced. - **kwargs: Additional parameters passed to ``open()``. + **kwargs: Additional parameters passed to ``open()``, except + ``compression``, with which the value is compressed before + it is written, as in :meth:`S3FileSystem.pipe_file`. Raises: FileExistsError: If the mode is "create" and the path already exists. - ValueError: If the data takes more than - ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. + ValueError: If the compression is not supported, or if the data + takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ + compression = kwargs.pop("compression", None) + if compression is not None: + # See S3FileSystem.pipe_file. + value = _compress(path, value, compression) block_size = kwargs.get("block_size") or self._sync_fs.default_block_size # The size in bytes; the length of a memoryview counts its items. self._sync_fs._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 4eed9dd0..c5ffb61e 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -2,6 +2,7 @@ import contextlib import functools import gc +import gzip import io import os import re @@ -1536,6 +1537,73 @@ def test_pipe_file_failed_write(self, intrans): fs._put_object.assert_not_called() fs._call.assert_not_called() + @pytest.mark.parametrize( + ("path", "compression"), [("s3://bucket/key", "gzip"), ("s3://bucket/key.gz", "infer")] + ) + @pytest.mark.parametrize("intrans", [False, True]) + @pytest.mark.parametrize("size", [1, S3FileSystem.DEFAULT_BLOCK_SIZE + 1]) + def test_pipe_file_compression(self, path, compression, intrans, size): + # GH-1037: the value is compressed before it is uploaded, on every + # path. The single-request path used to send compression to + # PutObject, which botocore rejects. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs._transaction = None + fs._put_object = mock.MagicMock() + value = b"a" * size + + with fs.transaction if intrans else contextlib.nullcontext(): + fs.pipe_file(path, value, compression=compression) + + # The compressed value fits in one block. + ((_, kwargs),) = fs._put_object.call_args_list + assert "compression" not in kwargs + assert gzip.decompress(kwargs["body"]) == value + + def test_pipe_file_compression_inferred_none(self): + # "infer" uploads the value as it is for a path without the + # extension of a codec, as open() does. + fs = self._make_fs() + fs._put_object = mock.MagicMock() + + fs.pipe_file("s3://bucket/key.txt", b"a", compression="infer") + + ((_, kwargs),) = fs._put_object.call_args_list + assert "compression" not in kwargs + assert kwargs["body"] == b"a" + + def test_pipe_file_unsupported_compression(self): + fs = self._make_fs() + + with pytest.raises(ValueError, match="not supported"): + fs.pipe_file("s3://bucket/key", b"a", compression="unknown") + fs._call.assert_not_called() + + @pytest.mark.parametrize("intrans", [False, True]) + def test_pipe_file_compression_failed_write(self, intrans): + # GH-1037: a failed write of a compressed value leaves the existing + # object unchanged. open() used to return a compression wrapper, + # without _close_without_commit(), and the object was replaced with + # an empty compressed one. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs._transaction = None + fs._put_object = mock.MagicMock() + # Random bytes stay larger than the block size when compressed, so + # that the buffered path also writes them outside a transaction. + value = b"a" if intrans else os.urandom(S3FileSystem.DEFAULT_BLOCK_SIZE + 1) + + with ( + mock.patch.object(S3File, "write", side_effect=RuntimeError("write failed")), + fs.transaction if intrans else contextlib.nullcontext(), + pytest.raises(RuntimeError, match="write failed"), + ): + fs.pipe_file("s3://bucket/key", value, compression="gzip") + gc.collect() + + fs._put_object.assert_not_called() + fs._call.assert_not_called() + def test_pipe_file_failed_write_aborts_multipart_upload(self): # GH-997: a write that fails after its multipart upload has started # aborts the upload instead of completing it. diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f65e72de..a0ac8fad 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import gzip import os import tempfile import threading @@ -269,6 +270,20 @@ def test_put_file_mode(self, tmp_path, mode): assert "mode" not in call.kwargs assert call.kwargs.get("IfNoneMatch") == ("*" if mode == "create" else None) + @pytest.mark.parametrize("intrans", [False, True]) + def test_pipe_file_compression(self, intrans): + # GH-1037: the value is compressed before it is uploaded, also in a + # transaction of this filesystem. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + put_object = fs._sync_fs._put_object = mock.MagicMock() + + with fs.transaction if intrans else contextlib.nullcontext(): + fs.pipe_file("s3://bucket/key", b"data", compression="gzip") + + ((_, kwargs),) = put_object.call_args_list + assert "compression" not in kwargs + assert gzip.decompress(kwargs["body"]) == b"data" + def test_transaction_pipe_file_write(self): # GH-997: in a transaction, a non-contiguous memoryview is written, # and a failed write does not replace the object with an empty one From ebb855d7cc6759eec2981f00f5d2601eb4bd27c7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 23:52:16 +0900 Subject: [PATCH 02/11] Keep the compressed data when a codec closes its buffer The zstandard stream writer, fsspec's zstd codec without compression.zstd or backports.zstd, closes the file that it writes to, so reading the compressed bytes from the BytesIO afterwards raised ValueError. Write to a BytesIO whose close() keeps the data. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 16 ++++++++++++++-- tests/pyathena/filesystem/test_s3.py | 16 +++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 3b84e08d..7d25f448 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -46,11 +46,23 @@ S3PutObject, S3StorageClass, ) -from pyathena.util import RetryConfig, retry_api_call +from pyathena.util import RetryConfig, override, retry_api_call _logger = logging.getLogger(__name__) +class _CompressedBuffer(BytesIO): + """A buffer that keeps its data when a codec closes it. + + Some codecs, such as the ``zstandard`` stream writer, close the file + that they write to when they are closed. + """ + + @override + def close(self) -> None: + pass + + def _compress( path: str, value: bytes | bytearray | memoryview, compression: str ) -> bytes | bytearray | memoryview: @@ -75,7 +87,7 @@ def _compress( if isinstance(value, memoryview) and not value.c_contiguous: # Codecs cannot compress a non-contiguous memoryview. value = value.tobytes() - buffer = BytesIO() + buffer = _CompressedBuffer() with compr[compression](buffer, mode="w") as f: f.write(value) return buffer.getvalue() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index c5ffb61e..2acc72ac 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -25,12 +25,13 @@ import pytest from botocore.stub import Stubber from fsspec import Callback +from fsspec.compression import compr from fsspec.dircache import DirCache from fsspec.implementations.dirfs import DirFileSystem import pyathena from pyathena.filesystem import register_s3_filesystem -from pyathena.filesystem.s3 import S3File, S3FileSystem +from pyathena.filesystem.s3 import S3File, S3FileSystem, _compress from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass @@ -4609,3 +4610,16 @@ def test_upload_chunk_multipart(self, autocommit): file.commit() file.fs._finish_multipart_upload.assert_called_once() file.fs._put_object.assert_not_called() + + +def test_compress_codec_closing_its_file(): + # GH-1037: some codecs, such as the zstandard stream writer, close the + # file that they write to when they are closed. + def closing_gzip(f, mode): + g = gzip.GzipFile(fileobj=f, mode=mode) + close = g.close + g.close = lambda: (close(), f.close()) + return g + + with mock.patch.dict(compr, {"closing": closing_gzip}): + assert gzip.decompress(_compress("s3://bucket/key", b"a", "closing")) == b"a" From 846227079ea492e3edddb55ea7fd4e0911d1bee8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 23:55:45 +0900 Subject: [PATCH 03/11] Document the compression argument of pipe_file() Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/filesystem.md b/docs/filesystem.md index d69a230a..2b29b359 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -83,6 +83,9 @@ block size (5 MiB by default); larger data is uploaded as a parallel multipart u through the buffered file path. Inside an [fsspec transaction](https://filesystem-spec.readthedocs.io/en/latest/features.html#transactions), writes are deferred until the transaction commits and are discarded on rollback. +With the `compression` argument (a codec of fsspec, or `"infer"` from the extension of +the path), `pipe`/`pipe_file` compress the data before uploading it, and the sizes in +this section apply to the compressed data. The block size for writing, given by the `block_size` argument of `open` or by the filesystem's `default_block_size`, must be between 5 MiB and 5 GiB, inclusive, the part From 135e64872fad15cc949a5f659f4a71f6fec49cc4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 00:03:44 +0900 Subject: [PATCH 04/11] Infer the compression from the path that open() infers it from fsspec's open() infers the codec from the path without the protocol, which also drops a trailing slash. pipe_file() inferred it from the path as given, so "key.gz/" was uploaded uncompressed. Strip the protocol first, also in the transaction path of AioS3FileSystem. Cover compressed multipart uploads and non-contiguous memoryviews too. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 5 ++-- pyathena/filesystem/s3_async.py | 2 +- tests/pyathena/filesystem/test_s3.py | 41 +++++++++++++++++++++++++++- 3 files changed, 44 insertions(+), 4 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 7d25f448..17ead468 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -69,7 +69,8 @@ def _compress( """Compress a value with the codec that ``open()`` uses for a compression. Args: - path: Path of the file, from which ``"infer"`` takes the codec. + path: Path of the file without the protocol, as ``open()`` strips + it, from which ``"infer"`` takes the codec. value: The bytes to compress. compression: Name of a codec in ``fsspec.compression.compr``, or ``"infer"`` to take it from the extension of the path. @@ -1821,7 +1822,7 @@ def pipe_file( if compression is not None: # Compressed up front, so that every path uploads the compressed # bytes, and open() returns the file instead of a wrapper. - value = _compress(path, value, compression) + value = _compress(self._strip_protocol(path), value, compression) block_size = kwargs.get("block_size") or self.default_block_size # The size in bytes; the length of a memoryview counts its items. self._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index fe941ae2..d8fea2c8 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -200,7 +200,7 @@ def _pipe_file_in_transaction( compression = kwargs.pop("compression", None) if compression is not None: # See S3FileSystem.pipe_file. - value = _compress(path, value, compression) + value = _compress(self._strip_protocol(path), value, compression) block_size = kwargs.get("block_size") or self._sync_fs.default_block_size # The size in bytes; the length of a memoryview counts its items. self._sync_fs._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 2acc72ac..08f1b2e1 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1539,7 +1539,14 @@ def test_pipe_file_failed_write(self, intrans): fs._call.assert_not_called() @pytest.mark.parametrize( - ("path", "compression"), [("s3://bucket/key", "gzip"), ("s3://bucket/key.gz", "infer")] + ("path", "compression"), + [ + ("s3://bucket/key", "gzip"), + ("s3://bucket/key.gz", "infer"), + # Inferred from the path without the trailing slash, as open() + # does. + ("s3://bucket/key.gz/", "infer"), + ], ) @pytest.mark.parametrize("intrans", [False, True]) @pytest.mark.parametrize("size", [1, S3FileSystem.DEFAULT_BLOCK_SIZE + 1]) @@ -1561,6 +1568,38 @@ def test_pipe_file_compression(self, path, compression, intrans, size): assert "compression" not in kwargs assert gzip.decompress(kwargs["body"]) == value + @pytest.mark.parametrize("intrans", [False, True]) + def test_pipe_file_compression_multipart(self, intrans): + # Compressed data larger than the block size is uploaded as a + # multipart upload. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs._transaction = None + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + fs._upload_part = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + fs._finish_multipart_upload = mock.MagicMock() + # Random bytes stay larger than the block size when compressed. + value = os.urandom(S3FileSystem.DEFAULT_BLOCK_SIZE + 1) + + with fs.transaction if intrans else contextlib.nullcontext(): + fs.pipe_file("s3://bucket/key", value, compression="gzip") + + body = b"".join(c.kwargs["body"] for c in fs._upload_part.call_args_list) + assert gzip.decompress(body) == value + fs._finish_multipart_upload.assert_called_once() + + def test_pipe_file_compression_non_contiguous_memoryview(self): + fs = self._make_fs() + fs._put_object = mock.MagicMock() + + fs.pipe_file("s3://bucket/key", memoryview(b"ab" * 4)[::2], compression="gzip") + + assert gzip.decompress(fs._put_object.call_args.kwargs["body"]) == b"aaaa" + def test_pipe_file_compression_inferred_none(self): # "infer" uploads the value as it is for a path without the # extension of a codec, as open() does. From 6d90268d601bd27f631049e1d72eb0e23325f451 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 00:12:27 +0900 Subject: [PATCH 05/11] Cover the inferred compression in an AioS3FileSystem transaction Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index a0ac8fad..880be02e 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -270,15 +270,24 @@ def test_put_file_mode(self, tmp_path, mode): assert "mode" not in call.kwargs assert call.kwargs.get("IfNoneMatch") == ("*" if mode == "create" else None) + @pytest.mark.parametrize( + ("path", "compression"), + [ + ("s3://bucket/key", "gzip"), + # Inferred from the path without the trailing slash, as open() + # does. + ("s3://bucket/key.gz/", "infer"), + ], + ) @pytest.mark.parametrize("intrans", [False, True]) - def test_pipe_file_compression(self, intrans): + def test_pipe_file_compression(self, path, compression, intrans): # GH-1037: the value is compressed before it is uploaded, also in a # transaction of this filesystem. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) put_object = fs._sync_fs._put_object = mock.MagicMock() with fs.transaction if intrans else contextlib.nullcontext(): - fs.pipe_file("s3://bucket/key", b"data", compression="gzip") + fs.pipe_file(path, b"data", compression=compression) ((_, kwargs),) = put_object.call_args_list assert "compression" not in kwargs From 407ec3edcd953f420b05c877eccda1abd7b7444d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 00:35:27 +0900 Subject: [PATCH 06/11] Route pipe_file() values by their size in bytes pipe_file() chose between PutObject and the buffered path by len(value), which counts the items of a memoryview. A typed memoryview larger than the block size in bytes, but not in items, went to PutObject, which accepts at most 5 GiB. Compare the size in bytes, as the multipart size check already does. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 5 +++-- tests/pyathena/filesystem/test_s3.py | 28 +++++++++++++++++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 17ead468..9025eead 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1825,8 +1825,9 @@ def pipe_file( value = _compress(self._strip_protocol(path), value, compression) block_size = kwargs.get("block_size") or self.default_block_size # The size in bytes; the length of a memoryview counts its items. - self._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) - if self._intrans or len(value) > min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): + size = memoryview(value).nbytes + self._check_multipart_upload_size(path, size, block_size) + if self._intrans or size > min(block_size, self.MULTIPART_UPLOAD_MAX_PART_SIZE): # Defer to the buffered open() path, which keeps the # deferred-commit semantics of fsspec transactions and uploads # large data as a parallel multipart upload. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 08f1b2e1..17af22e3 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1478,16 +1478,38 @@ def test_pipe_file_invalid_path_raises(self): fs.pipe_file("s3://bucket/key?versionId=12345abcde", b"data") def test_pipe_file_non_contiguous_memoryview(self): - # A non-contiguous memoryview within the block size in items, 4 items - # of 8 bytes here, is uploaded with PutObject. + # A non-contiguous memoryview within the block size, 4 items of 2 + # bytes here, is uploaded with PutObject. fs = self._make_fs() fs._put_object = mock.MagicMock() value = memoryview(b"ab" * 8).cast("H")[::2] - fs.pipe_file("s3://bucket/key", value, block_size=6) + fs.pipe_file("s3://bucket/key", value, block_size=8) fs._put_object.assert_called_once_with(bucket="bucket", key="key", body=b"ab" * 4) + def test_pipe_file_memoryview_routed_by_bytes(self): + # A memoryview larger than the block size in bytes, but not in items, + # is uploaded as a multipart upload. Its item count used to route it + # to PutObject, which accepts at most 5 GiB. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs._put_object = mock.MagicMock() + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + fs._upload_part = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + fs._finish_multipart_upload = mock.MagicMock() + data = b"a" * (S3FileSystem.DEFAULT_BLOCK_SIZE + 4) + + fs.pipe_file("s3://bucket/key", memoryview(data).cast("I")) + + fs._put_object.assert_not_called() + assert b"".join(c.kwargs["body"] for c in fs._upload_part.call_args_list) == data + fs._finish_multipart_upload.assert_called_once() + def test_pipe_file_small_drops_max_workers(self): fs = self._make_fs() fs._put_object = mock.MagicMock() From edc338db70f70ce38d2267953be9d898119c1c33 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 00:57:03 +0900 Subject: [PATCH 07/11] Make the compression a public CompressedBuffer API _compress() took the path only to resolve "infer" and returned the value unchanged when no codec was inferred. Resolve the codec in the callers with get_compression(), as open() does, and compress with the public CompressedBuffer.compress(), which takes a codec name and always returns bytes. s3_async no longer imports a private name from s3. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 61 +++++++++++++--------------- pyathena/filesystem/s3_async.py | 9 ++-- tests/pyathena/filesystem/test_s3.py | 43 +++++++++++++++----- 3 files changed, 65 insertions(+), 48 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 9025eead..6fe5b69c 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -51,47 +51,40 @@ _logger = logging.getLogger(__name__) -class _CompressedBuffer(BytesIO): - """A buffer that keeps its data when a codec closes it. +class CompressedBuffer(BytesIO): + """An in-memory buffer of data compressed with a codec of fsspec. - Some codecs, such as the ``zstandard`` stream writer, close the file - that they write to when they are closed. + The buffer keeps its data when it is closed, as some codecs, such as + the ``zstandard`` stream writer, close the file that they write to. """ @override def close(self) -> None: pass + @classmethod + def compress(cls, value: bytes | bytearray | memoryview, compression: str) -> bytes: + """Compress a value with a codec of fsspec. -def _compress( - path: str, value: bytes | bytearray | memoryview, compression: str -) -> bytes | bytearray | memoryview: - """Compress a value with the codec that ``open()`` uses for a compression. - - Args: - path: Path of the file without the protocol, as ``open()`` strips - it, from which ``"infer"`` takes the codec. - value: The bytes to compress. - compression: Name of a codec in ``fsspec.compression.compr``, or - ``"infer"`` to take it from the extension of the path. + Args: + value: The bytes to compress. + compression: Name of a codec in ``fsspec.compression.compr``. - Returns: - The compressed bytes, or the value itself when ``"infer"`` finds no - codec for the path. + Returns: + The compressed bytes. - Raises: - ValueError: If the codec is not supported. - """ - compression = get_compression(path, compression) - if compression is None: - return value - if isinstance(value, memoryview) and not value.c_contiguous: - # Codecs cannot compress a non-contiguous memoryview. - value = value.tobytes() - buffer = _CompressedBuffer() - with compr[compression](buffer, mode="w") as f: - f.write(value) - return buffer.getvalue() + Raises: + ValueError: If the codec is not supported. + """ + if compression not in compr: + raise ValueError(f"Compression type {compression} not supported") + if isinstance(value, memoryview) and not value.c_contiguous: + # Codecs cannot compress a non-contiguous memoryview. + value = value.tobytes() + buffer = cls() + with compr[compression](buffer, mode="w") as f: + f.write(value) + return buffer.getvalue() class S3FileSystem(AbstractFileSystem): @@ -1818,11 +1811,13 @@ def pipe_file( version, if the compression is not supported, or if the data takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ - compression = kwargs.pop("compression", None) + # The codec is taken as open() takes it, also from the extension of + # the path without the protocol for "infer". + compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None)) if compression is not None: # Compressed up front, so that every path uploads the compressed # bytes, and open() returns the file instead of a wrapper. - value = _compress(self._strip_protocol(path), value, compression) + value = CompressedBuffer.compress(value, compression) block_size = kwargs.get("block_size") or self.default_block_size # The size in bytes; the length of a memoryview counts its items. size = memoryview(value).nbytes diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index d8fea2c8..9e8d5b5b 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -18,8 +18,9 @@ from fsspec.asyn import AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK +from fsspec.core import get_compression -from pyathena.filesystem.s3 import S3File, S3FileSystem, _compress +from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import ( S3Metadata, @@ -197,10 +198,10 @@ def _pipe_file_in_transaction( ValueError: If the compression is not supported, or if the data takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ - compression = kwargs.pop("compression", None) + # See S3FileSystem.pipe_file. + compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None)) if compression is not None: - # See S3FileSystem.pipe_file. - value = _compress(self._strip_protocol(path), value, compression) + value = CompressedBuffer.compress(value, compression) block_size = kwargs.get("block_size") or self._sync_fs.default_block_size # The size in bytes; the length of a memoryview counts its items. self._sync_fs._check_multipart_upload_size(path, memoryview(value).nbytes, block_size) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 17af22e3..d244d35b 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,9 +1,11 @@ import asyncio +import bz2 import contextlib import functools import gc import gzip import io +import lzma import os import re import sys @@ -31,7 +33,7 @@ import pyathena from pyathena.filesystem import register_s3_filesystem -from pyathena.filesystem.s3 import S3File, S3FileSystem, _compress +from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass @@ -4673,14 +4675,33 @@ def test_upload_chunk_multipart(self, autocommit): file.fs._put_object.assert_not_called() -def test_compress_codec_closing_its_file(): - # GH-1037: some codecs, such as the zstandard stream writer, close the - # file that they write to when they are closed. - def closing_gzip(f, mode): - g = gzip.GzipFile(fileobj=f, mode=mode) - close = g.close - g.close = lambda: (close(), f.close()) - return g +class TestCompressedBuffer: + @pytest.mark.parametrize( + ("compression", "decompress"), + [("gzip", gzip.decompress), ("bz2", bz2.decompress), ("xz", lzma.decompress)], + ) + def test_compress(self, compression, decompress): + assert decompress(CompressedBuffer.compress(b"a" * 100, compression)) == b"a" * 100 + + def test_compress_non_contiguous_memoryview(self): + value = memoryview(b"ab" * 4)[::2] - with mock.patch.dict(compr, {"closing": closing_gzip}): - assert gzip.decompress(_compress("s3://bucket/key", b"a", "closing")) == b"a" + assert gzip.decompress(CompressedBuffer.compress(value, "gzip")) == b"aaaa" + + @pytest.mark.parametrize("compression", ["unknown", "infer"]) + def test_compress_unsupported(self, compression): + # "infer" is resolved from a path by the caller, as open() does. + with pytest.raises(ValueError, match="not supported"): + CompressedBuffer.compress(b"a", compression) + + def test_compress_codec_closing_its_file(self): + # GH-1037: some codecs, such as the zstandard stream writer, close the + # file that they write to when they are closed. + def closing_gzip(f, mode): + g = gzip.GzipFile(fileobj=f, mode=mode) + close = g.close + g.close = lambda: (close(), f.close()) + return g + + with mock.patch.dict(compr, {"closing": closing_gzip}): + assert gzip.decompress(CompressedBuffer.compress(b"a", "closing")) == b"a" From d8ec2207678815b2dee85386d379edd523fd804d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 01:08:12 +0900 Subject: [PATCH 08/11] Document the no-op close() of CompressedBuffer Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 6fe5b69c..a8dd630d 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -60,7 +60,7 @@ class CompressedBuffer(BytesIO): @override def close(self) -> None: - pass + """Do nothing, so that the data can still be read after a codec closes the buffer.""" @classmethod def compress(cls, value: bytes | bytearray | memoryview, compression: str) -> bytes: From 73cdb62537e951d8aa58a9baaf75cfe04599a9e1 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 01:19:09 +0900 Subject: [PATCH 09/11] Write pipe_file() values to the path without a trailing slash The single PutObject request wrote a path with a trailing slash to the key with the slash, while the buffered path, through open(), wrote it to the key without it, so the key depended on the size of the value. With the compression argument, compressing could now switch a call between the two. Normalize the path as open() does before choosing, so both paths write the same key. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 4 +-- pyathena/filesystem/s3.py | 11 +++++--- tests/pyathena/filesystem/test_s3.py | 39 +++++++++++++++++++++++----- 3 files changed, 41 insertions(+), 13 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 2b29b359..0190db65 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -109,8 +109,8 @@ and `open` treat `s3://YOUR_S3_BUCKET/dir/` as `s3://YOUR_S3_BUCKET/dir`: the ob `dir` if it exists, and otherwise the directory `dir`. An object whose key ends in a slash, such as a folder marker, is therefore not a file for these methods. Opening `dir/` for reading reads the object `dir` or raises `FileNotFoundError`, and opening it -for writing writes the object `dir`. A path with a `?versionId=` suffix keeps the slash -and refers to the object. `find`, and `ls` of the directory, list the object as a file +for writing, `pipe`, and `put` write the object `dir`. A path with a `?versionId=` +suffix keeps the slash and refers to the object. `find`, and `ls` of the directory, list the object as a file entry. `cat_file` uses the key as written. Without a `?versionId=` suffix, it reads such an object without a range, with a non-empty range of non-negative offsets, or with a negative `start` and no `end`, and raises `FileNotFoundError` for other ranges. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index a8dd630d..b310635f 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1786,7 +1786,8 @@ def pipe_file( and writes inside an fsspec transaction go through the buffered path, which uploads the data as a parallel multipart upload and keeps the deferred-commit semantics of transactions. A write that - fails on that path leaves the existing object unchanged. + fails on that path leaves the existing object unchanged. Both paths + write to the path without a trailing slash, as ``open()`` does. Args: path: S3 path (s3://bucket/key) to write to. @@ -1811,9 +1812,11 @@ def pipe_file( version, if the compression is not supported, or if the data takes more than ``MULTIPART_UPLOAD_MAX_PARTS`` blocks. """ - # The codec is taken as open() takes it, also from the extension of - # the path without the protocol for "infer". - compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None)) + # Normalized as open() normalizes it, so that the key written, and + # the codec that "infer" takes from its extension, do not depend on + # the path that the size of the value selects. + path = self._strip_protocol(path) + compression = get_compression(path, kwargs.pop("compression", None)) if compression is not None: # Compressed up front, so that every path uploads the compressed # bytes, and open() returns the file instead of a wrapper. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index d244d35b..d99057f9 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1490,6 +1490,30 @@ def test_pipe_file_non_contiguous_memoryview(self): fs._put_object.assert_called_once_with(bucket="bucket", key="key", body=b"ab" * 4) + @pytest.mark.parametrize("intrans", [False, True]) + @pytest.mark.parametrize("size", [1, S3FileSystem.DEFAULT_BLOCK_SIZE + 1]) + def test_pipe_file_trailing_slash(self, intrans, size): + # GH-1037: a path with a trailing slash is written without it, as + # open() writes it, whatever the size of the value. The single + # request used to write the key with the trailing slash. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs._transaction = None + fs._put_object = mock.MagicMock() + fs._create_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(upload_id="uploadid") + ) + fs._upload_part = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + fs._finish_multipart_upload = mock.MagicMock() + + with fs.transaction if intrans else contextlib.nullcontext(): + fs.pipe_file("s3://bucket/dir/key/", b"a" * size) + + calls = fs._put_object.call_args_list + fs._create_multipart_upload.call_args_list + assert [c.kwargs["key"] for c in calls] == ["dir/key"] + def test_pipe_file_memoryview_routed_by_bytes(self): # A memoryview larger than the block size in bytes, but not in items, # is uploaded as a multipart upload. Its item count used to route it @@ -1563,18 +1587,18 @@ def test_pipe_file_failed_write(self, intrans): fs._call.assert_not_called() @pytest.mark.parametrize( - ("path", "compression"), + ("path", "compression", "key"), [ - ("s3://bucket/key", "gzip"), - ("s3://bucket/key.gz", "infer"), - # Inferred from the path without the trailing slash, as open() - # does. - ("s3://bucket/key.gz/", "infer"), + ("s3://bucket/key", "gzip", "key"), + ("s3://bucket/key.gz", "infer", "key.gz"), + # Inferred from, and written to, the path without the trailing + # slash, as open() does. + ("s3://bucket/key.gz/", "infer", "key.gz"), ], ) @pytest.mark.parametrize("intrans", [False, True]) @pytest.mark.parametrize("size", [1, S3FileSystem.DEFAULT_BLOCK_SIZE + 1]) - def test_pipe_file_compression(self, path, compression, intrans, size): + def test_pipe_file_compression(self, path, compression, key, intrans, size): # GH-1037: the value is compressed before it is uploaded, on every # path. The single-request path used to send compression to # PutObject, which botocore rejects. @@ -1589,6 +1613,7 @@ def test_pipe_file_compression(self, path, compression, intrans, size): # The compressed value fits in one block. ((_, kwargs),) = fs._put_object.call_args_list + assert kwargs["key"] == key assert "compression" not in kwargs assert gzip.decompress(kwargs["body"]) == value From 67a3db4c921a4843cb8f37a137be9deb630b4107 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 01:26:16 +0900 Subject: [PATCH 10/11] Describe put() to a trailing-slash path as a write into the directory fsspec's put() treats a destination ending in a slash as a directory and writes the file under it, unlike put_file(). Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 0190db65..9a4c81ca 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -108,9 +108,10 @@ Paths are normalized as in fsspec, which drops a trailing slash, so `info`, `isf and `open` treat `s3://YOUR_S3_BUCKET/dir/` as `s3://YOUR_S3_BUCKET/dir`: the object `dir` if it exists, and otherwise the directory `dir`. An object whose key ends in a slash, such as a folder marker, is therefore not a file for these methods. Opening -`dir/` for reading reads the object `dir` or raises `FileNotFoundError`, and opening it -for writing, `pipe`, and `put` write the object `dir`. A path with a `?versionId=` -suffix keeps the slash and refers to the object. `find`, and `ls` of the directory, list the object as a file +`dir/` for reading reads the object `dir` or raises `FileNotFoundError`. Opening it for +writing, `pipe`, `pipe_file`, and `put_file` write the object `dir`, and `put` writes +the file into the directory `dir`. A path with a `?versionId=` suffix keeps the slash +and refers to the object. `find`, and `ls` of the directory, list the object as a file entry. `cat_file` uses the key as written. Without a `?versionId=` suffix, it reads such an object without a range, with a non-empty range of non-negative offsets, or with a negative `start` and no `end`, and raises `FileNotFoundError` for other ranges. From c6297d35b3552af208fb931ac8a60201178f47d2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 01:30:52 +0900 Subject: [PATCH 11/11] Describe only the PyAthena write methods for trailing-slash paths fsspec's put() resolves a destination ending in a slash itself, as a directory for a single source but not for paired lists, so the paragraph no longer describes it. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 9a4c81ca..9c05b9ce 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -109,9 +109,8 @@ and `open` treat `s3://YOUR_S3_BUCKET/dir/` as `s3://YOUR_S3_BUCKET/dir`: the ob `dir` if it exists, and otherwise the directory `dir`. An object whose key ends in a slash, such as a folder marker, is therefore not a file for these methods. Opening `dir/` for reading reads the object `dir` or raises `FileNotFoundError`. Opening it for -writing, `pipe`, `pipe_file`, and `put_file` write the object `dir`, and `put` writes -the file into the directory `dir`. A path with a `?versionId=` suffix keeps the slash -and refers to the object. `find`, and `ls` of the directory, list the object as a file +writing, `pipe`, `pipe_file`, and `put_file` write the object `dir`. A path with a +`?versionId=` suffix keeps the slash and refers to the object. `find`, and `ls` of the directory, list the object as a file entry. `cat_file` uses the key as written. Without a `?versionId=` suffix, it reads such an object without a range, with a non-empty range of non-negative offsets, or with a negative `start` and no `end`, and raises `FileNotFoundError` for other ranges.