pyathena s3 fsspec replacement file system not writing when used with transactions & a context manager

[dogversioning]

I have a project using fsspec via an abstraction layer. In one place, we provide an alias for fsspec's open method, as follows, using fsspec's transactions:

   @contextlib.contextmanager
    def open(self, mode: str = "r", **kwargs) -> typing.IO:
        """
        Opens the file.

        Writes are atomic by default (the target path will be unaffected until the open is closed).
        Compression will be automatically inferred based on suffix.
        """
        with contextlib.ExitStack() as stack:
            if "w" in mode:
                # fsspec is atomic per-transaction.
                # If an error occurs inside the transaction, partial writes will be discarded.
                # But we only want a transaction if we're writing - read transactions can error out
                stack.enter_context(self.fs.transaction)

            yield stack.enter_context(self.open_direct(mode, **kwargs))

    def open_direct(self, mode: str = "r", **kwargs) -> typing.IO:
        """
        Opens the file directly, without a context manager.

        You must manually call close() yourself.
        The operation will not be atomic.
        """
        # allow callers to override these defaults if they want
        kwargs.setdefault("compression", "infer")
        kwargs.setdefault("encoding", "utf8")
        #print(inspect.stack())
        return self.fs.open(self._path, mode=mode, **kwargs)

When we write to S3 using the transaction using fsspec directly, this works fine. When pyathena replaces fsspec, it this method creates empty files.

I'm working around this currently in the method described in this issue, manually instantiating a s3fs file system.

I think it would be preferable, from a maintenance standpoint, to not entirely replace fsspec when loading a PandasCursor, which would resolve this issue. Otherwise, to maintain API compatibility, it would be nice to support the usage of transactions in this way.

[laughingman7743]

Thanks for the detailed report — I can reproduce this and you're right that it's an implementation gap in PyAthena's S3 filesystem, not something you should have to work around.

Root cause

PyAthena ships its own fsspec.AbstractFileSystem implementation (pyathena/filesystem/s3.py) rather than delegating to s3fs. The S3File.commit() / discard() deferred-write path that fsspec transactions rely on is only correctly implemented for the multipart (large file) case. For files smaller than the block size (default 5 MiB) written inside a transaction, an empty object is created.

Mechanism

1. S3File.close() calls fsspec's flush(force=True).
2. flush() calls _upload_chunk(final=True) and then resets the buffer: self.buffer = io.BytesIO() (this is in fsspec core, AbstractBufferedFile.flush, and we can't change it).
3. For a small file, _upload_chunk defers the upload to commit() when autocommit=False (the transaction case) — it does not upload anything itself (pyathena/filesystem/s3.py:1411-1416).
4. On transaction exit, Transaction.complete() calls f.commit().
5. commit()'s small-file branch reads from self.buffer again (pyathena/filesystem/s3.py:1461-1476) — but the buffer was already reset in step 2, so body=b"" is sent to PutObject, producing an empty object.

For large files this works because the parts are already uploaded to S3 before commit() runs; commit() only needs to call CompleteMultipartUpload. Only the small-file deferred path loses the data.

Minimal reproduction (no AWS needed — pure buffer flow)

Same bytes (b"hello world"), the only difference is autocommit:

Path	body sent to PutObject
Normal open() (autocommit=True)	b"hello world" ✓
Transaction (autocommit=False)	b"" ✗ (empty object)

End-to-end this matches your case:

fsspec.register_implementation("s3", "pyathena.filesystem.s3.S3FileSystem", clobber=True)
fs = ...  # PyAthena S3FileSystem
with fs.transaction:
    with fs.open("s3://bucket/key.txt", "wb") as f:
        f.write(b"hello world")   # < 5 MiB
# -> object exists but is 0 bytes

Why we won't drop the custom filesystem

The custom S3FileSystem is intentional: it exists to avoid the s3fs / aiobotocore dependency and interference problems that motivated it (e.g. #344, #465, #476, #272). So the fix is to support transactions properly in our own implementation rather than handing filesystem ownership back to s3fs.

Implementation plan (for a follow-up PR)

Snapshot the buffered bytes at flush time (before fsspec resets the buffer) and have commit() use the snapshot.

1. S3File.__init__ (around pyathena/filesystem/s3.py:1356-1357): initialize self._commit_buffer: bytes | None = None.

2. S3File._upload_chunk small-file branch (pyathena/filesystem/s3.py:1411-1416): capture the bytes on the final flush, regardless of autocommit, so a deferred commit() can use them:

   if self.tell() < self.blocksize:
       # Files smaller than block size in size cannot be multipart uploaded.
       if final:
           # fsspec's flush() resets self.buffer right after this returns.
           # When the commit is deferred (autocommit=False, i.e. inside a
           # transaction), commit() would otherwise read the already-reset
           # buffer and upload an empty object. Snapshot the bytes now.
           self.buffer.seek(0)
           self._commit_buffer = self.buffer.read()
           if self.autocommit:
               self.commit()
       return True

3. S3File.commit small-file branch (pyathena/filesystem/s3.py:1466-1476): prefer the snapshot, falling back to the buffer:

   elif not self.multipart_upload_parts:
       if self.buffer is not None:
           if self._commit_buffer is not None:
               data = self._commit_buffer
           else:
               self.buffer.seek(0)
               data = self.buffer.read()
           self.fs._put_object(
               bucket=self.bucket, key=self.key, body=data,
               **self.s3_additional_kwargs,
           )

Notes:

• The tell() == 0 branch (intentional empty file via touch) is unaffected: the snapshot is b"" and that branch is selected first.
• discard() (transaction rollback) for small files already does nothing — no multipart, no PutObject — which is the correct atomic behavior (the object is never created on rollback).
• The multipart/large-file transaction path already works and is untouched.

Test plan

• Unit test (no AWS): drive S3File with mocked _put_object / touch, write < blocksize bytes with autocommit=False, then call commit() (simulating Transaction.complete()), and assert the real bytes are sent — covering the regression directly.
• Integration test mirroring test_write in tests/pyathena/filesystem/test_s3.py: write a small file inside with fs.transaction: and assert the round-tripped content matches.
• Also assert rollback: raising inside the transaction leaves no object (small-file discard() path).

[laughingman7743]

Planning to implement this over the weekend. I'll open a draft PR following the plan above.

[dogversioning]

Why we won't drop the custom filesystem

The custom S3FileSystem is intentional: it exists to avoid the s3fs / aiobotocore dependency and interference problems that motivated it (e.g. #344, #465, #476, #272). So the fix is to support transactions properly in our own implementation rather than handing filesystem ownership back to s3fs.

This is reasonable (and, indeed managing those dependencies is a headache with pip - as an FYI of a thing to keep an eye on, I found out in my testing of this that uv is much better at resolving this dependency chain), but I will note that replacing it by namespace interferes with fsspec's registry lookup. If you were to try and fetch an S3 filesystem in this way:

import fsspec

fs = fsspec.filesystem('s3')

I think you would reasonably expect to get back the native s3fs filesystem by calling this function, and not the pyathena version of the same. This also impacts debugging - I was trying to use the s3fs log debugging, and was getting confused by no output when S3FS_LOGGING_LEVEL was set to DEBUG.

If you want to keep things the way they are now, it might be nice to emit a suppressible message on registry lookup, and check for that environment variable and log out that the swap has occurred, just so a user can be informed and react in a way that makes sense for their context.