diff --git a/docs/filesystem.md b/docs/filesystem.md index d69a230a6..9c05b9ce0 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 @@ -105,9 +108,9 @@ 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 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 +`dir/` for reading reads the object `dir` or raises `FileNotFoundError`. Opening it for +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. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 5434a2973..b310635ff 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 @@ -44,11 +46,47 @@ 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): + """An in-memory buffer of data compressed with a codec of fsspec. + + 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: + """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: + """Compress a value with a codec of fsspec. + + Args: + value: The bytes to compress. + compression: Name of a codec in ``fsspec.compression.compr``. + + Returns: + The compressed bytes. + + 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): """A filesystem interface for Amazon S3 that implements the fsspec protocol. @@ -1748,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. @@ -1760,20 +1799,33 @@ 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. """ + # 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. + 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. - 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/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index a7d933bc1..9e8d5b5b1 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 +from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import ( S3Metadata, @@ -187,14 +188,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. """ + # See S3FileSystem.pipe_file. + compression = get_compression(self._strip_protocol(path), kwargs.pop("compression", None)) + if compression is not None: + 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 4eed9dd05..d99057f9e 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,8 +1,11 @@ import asyncio +import bz2 import contextlib import functools import gc +import gzip import io +import lzma import os import re import sys @@ -24,12 +27,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 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 @@ -1476,16 +1480,62 @@ 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) + @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 + # 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() @@ -1536,6 +1586,113 @@ def test_pipe_file_failed_write(self, intrans): fs._put_object.assert_not_called() fs._call.assert_not_called() + @pytest.mark.parametrize( + ("path", "compression", "key"), + [ + ("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, 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. + 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 kwargs["key"] == key + 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. + 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. @@ -4541,3 +4698,35 @@ def test_upload_chunk_multipart(self, autocommit): file.commit() file.fs._finish_multipart_upload.assert_called_once() file.fs._put_object.assert_not_called() + + +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] + + 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" diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f65e72de1..880be02ef 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,29 @@ 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, 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(path, b"data", compression=compression) + + ((_, 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