Keep the existing object when appending within a larger block size

pyathena/filesystem/s3.py

@@ -2111,8 +2111,9 @@ def __init__(
         In read mode, the object is looked up with ``info()`` and the reads
         are made conditional on its ETag (``IfMatch``). In append mode, an
         existing object smaller than ``MULTIPART_UPLOAD_MIN_PART_SIZE`` is
-        read into the write buffer; a larger one is copied as the first parts
-        once a multipart upload starts.
+        read into the write buffer; a larger one is copied with
+        ``UploadPartCopy`` as the first parts of a multipart upload, whatever
+        the block size.
 
         Args:
             fs: The filesystem that the file belongs to.
@@ -2188,11 +2189,15 @@ def __init__(
                 self.s3_additional_kwargs.update({"IfMatch": etag})
             self._details = info
         elif "a" in mode and self.fs.exists(path):
-            self.append_block = True
             info = self.fs.info(self.path, version_id=self.version_id)
             loc = info.get("size", 0)
             if loc < self.fs.MULTIPART_UPLOAD_MIN_PART_SIZE:
+                # Too small to be a part of a multipart upload: rewrite it
+                # from the buffer.
                 self.write(self.fs.cat(self.path))
+            else:
+                # Copied with UploadPartCopy as the leading part(s).
+                self.append_block = True
             self.loc = loc
             self.s3_additional_kwargs.update(info.to_api_repr())
             self._details = info
@@ -2208,8 +2213,10 @@ def close(self) -> None:
         self._executor.shutdown()
 
     def _initiate_upload(self) -> None:
-        if self.tell() < self.blocksize:
+        if not self.append_block and self.tell() < self.blocksize:
             # Files smaller than block size in size cannot be multipart uploaded.
+            # An append to an object copied with UploadPartCopy always uses
+            # a multipart upload, whatever the block size.
             return
 
         self.multipart_upload = self.fs._create_multipart_upload(
@@ -2259,7 +2266,7 @@ def _upload_chunk(self, final: bool = False) -> bool:
         # can still read the bytes; resetting it there would upload an empty
         # object for small files. Mid-stream chunks (final=False) return True so
         # fsspec clears the already-uploaded buffer between parts.
-        if self.tell() < self.blocksize:
+        if not self.append_block and self.tell() < self.blocksize:
             # Files smaller than block size in size cannot be multipart uploaded.
             if self.autocommit and final:
                 self.commit()
@@ -2360,12 +2367,19 @@ def discard(self) -> None:
         if self.multipart_upload:
             for f in self.multipart_upload_parts:
                 f.cancel()
+            # s3_additional_kwargs also holds object parameters (e.g., the
+            # existing object's metadata in append mode) that
+            # AbortMultipartUpload rejects.
             self.fs._call(
                 "abort_multipart_upload",
                 Bucket=self.bucket,
                 Key=self.key,
                 UploadId=self.multipart_upload.upload_id,
-                **self.s3_additional_kwargs,
+                **{
+                    k: v
+                    for k, v in self.s3_additional_kwargs.items()
+                    if k in ("RequestPayer", "ExpectedBucketOwner")
+                },
             )
 
         self.multipart_upload = None

tests/pyathena/filesystem/test_s3.py

@@ -676,6 +676,65 @@ def test_append(self, fs, base, exp):
             assert len(actual) == len(data + extra)
             assert actual == data + extra
 
+    @pytest.mark.parametrize(
+        ("size", "extra_size", "block_size"),
+        [
+            # GH-921: an existing object of at least 5 MiB, appended within a
+            # larger block size, is copied with UploadPartCopy.
+            (6 * 2**20, 5, 16 * 2**20),
+            # An existing object smaller than 5 MiB is rewritten from the
+            # buffer, not copied as well, when the append crosses the block size.
+            (2**10, 5 * 2**20, None),
+        ],
+    )
+    def test_append_with_block_size(self, fs, size, extra_size, block_size):
+        data = b"a" * size
+        extra = b"b" * extra_size
+        path = (
+            f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/"
+            f"filesystem/test_append_with_block_size/{uuid.uuid4()}"
+        )
+        fs.pipe_file(path, data)
+        with fs.open(path, "ab", block_size=block_size) as f:
+            f.write(extra)
+        # Check the size and the bytes at the ends and around the boundary
+        # instead of reading the whole object back, to keep the transfer small.
+        assert fs.info(path, refresh=True).size == size + extra_size
+        assert fs.cat_file(path, start=0, end=1) == b"a"
+        assert fs.cat_file(path, start=size - 1, end=size + 1) == b"ab"
+        assert fs.cat_file(path, start=-1) == b"b"
+
+    @pytest.mark.parametrize("block_size", [None, 16 * 2**20])
+    def test_append_transaction_rollback(self, fs, block_size):
+        # Raising inside the transaction aborts the multipart upload that
+        # copies the existing object and leaves the object unchanged.
+        data = b"a" * (6 * 2**20)
+        path = (
+            f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/"
+            f"filesystem/test_append_transaction_rollback/{uuid.uuid4()}"
+        )
+        fs.pipe_file(path, data)
+        before = fs.info(path, refresh=True)
+
+        def append_then_fail():
+            with fs.transaction:
+                f = fs.open(path, "ab", block_size=block_size)
+                f.write(b"b" * 5)
+                f.close()
+                raise RuntimeError("rollback")
+
+        with pytest.raises(RuntimeError):
+            append_then_fail()
+        # A committed append (a multipart upload, or the appended bytes alone)
+        # would change the ETag and the size, so the object is not read back.
+        after = fs.info(path, refresh=True)
+        assert (after.etag, after.last_modified, after.size) == (
+            before.etag,
+            before.last_modified,
+            before.size,
+        )
+        assert not fs.list_multipart_uploads(path)
+
     def test_ls_buckets(self, fs):
         fs.invalidate_cache()
         actual = fs.ls("s3://")
@@ -1511,6 +1570,7 @@ def _make_write_file(data: bytes, autocommit: bool):
         file.s3_additional_kwargs = {}
         file.autocommit = autocommit
         file.blocksize = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE
+        file.append_block = False
         file.multipart_upload = None
         file.multipart_upload_parts = []
         file.buffer = io.BytesIO(data)
@@ -1534,6 +1594,102 @@ def _make_multipart_write_file(data: bytes, autocommit: bool):
         )
         return file
 
+    @staticmethod
+    def _make_append_fs(existing: bytes):
+        # A mocked filesystem holding an existing object, with a minimum part
+        # size of 4 bytes so that the append paths can be exercised with tiny
+        # data and no AWS access.
+        fs = mock.MagicMock(spec=S3FileSystem)
+        fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4
+        fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 64
+        fs.exists.return_value = True
+        fs.info.return_value = S3Object(
+            init={"ContentLength": len(existing)},
+            type=S3ObjectType.S3_OBJECT_TYPE_FILE,
+            bucket="bucket",
+            key="key.txt",
+        )
+        fs.cat.return_value = existing
+        fs._create_multipart_upload.return_value = SimpleNamespace(upload_id="uploadid")
+
+        def part(**kw):
+            return SimpleNamespace(etag=f'"e{kw["part_number"]}"', part_number=kw["part_number"])
+
+        fs._upload_part.side_effect = part
+        fs._upload_part_copy.side_effect = part
+        return fs
+
+    @staticmethod
+    def _uploaded_object(fs, existing: bytes) -> bytes:
+        # Rebuild the object S3 would store from the mocked upload calls.
+        # A part copy without a range copies the whole existing object.
+        if fs._put_object.called:
+            fs._create_multipart_upload.assert_not_called()
+            return fs._put_object.call_args.kwargs["body"]
+        fs._finish_multipart_upload.assert_called_once()
+        parts = [(c.kwargs["part_number"], existing) for c in fs._upload_part_copy.call_args_list]
+        parts += [
+            (c.kwargs["part_number"], c.kwargs["body"]) for c in fs._upload_part.call_args_list
+        ]
+        part_numbers = sorted(n for n, _ in parts)
+        assert part_numbers == list(range(1, len(parts) + 1))
+        return b"".join(body for _, body in sorted(parts))
+
+    @pytest.mark.parametrize(
+        ("existing", "appended", "multipart", "part_copy"),
+        [
+            # Smaller than the minimum part size: read into the buffer.
+            (b"aa", b"bb", False, False),
+            # GH-921: an existing object of at least the minimum part size is
+            # copied with UploadPartCopy even when the block size is larger
+            # than the whole object.
+            (b"a" * 6, b"bb", True, True),
+            (b"a" * 6, b"", True, True),
+            # An existing object read into the buffer is not copied again
+            # when the append crosses the block size.
+            (b"aa", b"b" * 16, True, False),
+        ],
+    )
+    def test_append(self, existing, appended, multipart, part_copy):
+        fs = self._make_append_fs(existing)
+
+        with S3File(fs, "s3://bucket/key.txt", mode="ab", block_size=16) as f:
+            f.write(appended)
+
+        assert self._uploaded_object(fs, existing) == existing + appended
+        assert fs._create_multipart_upload.called is multipart
+        assert fs._upload_part_copy.called is part_copy
+        fs.touch.assert_not_called()
+
+    def test_append_discard(self):
+        # Rolling back an append aborts its multipart upload without the
+        # existing object's metadata, which AbortMultipartUpload rejects,
+        # but with the request parameters it accepts.
+        fs = self._make_append_fs(b"a" * 6)
+        f = S3File(
+            fs,
+            "s3://bucket/key.txt",
+            mode="ab",
+            block_size=16,
+            autocommit=False,
+            s3_additional_kwargs={"RequestPayer": "requester", "ExpectedBucketOwner": "123"},
+        )
+        f.write(b"bb")
+        f.close()
+
+        f.discard()
+
+        fs._call.assert_called_once_with(
+            "abort_multipart_upload",
+            Bucket="bucket",
+            Key="key.txt",
+            UploadId="uploadid",
+            RequestPayer="requester",
+            ExpectedBucketOwner="123",
+        )
+        fs._finish_multipart_upload.assert_not_called()
+        fs._put_object.assert_not_called()
+
     @pytest.mark.parametrize(
         ("objects", "target"),
         [