Compress pipe_file() values up front, and route and key them consistently

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.

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.

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)