From 1a08eb1f8bb59dff9b0c7c3fb6b087129271bd37 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 10:30:49 +0900
Subject: [PATCH 1/7] Align AioS3FileSystem with S3FileSystem for cancellation
 and arguments (GH-1337)

A cancelled coroutine of the filesystem now waits for its blocking call,
whose stop scope keeps a write from sending further parts or committing,
before it re-raises. The batch coroutines of fsspec cancel their
per-object tasks and wait for them when they are cancelled.

A recursive _put() no longer creates a directory, and thereby looks up or
creates the bucket, for each local directory. _expand_path() runs
S3FileSystem.expand_path(), so that its keyword arguments reach the
listing, and S3FileSystem.glob() lists only the keys that start with the
stem before the first wildcard, as fsspec's asynchronous glob does.
_rm() takes batch_size and the deprecated max_workers as S3FileSystem.rm()
does, and _cp_file() no longer sends the unused recursive, maxdepth, and
on_error arguments to CopyObject.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/aio.md                                   |  29 +-
 pyathena/filesystem/s3.py                     | 120 +++++--
 pyathena/filesystem/s3_async.py               | 292 +++++++++++++++---
 pyathena/filesystem/s3_executor.py            |  26 +-
 tests/pyathena/filesystem/test_s3.py          |  11 +
 tests/pyathena/filesystem/test_s3_async.py    | 252 ++++++++++++++-
 tests/pyathena/filesystem/test_s3_executor.py |  44 ++-
 7 files changed, 679 insertions(+), 95 deletions(-)

diff --git a/docs/aio.md b/docs/aio.md
index d37bdfe5d..85d45673d 100644
--- a/docs/aio.md
+++ b/docs/aio.md
@@ -283,7 +283,7 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized
 | **Blocking calls** | Caller's thread | Thread pool of the filesystem |
 | **Parallelism** | Request pool of the filesystem | `asyncio.gather` on a thread pool; requests in the request pool |
 | **File handles** | `S3File` with thread pool | `AioS3File` with thread pool; `AioS3StreamedFile` from `open_async()` |
-| **Bulk delete** | Request pool, `max_workers` requests at once | Request pool, `max_workers` requests at once |
+| **Bulk delete** | Request pool, `batch_size` requests at once | Request pool, `batch_size` requests at once |
 | **Multipart copy** | Request pool, `max_workers` requests per copy | Request pool, `max_workers` requests per copy |
 | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) |
 
@@ -330,14 +330,25 @@ async with AioS3FileSystem(asynchronous=True) as fs:
 Cancelling a coroutine of the filesystem also stops the retries of the S3 requests that
 its call makes in these threads, including the part uploads and range reads that the
 call runs in parallel: no further attempt is made, and the wait before the next attempt
-ends. A request that has already been sent still runs to its end in its thread. The
-requests that run to their end after a cancellation keep their retries: the commit of a
-file, the CompleteMultipartUpload request of a multipart copy, and the aborts of
-multipart uploads.
-The range reads of a file from `open_async()`, the DeleteObjects requests of `_rm()`, and the annotation
-copies of a multipart copy are waited for before the cancellation is raised, also when
-the shutdown of the event loop at the end of `asyncio.run()` cancels the coroutine; a
-second cancellation stops that wait.
+ends. A write, such as one of `_put_file()` or `_pipe_file()`, sends no further part and
+aborts its multipart upload instead of completing it, and `_get_file()` reads and writes
+no further block. A request that has already been sent still runs to its end in its
+thread. The requests that run to their end after a cancellation keep their retries: the
+commit of a file, the CompleteMultipartUpload request of a multipart copy, and the
+aborts of multipart uploads.
+The call, the range reads of a file from `open_async()`, the DeleteObjects requests of
+`_rm()`, and the annotation copies of a multipart copy are waited for before the
+cancellation is raised, also when the shutdown of the event loop at the end of
+`asyncio.run()` cancels the coroutine; a second cancellation stops that wait.
+
+`_cat()`, `_cat_ranges()`, `_pipe()`, `_put()`, `_get()`, `_copy()`, and `_mv()` run the
+coroutines of their objects as tasks, `batch_size` of them at a time as in fsspec, except
+`_mv()`, which runs all its copies at once. Cancelling one of them cancels those tasks
+and waits for them before the cancellation is raised, also when the shutdown of the
+event loop cancels them together with the coroutine.
+A recursive `_put()` sends no request for the local directories. `_rm()` sends
+`batch_size` DeleteObjects requests at a time, `max_workers` by default and with `-1`; its
+`max_workers` argument is a deprecated alias of `batch_size`.
 
 ### Executor strategy
 
diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py
index af67f7fd6..a142e0b84 100644
--- a/pyathena/filesystem/s3.py
+++ b/pyathena/filesystem/s3.py
@@ -1127,6 +1127,31 @@ def find(
             return {f.name: f for f in files}
         return [f.name for f in files]
 
+    def glob(self, path, maxdepth=None, **kwargs):
+        """Find the paths that match a glob pattern.
+
+        As fsspec's ``glob()``, except that only the keys that start with the
+        part of the last path component before the first wildcard are listed,
+        as with fsspec's asynchronous ``_glob()``.
+
+        Args:
+            path: S3 glob pattern.
+            maxdepth: Maximum depth of the paths below the literal part of the
+                pattern, at least 1.
+            **kwargs: Additional arguments passed to ``find``, such as
+                ``detail``.
+
+        Returns:
+            The matching paths, or a dictionary of them and their objects with
+            ``detail=True``.
+        """
+        stripped = self._strip_protocol(path)
+        wildcards = [i for i in (stripped.find(c) for c in "*?[") if i >= 0]
+        stem = stripped[: min(wildcards)].rpartition("/")[2] if wildcards else ""
+        if stem and "prefix" not in kwargs:
+            kwargs["prefix"] = stem
+        return super().glob(path, maxdepth=maxdepth, **kwargs)
+
     def expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> list[str]:
         """Expand glob patterns and directories into the paths they match.
 
@@ -1288,18 +1313,35 @@ def rm(self, path, recursive=False, maxdepth=None, batch_size=None, **kwargs) ->
             TypeError: If both ``batch_size`` and ``max_workers`` are given.
             OSError: If S3 could not delete some of the objects.
         """
+        max_workers = self._delete_batch_limit(batch_size, kwargs)
+        paths = self._delete_paths(path, recursive=recursive, maxdepth=maxdepth)
+        self._delete_objects(paths, max_workers=max_workers, **kwargs)
+
+    def _delete_batch_limit(self, batch_size: int | None, kwargs: dict[str, Any]) -> int:
+        """Return the number of DeleteObjects requests of ``rm()`` to send at once.
+
+        Args:
+            batch_size: The ``batch_size`` argument of ``rm()``.
+            kwargs: The other arguments of ``rm()``, from which its deprecated
+                alias ``max_workers`` is removed.
+
+        Returns:
+            The number of requests.
+
+        Raises:
+            ValueError: If the number is neither positive, -1, nor None.
+            TypeError: If both ``batch_size`` and ``max_workers`` are given.
+        """
         if "max_workers" in kwargs:
             if batch_size is not None:
                 raise TypeError("rm() got both batch_size and its alias max_workers.")
             warnings.warn(
                 "The max_workers argument of rm() is deprecated; use batch_size.",
                 DeprecationWarning,
-                stacklevel=2,
+                stacklevel=3,
             )
             batch_size = kwargs.pop("max_workers")
-        max_workers = self._batch_limit(batch_size)
-        paths = self._delete_paths(path, recursive=recursive, maxdepth=maxdepth)
-        self._delete_objects(paths, max_workers=max_workers, **kwargs)
+        return self._batch_limit(batch_size)
 
     def _delete_paths(
         self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None
@@ -2671,32 +2713,7 @@ def put(
                 ``batch_size`` is neither positive, -1, nor None.
         """
         limit = self._batch_limit(batch_size)
-        if isinstance(lpath, list) and isinstance(rpath, list):
-            pairs = [(lp, self._strip_protocol(rp)) for lp, rp in zip(lpath, rpath, strict=False)]
-        else:
-            # As fsspec's put() pairs them.
-            source_is_str = isinstance(lpath, str)
-            if source_is_str:
-                lpath = make_path_posix(lpath)
-            local = LocalFileSystem()
-            lpaths = local.expand_path(lpath, recursive=recursive, maxdepth=maxdepth, **kwargs)
-            if source_is_str and (not recursive or maxdepth is not None):
-                # A string source without recursion does not copy directories.
-                lpaths = [p for p in lpaths if not (trailing_sep(p) or local.isdir(p))]
-                if not lpaths:
-                    return
-            dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or self.isdir(rpath))
-            rpath = (
-                self._strip_protocol(rpath)
-                if isinstance(rpath, str)
-                else [self._strip_protocol(p) for p in rpath]
-            )
-            exists = source_is_str and (
-                (has_magic(lpath) and len(lpaths) == 1)
-                or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath))
-            )
-            rpaths = other_paths(lpaths, rpath, exists=exists, flatten=not source_is_str)
-            pairs = list(zip(lpaths, rpaths, strict=True))
+        pairs = self._put_pairs(lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs)
         if not pairs:
             return
         if len({rp for _, rp in pairs}) != len(pairs):
@@ -2719,6 +2736,49 @@ def put_pair(pair: tuple[str, str]) -> None:
         locked.set_size(len(pairs))
         self._run_batch(put_pair, pairs, concurrency)
 
+    def _put_pairs(
+        self, lpath, rpath, recursive=False, maxdepth=None, **kwargs
+    ) -> list[tuple[str, str]]:
+        """Pair the local sources of ``put()`` with their S3 destinations as fsspec does.
+
+        Args:
+            lpath: Local source path, glob pattern, or list of them.
+            rpath: S3 destination path, or list of paths when ``lpath`` is a
+                list.
+            recursive: Whether to include the directories with their contents.
+            maxdepth: Maximum depth of a recursive copy.
+            **kwargs: Additional arguments passed to the ``expand_path()`` of
+                the local filesystem.
+
+        Returns:
+            The local paths, including the directories of a recursive copy,
+            with their destinations without the protocol.
+        """
+        if isinstance(lpath, list) and isinstance(rpath, list):
+            return [(lp, self._strip_protocol(rp)) for lp, rp in zip(lpath, rpath, strict=False)]
+        source_is_str = isinstance(lpath, str)
+        if source_is_str:
+            lpath = make_path_posix(lpath)
+        local = LocalFileSystem()
+        lpaths = local.expand_path(lpath, recursive=recursive, maxdepth=maxdepth, **kwargs)
+        if source_is_str and (not recursive or maxdepth is not None):
+            # A string source without recursion does not copy directories.
+            lpaths = [p for p in lpaths if not (trailing_sep(p) or local.isdir(p))]
+            if not lpaths:
+                return []
+        dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or self.isdir(rpath))
+        rpath = (
+            self._strip_protocol(rpath)
+            if isinstance(rpath, str)
+            else [self._strip_protocol(p) for p in rpath]
+        )
+        exists = source_is_str and (
+            (has_magic(lpath) and len(lpaths) == 1)
+            or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath))
+        )
+        rpaths = other_paths(lpaths, rpath, exists=exists, flatten=not source_is_str)
+        return list(zip(lpaths, rpaths, strict=True))
+
     def put_file(
         self,
         lpath: str,
diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 34add38cb..d27b5be9d 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -11,15 +11,16 @@
 
 import asyncio
 import contextlib
+import contextvars
 import functools
 import logging
 import mimetypes
 import os
 import warnings
-from collections.abc import Callable, Iterable
+from collections.abc import Awaitable, Callable, Iterable
 from concurrent.futures import Future, wait
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, sync
 from fsspec.callbacks import _DEFAULT_CALLBACK
@@ -51,6 +52,68 @@
 _logger = logging.getLogger(__name__)
 
 
+class _S3BatchTasks:
+    """The tasks of the per-object calls of a batch coroutine of fsspec.
+
+    fsspec runs the per-object coroutines of its batch coroutines, such as
+    ``_cat()`` and ``_put()``, as tasks, and a cancellation of the batch
+    coroutine leaves them running. Each per-object coroutine of
+    ``AioS3FileSystem`` adds its task with :meth:`track` to the batch that
+    runs in its context, which cancels them and waits for them when the
+    batch coroutine fails or is cancelled.
+    """
+
+    _current: ClassVar[contextvars.ContextVar[_S3BatchTasks | None]] = contextvars.ContextVar(
+        "_S3BatchTasks", default=None
+    )
+
+    def __init__(self) -> None:
+        """Initialize the batch with no task."""
+        self._tasks: set[asyncio.Task[Any]] = set()
+        self._owner: asyncio.Task[Any] | None = None
+
+    @classmethod
+    def track(cls) -> None:
+        """Add the current task to the batch that runs in its context, if any.
+
+        A task that the batch coroutine itself runs in is not added.
+        """
+        batch = cls._current.get()
+        task = asyncio.current_task()
+        if batch is not None and task is not None and task is not batch._owner:
+            batch._tasks.add(task)
+            task.add_done_callback(batch._tasks.discard)
+
+    async def run(self, batch: Awaitable[T]) -> T:
+        """Await a batch coroutine, and on its failure or cancellation wait for its tasks.
+
+        A task that is already being cancelled, as the shutdown of the event
+        loop with ``asyncio.run()`` cancels every task, is not cancelled
+        again, so that the wait for its own calls is not stopped. A repeated
+        cancellation stops the wait for the tasks.
+
+        Args:
+            batch: The batch coroutine, created and not yet started.
+
+        Returns:
+            The return value of the batch coroutine.
+        """
+        self._owner = asyncio.current_task()
+        token = self._current.set(self)
+        try:
+            return await batch
+        except BaseException:
+            tasks = [task for task in self._tasks if not task.done()]
+            for task in tasks:
+                if not task.cancelling():
+                    task.cancel()
+            if tasks:
+                await asyncio.wait(tasks)
+            raise
+        finally:
+            self._current.reset(token)
+
+
 class AioS3FileSystem(AsyncFileSystem):
     """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
 
@@ -203,6 +266,7 @@ async def _ls(self, path: str, detail: bool = False, **kwargs) -> list[S3Object]
     async def _cat_file(
         self, path: str, start: int | None = None, end: int | None = None, **kwargs
     ) -> bytes:
+        _S3BatchTasks.track()
         return await self._thread_pool.run(
             self._sync_fs.cat_file, path, start=start, end=end, **kwargs
         )
@@ -216,6 +280,7 @@ async def _rm_file(self, path: str, **kwargs) -> None:
     async def _pipe_file(
         self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs
     ) -> None:
+        _S3BatchTasks.track()
         if self._intrans:
             # The transaction belongs to this filesystem, not to the internal
             # S3FileSystem, so write through open() to defer the commit to it.
@@ -262,6 +327,7 @@ async def _put_file(
         mode: str = "overwrite",
         **kwargs,
     ) -> None:
+        _S3BatchTasks.track()
         if self._intrans:
             # See _pipe_file.
             await self._thread_pool.run(
@@ -327,6 +393,7 @@ def _put_file_in_transaction(
         self.invalidate_cache(rpath)
 
     async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
+        _S3BatchTasks.track()
         await self._thread_pool.run(
             self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs
         )
@@ -350,34 +417,45 @@ async def _rm(
         """Delete objects with DeleteObjects requests.
 
         See :meth:`S3FileSystem.rm`. The DeleteObjects requests run in
-        parallel, at most ``max_workers`` at once, in the request pool of the
+        parallel, at most ``batch_size`` at once, in the request pool of the
         internal ``S3FileSystem``. A cancellation cancels the requests that
         have not started and waits for the others.
 
         Args:
             path: S3 path (s3://bucket/key) or list of paths to delete.
             recursive: Whether to delete all objects below the paths.
-            batch_size: Accepted for fsspec compatibility; not used.
+            batch_size: Number of DeleteObjects requests to send at the same
+                time. None or -1 is ``max_workers``; the ``batch_size`` of
+                the filesystem, which is that of fsspec's other batch
+                coroutines, does not apply. ``max_workers`` is a deprecated
+                alias.
             maxdepth: Maximum depth to expand when ``recursive`` is True.
             **kwargs: Additional parameters passed to the DeleteObjects API.
                 ``Quiet`` (default True) sets the quiet mode of the requests.
 
         Raises:
-            ValueError: If a path is a bucket.
+            ValueError: If a path is a bucket, or ``batch_size`` is neither
+                positive, -1, nor None.
+            TypeError: If both ``batch_size`` and ``max_workers`` are given.
             OSError: If S3 could not delete some of the objects.
         """
+        max_workers = self._sync_fs._delete_batch_limit(batch_size, kwargs)
         paths = await self._thread_pool.run(
             self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth
         )
-        await self._delete_objects(paths, **kwargs)
+        await self._delete_objects(paths, max_workers=max_workers, **kwargs)
 
-    async def _delete_objects(self, paths: list[str], **kwargs) -> None:
+    async def _delete_objects(
+        self, paths: list[str], max_workers: int | None = None, **kwargs
+    ) -> None:
         """Delete objects with parallel DeleteObjects requests.
 
         See :meth:`S3FileSystem._delete_objects`.
 
         Args:
             paths: Paths of the objects to delete. Bucket paths are skipped.
+            max_workers: Maximum number of parallel requests. Defaults to the
+                ``max_workers`` of the internal ``S3FileSystem``.
             **kwargs: Additional parameters passed to the DeleteObjects API.
                 ``Quiet`` (default True) sets the quiet mode of the requests.
 
@@ -394,6 +472,7 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None:
         results = await self._map_async(
             functools.partial(self._sync_fs._delete_batch, **params),
             batches,
+            max_workers=max_workers,
             stop_on_failure=False,
         )
         self._sync_fs._raise_delete_errors(results.outcomes)
@@ -431,8 +510,9 @@ async def _map_async(
     async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None:
         """Move files from one S3 location to another.
 
-        See :meth:`S3FileSystem.mv`. The copies and the deletions run in
-        parallel with ``asyncio.gather``.
+        See :meth:`S3FileSystem.mv`. The copies run in parallel with
+        ``asyncio.gather``, whose cancellation cancels them and waits for
+        them, and the deletions in the request pool.
 
         Args:
             path1: Source S3 path, glob pattern, or list of paths.
@@ -506,7 +586,10 @@ async def _copy(
         """Copy files within S3.
 
         See :meth:`S3FileSystem.copy`. The copies run as in fsspec's
-        ``_copy()``.
+        ``_copy()``, ``batch_size`` of them at once as tasks. A cancellation
+        cancels the copies and waits for them; a copy that is already being
+        cancelled, such as by the shutdown of the event loop, is not
+        cancelled again.
 
         Args:
             path1: Source S3 path, glob pattern, or list of them.
@@ -527,14 +610,16 @@ async def _copy(
             if not pairs:
                 return
             path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs]
-        await super()._copy(
-            path1,
-            path2,
-            recursive=recursive,
-            on_error=on_error,
-            maxdepth=maxdepth,
-            batch_size=batch_size,
-            **kwargs,
+        await _S3BatchTasks().run(
+            super()._copy(
+                path1,
+                path2,
+                recursive=recursive,
+                on_error=on_error,
+                maxdepth=maxdepth,
+                batch_size=batch_size,
+                **kwargs,
+            )
         )
 
     async def _get(
@@ -543,7 +628,8 @@ async def _get(
         """Copy files from S3 to the local filesystem.
 
         See :meth:`S3FileSystem.get`. The downloads run as in fsspec's
-        ``_get()``.
+        ``_get()``, and a cancellation cancels them and waits for them, as
+        in :meth:`_copy`.
 
         Args:
             rpath: Source S3 path, glob pattern, or list of them.
@@ -570,11 +656,137 @@ async def _get(
             check_contained(root, lpath)
             if not rpath:
                 return
-        await super()._get(
-            rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs
+        await _S3BatchTasks().run(
+            super()._get(
+                rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs
+            )
         )
 
-    async def _cp_file(self, path1: str, path2: str, **kwargs) -> None:
+    async def _put(
+        self,
+        lpath,
+        rpath,
+        recursive=False,
+        callback=_DEFAULT_CALLBACK,
+        batch_size=None,
+        maxdepth=None,
+        **kwargs,
+    ) -> None:
+        """Copy files from the local filesystem to S3.
+
+        See :meth:`S3FileSystem.put`. The uploads run as in fsspec's
+        ``_put()``, and a cancellation cancels them and waits for them, as
+        in :meth:`_copy`. Unlike fsspec's ``_put()``, which creates a
+        directory for each local directory, it makes no other request: S3
+        has no directories, and creating one checks or creates the bucket.
+
+        Args:
+            lpath: Local source path, glob pattern, or list of them.
+            rpath: S3 destination path, or list of paths when ``lpath`` is a
+                list.
+            recursive: Whether to copy the directories with their contents.
+            callback: Progress callback.
+            batch_size: Number of uploads to run at the same time.
+            maxdepth: Maximum depth of a recursive copy.
+            **kwargs: Additional parameters passed to ``_put_file()``.
+        """
+        pairs = await self._thread_pool.run(
+            self._sync_fs._put_pairs, lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs
+        )
+        files = [(lp, rp) for lp, rp in pairs if not os.path.isdir(lp)]
+        await _S3BatchTasks().run(
+            super()._put(
+                [lp for lp, _ in files],
+                [rp for _, rp in files],
+                callback=callback,
+                batch_size=batch_size,
+                **kwargs,
+            )
+        )
+
+    async def _pipe(self, path, value=None, batch_size=None, **kwargs) -> None:
+        """Write bytes into one or more paths.
+
+        See :meth:`S3FileSystem.pipe`. The writes run as in fsspec's
+        ``_pipe()``, and a cancellation cancels them and waits for them, as
+        in :meth:`_copy`.
+
+        Args:
+            path: S3 path, or a dictionary of paths and their bytes.
+            value: The bytes to write when ``path`` is a path.
+            batch_size: Number of writes to run at the same time.
+            **kwargs: Additional parameters passed to ``_pipe_file()``.
+        """
+        await _S3BatchTasks().run(super()._pipe(path, value, batch_size=batch_size, **kwargs))
+
+    async def _cat(self, path, recursive=False, on_error="raise", batch_size=None, **kwargs):
+        """Read the contents of one or more paths.
+
+        See :meth:`S3FileSystem.cat`. The reads run as in fsspec's
+        ``_cat()``, and a cancellation cancels them and waits for them, as
+        in :meth:`_copy`.
+
+        Args:
+            path: S3 path, glob pattern, or list of them.
+            recursive: Whether to read the files below the directories.
+            on_error: ``"raise"``, ``"omit"``, or ``"return"``.
+            batch_size: Number of reads to run at the same time.
+            **kwargs: Additional parameters passed to ``_cat_file()``.
+
+        Returns:
+            The bytes of a single path, or a dictionary of the paths and
+            their bytes.
+        """
+        return await _S3BatchTasks().run(
+            super()._cat(
+                path, recursive=recursive, on_error=on_error, batch_size=batch_size, **kwargs
+            )
+        )
+
+    async def _cat_ranges(
+        self,
+        paths,
+        starts,
+        ends,
+        max_gap=None,
+        batch_size=None,
+        on_error="return",
+        **kwargs,
+    ):
+        """Read byte ranges of one or more paths.
+
+        See :meth:`S3FileSystem.cat_ranges`. The reads run as in fsspec's
+        ``_cat_ranges()``, and a cancellation cancels them and waits for
+        them, as in :meth:`_copy`.
+
+        Args:
+            paths: List of S3 paths.
+            starts: Start offset of each range, or one for all of them.
+            ends: End offset of each range, or one for all of them.
+            max_gap: Not supported; must be None.
+            batch_size: Number of reads to run at the same time.
+            on_error: ``"return"`` to return the error of a read in place of
+                its bytes; otherwise the first error is raised.
+            **kwargs: Additional parameters passed to ``_cat_file()``.
+
+        Returns:
+            The bytes or the error of each range, in order.
+        """
+        return await _S3BatchTasks().run(
+            super()._cat_ranges(
+                paths,
+                starts,
+                ends,
+                max_gap=max_gap,
+                batch_size=batch_size,
+                on_error=on_error,
+                **kwargs,
+            )
+        )
+
+    async def _cp_file(
+        self, path1: str, path2: str, recursive=False, maxdepth=None, on_error=None, **kwargs
+    ) -> None:
         """Copy an S3 object, using async parallel multipart upload for large files.
 
         A directory ``path1``, which recursive ``copy()`` passes along with
@@ -583,11 +795,16 @@ async def _cp_file(self, path1: str, path2: str, **kwargs) -> None:
         Args:
             path1: Source S3 path (s3://bucket/key).
             path2: Destination S3 path (s3://bucket/key).
-            **kwargs: Additional S3 copy parameters.
+            recursive: Unused parameter for fsspec compatibility.
+            maxdepth: Unused parameter for fsspec compatibility.
+            on_error: Unused parameter for fsspec compatibility.
+            **kwargs: Additional S3 copy parameters, as for
+                :meth:`S3FileSystem.cp_file`.
 
         Raises:
             ValueError: If trying to copy to a versioned file or copy buckets.
         """
+        _S3BatchTasks.track()
         await self._copy_file(path1, path2, **kwargs)
 
     async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool:
@@ -795,16 +1012,15 @@ async def _find(
     async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) -> list[str]:
         """Expand glob patterns and directories into the paths they match.
 
-        See :meth:`S3FileSystem.expand_path`. The other paths are expanded by
-        fsspec's ``_expand_path()``, whose glob lists only the keys under the
-        literal part of the pattern.
+        Runs :meth:`S3FileSystem.expand_path`, as :meth:`_find` runs
+        ``find()``.
 
         Args:
             path: S3 path, glob pattern, or list of them.
             recursive: Whether to include the paths below the directories.
             maxdepth: Maximum depth of the expansion, at least 1.
-            **kwargs: Additional arguments passed to fsspec's
-                ``_expand_path()``, such as ``assume_literal``.
+            **kwargs: Additional arguments passed to ``glob`` and ``find``,
+                and ``assume_literal``.
 
         Returns:
             The sorted matching paths.
@@ -813,23 +1029,9 @@ async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs) ->
             ValueError: If ``maxdepth`` is less than 1.
             FileNotFoundError: If nothing matches.
         """
-        if maxdepth is not None and maxdepth < 1:
-            raise ValueError("maxdepth must be at least 1")
-        versions, others = self._sync_fs._split_version_paths(path)
-        out = {p for p in versions if not recursive or await self._exists(p)}
-        if others:
-            try:
-                out.update(
-                    await super()._expand_path(
-                        others, recursive=recursive, maxdepth=maxdepth, **kwargs
-                    )
-                )
-            except FileNotFoundError:
-                if not out:
-                    raise
-        if not out:
-            raise FileNotFoundError(path)
-        return sorted(out)
+        return await self._thread_pool.run(
+            self._sync_fs.expand_path, path, recursive=recursive, maxdepth=maxdepth, **kwargs
+        )
 
     def _open(
         self,
diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py
index 832d5ffa0..95e8760f1 100644
--- a/pyathena/filesystem/s3_executor.py
+++ b/pyathena/filesystem/s3_executor.py
@@ -498,10 +498,14 @@ def shutdown(self, wait: bool = True) -> None:
     async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
         """Run a blocking call in the executor and wait for it.
 
-        Cancelling the waiting task cancels a call that has not started, and
-        stops the retries of the AWS API calls that a running call makes, as
-        :func:`~pyathena.aio.util.run_in_thread` does; the call itself keeps
-        running until it returns or raises.
+        Cancelling the waiting task cancels a call that has not started. A
+        running call makes no further attempt of its AWS API calls, as with
+        :func:`~pyathena.aio.util.run_in_thread`, and a write in it sends no
+        further part and does not commit (see
+        :func:`~pyathena.util.raise_if_stopped`); the task waits for the call
+        to end before it re-raises the cancellation, also when the event loop
+        shuts down with ``asyncio.run()``. A repeated cancellation stops that
+        wait, and the call then ends in the executor.
 
         Args:
             fn: The blocking callable.
@@ -531,12 +535,24 @@ def call() -> T:
                 with lock:
                     returned = True
 
+        # Created before it is scheduled, so that a cancellation while it is
+        # being scheduled can still cancel it or wait for it.
+        future: Future[T] = Future()
+        waiter = self.waiter(future)
         try:
-            return await asyncio.wrap_future(self.submit(call))
+            self.submit_to(future, call)
+            # Unlike awaiting the waiter, a cancellation of this task leaves
+            # the waiter to the wait below.
+            await asyncio.wait([waiter])
+            return waiter.result()
         except asyncio.CancelledError:
             with lock:
                 if not returned:
                     stop.set()
+            future.cancel()
+            # The states are also checked every second, as in map_async().
+            while not future.done():
+                await asyncio.wait([waiter], timeout=1.0)
             raise
 
     async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py
index 06a291ee1..aaaac7a41 100644
--- a/tests/pyathena/filesystem/test_s3.py
+++ b/tests/pyathena/filesystem/test_s3.py
@@ -1056,6 +1056,17 @@ def test_glob_matches_fsspec(self, pattern, maxdepth):
         expected = [p.lstrip("/") for p in memory.glob(f"/bucket/{pattern}", maxdepth=maxdepth)]
         assert sorted(fs.glob(f"s3://bucket/{pattern}", maxdepth=maxdepth)) == expected
 
+    def test_glob_lists_stem_prefix(self):
+        # GH-1337: glob() lists only the keys that start with the stem before
+        # the first wildcard, as fsspec's asynchronous glob does, which
+        # AioS3FileSystem used before it expanded its paths with this
+        # filesystem. It listed every key of the directory.
+        fs = self._make_fs()
+        self._serve_keys(fs, self.FIND_KEYS)
+
+        assert fs.glob("s3://bucket/dir/s*") == ["bucket/dir/sub"]
+        assert [c.kwargs["Prefix"] for c in fs._call.call_args_list] == ["dir/s"]
+
     def test_find_withdirs_omits_bucket(self):
         fs = self._make_fs()
         self._serve_keys(fs, self.FIND_KEYS)
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 95cfaca98..e54fac525 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -11,7 +11,7 @@
 import uuid
 import warnings
 import weakref
-from concurrent.futures import Future, ThreadPoolExecutor, wait
+from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait
 from datetime import UTC, datetime
 from types import SimpleNamespace
 from unittest import mock
@@ -34,7 +34,7 @@
 )
 from pyathena.filesystem.s3_path import S3Path
 from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession
-from pyathena.util import RetryConfig
+from pyathena.util import RetryConfig, raise_if_stopped
 from tests import ENV
 from tests.pyathena.conftest import connect
 from tests.pyathena.util import (
@@ -1780,6 +1780,249 @@ async def test_rm_maxdepth(self):
             await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10)
             stubber.assert_no_pending_responses()
 
+    @pytest.mark.asyncio
+    async def test_put_file_cancelled_stops_upload(self, tmp_path):
+        # GH-1337: a cancelled _put_file() sends no further part, does not
+        # complete the upload, aborts it, and raises once the abort has ended.
+        # The upload used to run on in its thread and commit.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE
+        local = tmp_path / "local"
+        local.write_bytes(bytes(3 * block_size + 1))
+        started = threading.Event()
+        release = threading.Event()
+        requests = []
+
+        def call(method, **request):
+            name = method.__name__
+            requests.append(name)
+            if name == "create_multipart_upload":
+                return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}
+            if name == "upload_part":
+                started.set()
+                release.wait(EVENT_TIMEOUT)
+                return {"ETag": '"e"'}
+            return {}
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        task = asyncio.ensure_future(
+            fs._put_file(str(local), "s3://bucket/key", block_size=block_size, max_workers=1)
+        )
+        try:
+            assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT)
+            task.cancel()
+            done, _ = await asyncio.wait([task], timeout=0.1)
+            assert not done
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await await_bounded(task)
+        finally:
+            release.set()
+        assert requests == ["create_multipart_upload", "upload_part", "abort_multipart_upload"]
+
+    @pytest.mark.parametrize("shutdown", [False, True])
+    @pytest.mark.parametrize(
+        "method", ["_cat", "_cat_ranges", "_pipe", "_put", "_get", "_copy", "_mv"]
+    )
+    def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
+        # GH-1337: a cancelled batch coroutine cancels its per-object calls,
+        # which make no further request, and raises once they have ended.
+        # fsspec's batch coroutines left them running after the
+        # cancellation. When the shutdown of the event loop at the end of
+        # asyncio.run() cancels the batch and its calls together, a call is
+        # not cancelled again, which would stop its wait.
+        fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
+        sync_fs = fs._sync_fs
+        release = threading.Event()
+        lock = threading.Lock()
+        state = {"started": 0, "ended": 0, "stopped": 0}
+
+        def blocking(result):
+            def call(*args, **kwargs):
+                with lock:
+                    state["started"] += 1
+                release.wait(EVENT_TIMEOUT)
+                try:
+                    raise_if_stopped("continue")
+                except CancelledError:
+                    with lock:
+                        state["stopped"] += 1
+                with lock:
+                    state["ended"] += 1
+                return result
+
+            return call
+
+        sources = ["s3://bucket/a", "s3://bucket/b"]
+        destinations = ["s3://bucket/c", "s3://bucket/d"]
+        local = [str(tmp_path / "a"), str(tmp_path / "b")]
+        sync_fs.cat_file = blocking(b"")
+        sync_fs.pipe_file = blocking(None)
+        sync_fs.put_file = blocking(None)
+        sync_fs.get_file = blocking(None)
+        sync_fs.info = blocking({"name": "bucket/a", "type": "file", "size": 0})
+        sync_fs._move_pairs = mock.MagicMock(
+            return_value=list(zip(sources, destinations, strict=True))
+        )
+        batches = {
+            "_cat": lambda: fs._cat(sources),
+            "_cat_ranges": lambda: fs._cat_ranges(sources, 0, 1),
+            "_pipe": lambda: fs._pipe(dict.fromkeys(sources, b"a")),
+            "_put": lambda: fs._put(local, sources),
+            "_get": lambda: fs._get(sources, local),
+            "_copy": lambda: fs._copy(sources, destinations),
+            "_mv": lambda: fs._mv(sources, destinations),
+        }
+
+        async def main():
+            task = asyncio.ensure_future(batches[method]())
+            await wait_until(lambda: state["started"] == 2)
+            if shutdown:
+                # asyncio.run() cancels the batch and its calls once main()
+                # has returned.
+                threading.Timer(0.2, release.set).start()
+                return
+            task.cancel()
+            done, _ = await asyncio.wait([task], timeout=0.1)
+            assert not done
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await await_bounded(task)
+            assert state["ended"] == 2
+
+        try:
+            run_bounded(main)
+        finally:
+            release.set()
+        assert state == {"started": 2, "ended": 2, "stopped": 2}
+
+    @pytest.mark.parametrize("method", ["_pipe", "_put", "_get"])
+    def test_batch_failure_waits_for_calls(self, tmp_path, method):
+        # GH-1337: when a per-object call of _pipe(), _put(), or _get() fails,
+        # fsspec cancels the others, which make no further request, and the
+        # failure is raised once they have ended. They used to run on, and
+        # commit, after the failure was raised.
+        fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
+        sync_fs = fs._sync_fs
+        release = threading.Event()
+        lock = threading.Lock()
+        state = {"started": 0, "ended": 0, "stopped": 0}
+
+        def call(source, *args, **kwargs):
+            if source.endswith("/a"):
+                raise ValueError("failed")
+            with lock:
+                state["started"] += 1
+            release.wait(EVENT_TIMEOUT)
+            try:
+                raise_if_stopped("continue")
+            except CancelledError:
+                with lock:
+                    state["stopped"] += 1
+            with lock:
+                state["ended"] += 1
+
+        sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = call
+        sources = ["s3://bucket/b", "s3://bucket/a"]
+        local = [str(tmp_path / "b"), str(tmp_path / "a")]
+        batches = {
+            "_pipe": lambda: fs._pipe(dict.fromkeys(sources, b"a")),
+            "_put": lambda: fs._put(local, sources),
+            "_get": lambda: fs._get(sources, local),
+        }
+
+        async def main():
+            task = asyncio.ensure_future(batches[method]())
+            await wait_until(lambda: state["started"] == 1)
+            done, _ = await asyncio.wait([task], timeout=0.1)
+            assert not done
+            release.set()
+            with pytest.raises(ValueError, match="failed"):
+                await await_bounded(task)
+            assert state == {"started": 1, "ended": 1, "stopped": 1}
+
+        try:
+            run_bounded(main)
+        finally:
+            release.set()
+
+    @pytest.mark.asyncio
+    async def test_put_recursive_creates_no_directories(self, tmp_path):
+        # GH-1337: a recursive _put() sends no request for the local
+        # directories. fsspec's _put() created a directory for each, which
+        # looked up the bucket, created it if it was missing and allowed, and
+        # otherwise raised PermissionError before any upload.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        store, requests = self._serve_keys(fs, set())
+        (tmp_path / "src" / "sub").mkdir(parents=True)
+        (tmp_path / "src" / "a").write_bytes(b"a")
+        (tmp_path / "src" / "sub" / "b").write_bytes(b"b")
+
+        await fs._put(f"{tmp_path}/src/", "s3://bucket/dst/", recursive=True)
+        assert store == {"dst/a", "dst/sub/b"}
+        assert [name for name, _ in requests] == ["put_object", "put_object"]
+
+    @pytest.mark.asyncio
+    async def test_expand_path_kwargs(self):
+        # GH-1337: the keyword arguments of _expand_path() reach the listing,
+        # as with S3FileSystem.expand_path(). fsspec's _expand_path() takes
+        # none, so they raised TypeError.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        self._serve_keys(fs, {"dir/a1", "dir/b1"})
+
+        assert await fs._expand_path("s3://bucket/dir", recursive=True, prefix="a") == [
+            "bucket/dir",
+            "bucket/dir/a1",
+        ]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("argument", ["batch_size", "max_workers"])
+    async def test_rm_batch_size(self, argument):
+        # GH-1337: batch_size of _rm() is the number of DeleteObjects requests
+        # sent at once, and max_workers its deprecated alias, as for
+        # S3FileSystem.rm(). batch_size was ignored, and max_workers was sent
+        # to DeleteObjects, which rejected it.
+        fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
+        overlap, state = overlapping_calls(2, timeout=0.2)
+        requests = []
+
+        def call(method, **request):
+            requests.append(request)
+            overlap()
+            return {}
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        with contextlib.ExitStack() as stack:
+            if argument == "max_workers":
+                stack.enter_context(pytest.warns(DeprecationWarning, match="use batch_size"))
+            await fs._rm(["s3://b1/a", "s3://b2/b"], **{argument: 1})
+        assert state["peak"] == 1
+        assert sorted(r["Bucket"] for r in requests) == ["b1", "b2"]
+        assert all(set(r) == {"Bucket", "Delete"} for r in requests)
+
+    @pytest.mark.asyncio
+    async def test_cp_file_fsspec_arguments(self):
+        # GH-1337: the recursive, maxdepth, and on_error arguments of
+        # _cp_file(), unused as with S3FileSystem.cp_file(), are not sent to
+        # CopyObject, which rejected them.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        _, requests = self._serve_keys(fs, {"key"})
+
+        await fs._cp_file(
+            "s3://bucket/key", "s3://bucket/dst", recursive=True, maxdepth=1, on_error="ignore"
+        )
+        assert [r for name, r in requests if name == "copy_object"] == [
+            {"CopySource": {"Bucket": "bucket", "Key": "key"}, "Bucket": "bucket", "Key": "dst"}
+        ]
+
     @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")])
     def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode):
         # GH-969: in a transaction, the open() parameters of put_file() go to
@@ -1852,8 +2095,9 @@ async def test_copy_version(self):
 
     @pytest.mark.asyncio
     async def test_expand_path_glob_lists_stem_prefix(self):
-        # Unversioned globs keep fsspec's async expansion, which lists only
-        # the keys that start with the stem before the first wildcard.
+        # A glob lists only the keys that start with the stem before the
+        # first wildcard, as fsspec's async expansion, which _expand_path()
+        # used before GH-1337, does.
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
         sync_fs._call = sync_fs._core.call = mock.MagicMock(
diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py
index 6925f3767..558d84961 100644
--- a/tests/pyathena/filesystem/test_s3_executor.py
+++ b/tests/pyathena/filesystem/test_s3_executor.py
@@ -40,6 +40,7 @@
     overlapping_calls,
     poll_until,
     retried,
+    run_bounded,
     run_in_forked_child,
     shutdown_bounded,
     sigint_when,
@@ -155,10 +156,13 @@ async def main():
             task = asyncio.ensure_future(run(work))
             await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT)
             task.cancel()
-            with pytest.raises(asyncio.CancelledError):
-                await task
+            # Lets the cancellation reach the task, which waits for a running
+            # call unless it runs shielded.
+            await asyncio.sleep(0.1)
             operation.release.set()
             started = time.monotonic()
+            with pytest.raises(asyncio.CancelledError):
+                await task
             assert await asyncio.to_thread(finished.wait, EVENT_TIMEOUT)
             return time.monotonic() - started
 
@@ -170,6 +174,42 @@ async def main():
 
         assert operation.calls == (3 if shielded else 1)
 
+    @pytest.mark.parametrize("cancellations", [1, 2])
+    def test_run_cancelled_waits_for_call(self, cancellations):
+        # GH-1337: a cancelled task raises once its running call has ended,
+        # so that nothing that the call does, such as the commit of a write,
+        # follows the cancellation. The task used to raise at once. A second
+        # cancellation stops the wait.
+        release = threading.Event()
+        started = threading.Event()
+        ended = threading.Event()
+        pool = S3ThreadPool(1)
+
+        def work():
+            started.set()
+            release.wait(EVENT_TIMEOUT)
+            ended.set()
+
+        async def main():
+            task = asyncio.ensure_future(pool.run(work))
+            assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT)
+            for _ in range(cancellations):
+                task.cancel()
+                done, _ = await asyncio.wait([task], timeout=0.1)
+            # Done before the call has ended only after the second one.
+            assert bool(done) == (cancellations == 2)
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await await_bounded(task)
+            if cancellations == 1:
+                assert ended.is_set()
+
+        try:
+            run_bounded(main)
+        finally:
+            release.set()
+            pool.shutdown()
+
     def test_run_cancelled_after_return(self):
         # GH-1203: a cancellation that arrives after the call has returned does
         # not stop the work that the call left running, such as the parts of a

From df64393031089f3bb95cee1f5758d24eed52632d Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 10:35:25 +0900
Subject: [PATCH 2/7] Hold the batch test calls until their stop instead of a
 timer (GH-1337)

The calls of the cancelled and failed batch tests were released by a
timer or after a fixed wait, which a slow event loop could reach before
the stop. They are now held until their stop scope stops them, and a
short time longer, so that a batch that does not wait for them returns
first.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/pyathena/filesystem/test_s3_async.py | 97 ++++++++++------------
 1 file changed, 42 insertions(+), 55 deletions(-)

diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index e54fac525..554ba0b13 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -49,6 +49,7 @@
     interrupt_each_opcode,
     mock_connection,
     overlapping_calls,
+    poll_until,
     retried,
     run_bounded,
     run_in_forked_child,
@@ -1822,6 +1823,34 @@ def call(method, **request):
             release.set()
         assert requests == ["create_multipart_upload", "upload_part", "abort_multipart_upload"]
 
+    @staticmethod
+    def _held_until_stopped(state, result=None):
+        # Return a blocking call that records in state that it started, is
+        # held until a stop scope of its thread stops, then 0.2 s longer, so
+        # that a caller that does not wait for it returns first, and records
+        # that it was stopped and ended. A call that is never stopped ends
+        # after EVENT_TIMEOUT without being recorded as stopped.
+        lock = threading.Lock()
+
+        def call(*args, **kwargs):
+            with lock:
+                state["started"] += 1
+            deadline = time.monotonic() + EVENT_TIMEOUT
+            while time.monotonic() < deadline:
+                try:
+                    raise_if_stopped("continue")
+                except CancelledError:
+                    with lock:
+                        state["stopped"] += 1
+                    time.sleep(0.2)
+                    break
+                time.sleep(0.01)
+            with lock:
+                state["ended"] += 1
+            return result
+
+        return call
+
     @pytest.mark.parametrize("shutdown", [False, True])
     @pytest.mark.parametrize(
         "method", ["_cat", "_cat_ranges", "_pipe", "_put", "_get", "_copy", "_mv"]
@@ -1835,34 +1864,15 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
         # not cancelled again, which would stop its wait.
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
-        release = threading.Event()
-        lock = threading.Lock()
         state = {"started": 0, "ended": 0, "stopped": 0}
-
-        def blocking(result):
-            def call(*args, **kwargs):
-                with lock:
-                    state["started"] += 1
-                release.wait(EVENT_TIMEOUT)
-                try:
-                    raise_if_stopped("continue")
-                except CancelledError:
-                    with lock:
-                        state["stopped"] += 1
-                with lock:
-                    state["ended"] += 1
-                return result
-
-            return call
-
         sources = ["s3://bucket/a", "s3://bucket/b"]
         destinations = ["s3://bucket/c", "s3://bucket/d"]
         local = [str(tmp_path / "a"), str(tmp_path / "b")]
-        sync_fs.cat_file = blocking(b"")
-        sync_fs.pipe_file = blocking(None)
-        sync_fs.put_file = blocking(None)
-        sync_fs.get_file = blocking(None)
-        sync_fs.info = blocking({"name": "bucket/a", "type": "file", "size": 0})
+        sync_fs.cat_file = self._held_until_stopped(state, b"")
+        sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = self._held_until_stopped(state)
+        sync_fs.info = self._held_until_stopped(
+            state, {"name": "bucket/a", "type": "file", "size": 0}
+        )
         sync_fs._move_pairs = mock.MagicMock(
             return_value=list(zip(sources, destinations, strict=True))
         )
@@ -1881,21 +1891,14 @@ async def main():
             await wait_until(lambda: state["started"] == 2)
             if shutdown:
                 # asyncio.run() cancels the batch and its calls once main()
-                # has returned.
-                threading.Timer(0.2, release.set).start()
+                # has returned, and returns once they have ended.
                 return
             task.cancel()
-            done, _ = await asyncio.wait([task], timeout=0.1)
-            assert not done
-            release.set()
             with pytest.raises(asyncio.CancelledError):
                 await await_bounded(task)
             assert state["ended"] == 2
 
-        try:
-            run_bounded(main)
-        finally:
-            release.set()
+        run_bounded(main)
         assert state == {"started": 2, "ended": 2, "stopped": 2}
 
     @pytest.mark.parametrize("method", ["_pipe", "_put", "_get"])
@@ -1906,23 +1909,15 @@ def test_batch_failure_waits_for_calls(self, tmp_path, method):
         # commit, after the failure was raised.
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
-        release = threading.Event()
-        lock = threading.Lock()
         state = {"started": 0, "ended": 0, "stopped": 0}
+        held = self._held_until_stopped(state)
 
         def call(source, *args, **kwargs):
             if source.endswith("/a"):
+                # Once the other call has started.
+                poll_until(lambda: state["started"] == 1, "the other call did not start")
                 raise ValueError("failed")
-            with lock:
-                state["started"] += 1
-            release.wait(EVENT_TIMEOUT)
-            try:
-                raise_if_stopped("continue")
-            except CancelledError:
-                with lock:
-                    state["stopped"] += 1
-            with lock:
-                state["ended"] += 1
+            return held(source, *args, **kwargs)
 
         sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = call
         sources = ["s3://bucket/b", "s3://bucket/a"]
@@ -1934,19 +1929,11 @@ def call(source, *args, **kwargs):
         }
 
         async def main():
-            task = asyncio.ensure_future(batches[method]())
-            await wait_until(lambda: state["started"] == 1)
-            done, _ = await asyncio.wait([task], timeout=0.1)
-            assert not done
-            release.set()
             with pytest.raises(ValueError, match="failed"):
-                await await_bounded(task)
+                await await_bounded(batches[method]())
             assert state == {"started": 1, "ended": 1, "stopped": 1}
 
-        try:
-            run_bounded(main)
-        finally:
-            release.set()
+        run_bounded(main)
 
     @pytest.mark.asyncio
     async def test_put_recursive_creates_no_directories(self, tmp_path):

From 511ffe40a7afc3f782a7e2bb6c4f321c26ca3d6d Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 10:55:47 +0900
Subject: [PATCH 3/7] Address the independent review of the aio alignment
 (GH-1337)

- A _put_file() or _pipe_file() in a transaction that is cancelled is
  discarded and dropped from the transaction once its write has ended.
  A write cancelled while its last part ran had nothing left to stop, and
  the transaction committed it.
- _put() and _pipe() return the values of their per-object calls, as
  fsspec's do.
- The pairing of put() strips the protocol of a single path-like
  destination as fsspec does, instead of iterating it.
- The batch tests hold their stopped calls until the test releases them,
  and check that the batch is still waiting, instead of relying on a
  fixed delay.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/aio.md                                |   3 +-
 pyathena/filesystem/s3.py                  |   7 +-
 pyathena/filesystem/s3_async.py            |  84 ++++++++++++--
 tests/pyathena/filesystem/test_s3_async.py | 126 ++++++++++++++++++---
 4 files changed, 187 insertions(+), 33 deletions(-)

diff --git a/docs/aio.md b/docs/aio.md
index 85d45673d..ca8f3a90c 100644
--- a/docs/aio.md
+++ b/docs/aio.md
@@ -332,7 +332,8 @@ its call makes in these threads, including the part uploads and range reads that
 call runs in parallel: no further attempt is made, and the wait before the next attempt
 ends. A write, such as one of `_put_file()` or `_pipe_file()`, sends no further part and
 aborts its multipart upload instead of completing it, and `_get_file()` reads and writes
-no further block. A request that has already been sent still runs to its end in its
+no further block. In a transaction, a cancelled `_put_file()` or `_pipe_file()` is
+discarded and left out of the commit. A request that has already been sent still runs to its end in its
 thread. The requests that run to their end after a cancellation keep their retries: the
 commit of a file, the CompleteMultipartUpload request of a multipart copy, and the
 aborts of multipart uploads.
diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py
index a142e0b84..74d0ebbe5 100644
--- a/pyathena/filesystem/s3.py
+++ b/pyathena/filesystem/s3.py
@@ -2767,11 +2767,8 @@ def _put_pairs(
             if not lpaths:
                 return []
         dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or self.isdir(rpath))
-        rpath = (
-            self._strip_protocol(rpath)
-            if isinstance(rpath, str)
-            else [self._strip_protocol(p) for p in rpath]
-        )
+        # A list of destinations or one path, also path-like, as in fsspec.
+        rpath = self._strip_protocol(rpath)
         exists = source_is_str and (
             (has_magic(lpath) and len(lpaths) == 1)
             or (not has_magic(lpath) and dest_is_dir and not trailing_sep(lpath))
diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index d27b5be9d..69e3e8b52 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -284,12 +284,56 @@ async def _pipe_file(
         if self._intrans:
             # The transaction belongs to this filesystem, not to the internal
             # S3FileSystem, so write through open() to defer the commit to it.
-            await self._thread_pool.run(self._pipe_file_in_transaction, path, value, mode, **kwargs)
+            await self._write_in_transaction(
+                functools.partial(self._pipe_file_in_transaction, path, value, mode, **kwargs)
+            )
             return
         await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
 
+    async def _write_in_transaction(self, write: Callable[[list[S3File]], None]) -> None:
+        """Run a write of a file of this filesystem's transaction in a thread of the filesystem.
+
+        A cancellation discards the file once the write has ended, and drops
+        it from the transaction, which would otherwise commit it: a write
+        whose last part was running when it was cancelled has nothing left
+        that its stop would skip.
+
+        Args:
+            write: The write, which appends the file that it opens to the
+                list that it receives.
+        """
+        files: list[S3File] = []
+        try:
+            await self._thread_pool.run(write, files)
+        except asyncio.CancelledError:
+            # A repeated cancellation does not stop the discard.
+            await self._thread_pool.run_shielded(self._discard_transaction_files, files)
+            raise
+
+    def _discard_transaction_files(self, files: list[S3File]) -> None:
+        """Drop files from this filesystem's transaction and discard them.
+
+        A failed discard is logged instead of raised, so that the
+        cancellation that discarded the files is raised.
+
+        Args:
+            files: The files.
+        """
+        for file in files:
+            with contextlib.suppress(ValueError):
+                self.transaction.files.remove(file)
+            try:
+                file.discard()
+            except Exception:
+                _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")
+
     def _pipe_file_in_transaction(
-        self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs
+        self,
+        path: str,
+        value: bytes | bytearray | memoryview,
+        mode: str,
+        files: list[S3File],
+        **kwargs,
     ) -> None:
         """Write bytes into the path as a file of this filesystem's transaction.
 
@@ -300,6 +344,7 @@ 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.
+            files: The list to which the opened file is appended.
             **kwargs: Additional parameters passed to ``open()``, except
                 ``compression``, with which the value is compressed before
                 it is written, as in :meth:`S3FileSystem.pipe_file`.
@@ -317,7 +362,9 @@ def _pipe_file_in_transaction(
         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.core.check_multipart_upload_size(memoryview(value).nbytes, block_size)
-        self.open(path, "xb" if mode == "create" else "wb", **kwargs)._write_and_close(value)
+        file = self.open(path, "xb" if mode == "create" else "wb", **kwargs)
+        files.append(file)
+        file._write_and_close(value)
 
     async def _put_file(
         self,
@@ -330,8 +377,10 @@ async def _put_file(
         _S3BatchTasks.track()
         if self._intrans:
             # See _pipe_file.
-            await self._thread_pool.run(
-                self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs
+            await self._write_in_transaction(
+                functools.partial(
+                    self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs
+                )
             )
             return
         await self._thread_pool.run(
@@ -339,7 +388,7 @@ async def _put_file(
         )
 
     def _put_file_in_transaction(
-        self, lpath: str, rpath: str, callback, mode: str, **kwargs
+        self, lpath: str, rpath: str, callback, mode: str, files: list[S3File], **kwargs
     ) -> None:
         """Upload a local file as a file of this filesystem's transaction.
 
@@ -354,6 +403,7 @@ def _put_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.
+            files: The list to which the opened file is appended.
             **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
                 The ``block_size``, ``max_workers``, and ``s3_additional_kwargs``
                 parameters of ``open()`` are also accepted.
@@ -383,13 +433,15 @@ def _put_file_in_transaction(
 
         # See S3FileSystem.put_file.
         with open(lpath, "rb") as local:
-            self.open(
+            file = self.open(
                 rpath,
                 "xb" if mode == "create" else "wb",
                 block_size=block_size,
                 max_workers=max_workers,
                 s3_additional_kwargs=s3_additional_kwargs,
-            )._write_file_and_close(local, callback)
+            )
+            files.append(file)
+            file._write_file_and_close(local, callback)
         self.invalidate_cache(rpath)
 
     async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
@@ -671,7 +723,7 @@ async def _put(
         batch_size=None,
         maxdepth=None,
         **kwargs,
-    ) -> None:
+    ) -> list[Any]:
         """Copy files from the local filesystem to S3.
 
         See :meth:`S3FileSystem.put`. The uploads run as in fsspec's
@@ -689,12 +741,15 @@ async def _put(
             batch_size: Number of uploads to run at the same time.
             maxdepth: Maximum depth of a recursive copy.
             **kwargs: Additional parameters passed to ``_put_file()``.
+
+        Returns:
+            The return values of the ``_put_file()`` calls, as in fsspec.
         """
         pairs = await self._thread_pool.run(
             self._sync_fs._put_pairs, lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs
         )
         files = [(lp, rp) for lp, rp in pairs if not os.path.isdir(lp)]
-        await _S3BatchTasks().run(
+        return await _S3BatchTasks().run(
             super()._put(
                 [lp for lp, _ in files],
                 [rp for _, rp in files],
@@ -704,7 +759,7 @@ async def _put(
             )
         )
 
-    async def _pipe(self, path, value=None, batch_size=None, **kwargs) -> None:
+    async def _pipe(self, path, value=None, batch_size=None, **kwargs) -> list[Any]:
         """Write bytes into one or more paths.
 
         See :meth:`S3FileSystem.pipe`. The writes run as in fsspec's
@@ -716,8 +771,13 @@ async def _pipe(self, path, value=None, batch_size=None, **kwargs) -> None:
             value: The bytes to write when ``path`` is a path.
             batch_size: Number of writes to run at the same time.
             **kwargs: Additional parameters passed to ``_pipe_file()``.
+
+        Returns:
+            The return values of the ``_pipe_file()`` calls, as in fsspec.
         """
-        await _S3BatchTasks().run(super()._pipe(path, value, batch_size=batch_size, **kwargs))
+        return await _S3BatchTasks().run(
+            super()._pipe(path, value, batch_size=batch_size, **kwargs)
+        )
 
     async def _cat(self, path, recursive=False, on_error="raise", batch_size=None, **kwargs):
         """Read the contents of one or more paths.
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 554ba0b13..669e262ea 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -13,6 +13,7 @@
 import weakref
 from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait
 from datetime import UTC, datetime
+from pathlib import PurePosixPath
 from types import SimpleNamespace
 from unittest import mock
 
@@ -1824,12 +1825,11 @@ def call(method, **request):
         assert requests == ["create_multipart_upload", "upload_part", "abort_multipart_upload"]
 
     @staticmethod
-    def _held_until_stopped(state, result=None):
-        # Return a blocking call that records in state that it started, is
-        # held until a stop scope of its thread stops, then 0.2 s longer, so
-        # that a caller that does not wait for it returns first, and records
-        # that it was stopped and ended. A call that is never stopped ends
-        # after EVENT_TIMEOUT without being recorded as stopped.
+    def _held_until_stopped(state, release, result=None):
+        # Return a blocking call that records in state that it started, waits
+        # until a stop scope of its thread stops, records that, and then
+        # waits for release. A call that is never stopped ends after
+        # EVENT_TIMEOUT without being recorded as stopped.
         lock = threading.Lock()
 
         def call(*args, **kwargs):
@@ -1842,7 +1842,7 @@ def call(*args, **kwargs):
                 except CancelledError:
                     with lock:
                         state["stopped"] += 1
-                    time.sleep(0.2)
+                    release.wait(EVENT_TIMEOUT)
                     break
                 time.sleep(0.01)
             with lock:
@@ -1865,13 +1865,15 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
         state = {"started": 0, "ended": 0, "stopped": 0}
+        release = threading.Event()
         sources = ["s3://bucket/a", "s3://bucket/b"]
         destinations = ["s3://bucket/c", "s3://bucket/d"]
         local = [str(tmp_path / "a"), str(tmp_path / "b")]
-        sync_fs.cat_file = self._held_until_stopped(state, b"")
-        sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = self._held_until_stopped(state)
+        held = self._held_until_stopped(state, release)
+        sync_fs.cat_file = self._held_until_stopped(state, release, b"")
+        sync_fs.pipe_file = sync_fs.put_file = sync_fs.get_file = held
         sync_fs.info = self._held_until_stopped(
-            state, {"name": "bucket/a", "type": "file", "size": 0}
+            state, release, {"name": "bucket/a", "type": "file", "size": 0}
         )
         sync_fs._move_pairs = mock.MagicMock(
             return_value=list(zip(sources, destinations, strict=True))
@@ -1885,21 +1887,40 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
             "_copy": lambda: fs._copy(sources, destinations),
             "_mv": lambda: fs._mv(sources, destinations),
         }
+        returned = threading.Event()
+        returned_while_held = []
+
+        def release_once_stopped():
+            # The shutdown: asyncio.run() must not return while the stopped
+            # calls are held.
+            poll_until(lambda: state["stopped"] == 2, "the calls were not stopped")
+            returned_while_held.append(returned.wait(0.1))
+            release.set()
 
         async def main():
             task = asyncio.ensure_future(batches[method]())
             await wait_until(lambda: state["started"] == 2)
             if shutdown:
                 # asyncio.run() cancels the batch and its calls once main()
-                # has returned, and returns once they have ended.
+                # has returned.
+                threading.Thread(target=release_once_stopped, daemon=True).start()
                 return
             task.cancel()
+            await wait_until(lambda: state["stopped"] == 2)
+            done, _ = await asyncio.wait([task], timeout=0.1)
+            assert not done
+            release.set()
             with pytest.raises(asyncio.CancelledError):
                 await await_bounded(task)
             assert state["ended"] == 2
 
-        run_bounded(main)
+        try:
+            run_bounded(main)
+            returned.set()
+        finally:
+            release.set()
         assert state == {"started": 2, "ended": 2, "stopped": 2}
+        assert returned_while_held == ([False] if shutdown else [])
 
     @pytest.mark.parametrize("method", ["_pipe", "_put", "_get"])
     def test_batch_failure_waits_for_calls(self, tmp_path, method):
@@ -1910,7 +1931,8 @@ def test_batch_failure_waits_for_calls(self, tmp_path, method):
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
         state = {"started": 0, "ended": 0, "stopped": 0}
-        held = self._held_until_stopped(state)
+        release = threading.Event()
+        held = self._held_until_stopped(state, release)
 
         def call(source, *args, **kwargs):
             if source.endswith("/a"):
@@ -1929,11 +1951,85 @@ def call(source, *args, **kwargs):
         }
 
         async def main():
+            task = asyncio.ensure_future(batches[method]())
+            await wait_until(lambda: state["stopped"] == 1)
+            done, _ = await asyncio.wait([task], timeout=0.1)
+            assert not done
+            release.set()
             with pytest.raises(ValueError, match="failed"):
-                await await_bounded(batches[method]())
+                await await_bounded(task)
             assert state == {"started": 1, "ended": 1, "stopped": 1}
 
-        run_bounded(main)
+        try:
+            run_bounded(main)
+        finally:
+            release.set()
+
+    @pytest.mark.asyncio
+    async def test_transaction_put_file_cancelled_is_discarded(self, tmp_path):
+        # GH-1337: a _put_file() in a transaction that is cancelled while its
+        # last part runs is discarded, and the transaction does not commit
+        # it. The part ended normally, so the write had nothing left to stop,
+        # and the file stayed in the transaction.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE
+        local = tmp_path / "local"
+        local.write_bytes(bytes(2 * block_size))
+        started = threading.Event()
+        release = threading.Event()
+        requests = []
+
+        def call(method, **request):
+            name = method.__name__
+            requests.append(name)
+            if name == "create_multipart_upload":
+                return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}
+            if name == "upload_part":
+                if request["PartNumber"] == 2:
+                    started.set()
+                    release.wait(EVENT_TIMEOUT)
+                return {"ETag": '"e"'}
+            return {}
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        try:
+            with fs.transaction:
+                task = asyncio.ensure_future(
+                    fs._put_file(
+                        str(local), "s3://bucket/key", block_size=block_size, max_workers=1
+                    )
+                )
+                assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT)
+                task.cancel()
+                await asyncio.sleep(0.1)
+                release.set()
+                with pytest.raises(asyncio.CancelledError):
+                    await await_bounded(task)
+        finally:
+            release.set()
+        assert requests == [
+            "create_multipart_upload",
+            "upload_part",
+            "upload_part",
+            "abort_multipart_upload",
+        ]
+
+    @pytest.mark.asyncio
+    async def test_put_pipe_return_values(self, tmp_path):
+        # GH-1337: _put() and _pipe() return the values of their per-object
+        # calls, as fsspec's do, and _put() accepts a path-like destination.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        store, _ = self._serve_keys(fs, set())
+        local = tmp_path / "local"
+        local.write_bytes(b"a")
+
+        assert await fs._pipe("s3://bucket/a", b"a") == [None]
+        assert await fs._put(str(local), PurePosixPath("bucket/b")) == [None]
+        assert store == {"a", "b"}
 
     @pytest.mark.asyncio
     async def test_put_recursive_creates_no_directories(self, tmp_path):

From 1bf31e1690c55d4e4d89ff3d0c2ae2868138d2f7 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 11:20:37 +0900
Subject: [PATCH 4/7] Discard a transaction write cancelled twice without a
 further wait (GH-1337)

The follow-up review found that a second cancellation could end the wait
for a transaction write before its file was registered, let the
transaction commit the file while its discard was queued, or start a new
wait behind the running write. _S3TransactionWrite now drops the file
from the transaction when the write is cancelled, or when it is opened
after that, and the file is discarded once by whichever of the write and
the cancelled coroutine ends last; a write that is still running
discards its own file, so the coroutine does not wait for it again.

The batch cancellation test records, in the event loop, how many calls
had ended when the batch task finished, instead of observing the return
of asyncio.run() from the test thread, and cancels the calls before the
batch in a new case, the order in which the shutdown of the event loop
can reach them.

The moves of the mv tests are compared by destination, since the copies
of S3FileSystem.mv() run in parallel since GH-1322.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py            | 158 ++++++++++++++++-----
 tests/pyathena/filesystem/test_s3.py       |  18 ++-
 tests/pyathena/filesystem/test_s3_async.py | 142 +++++++++++++++---
 3 files changed, 260 insertions(+), 58 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 69e3e8b52..412a511d8 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -16,6 +16,7 @@
 import logging
 import mimetypes
 import os
+import threading
 import warnings
 from collections.abc import Awaitable, Callable, Iterable
 from concurrent.futures import Future, wait
@@ -47,6 +48,8 @@
 if TYPE_CHECKING:
     from datetime import datetime
 
+    from fsspec.transaction import Transaction
+
     from pyathena.connection import Connection
 
 _logger = logging.getLogger(__name__)
@@ -114,6 +117,101 @@ async def run(self, batch: Awaitable[T]) -> T:
             self._current.reset(token)
 
 
+class _S3TransactionWrite:
+    """A write of a file of an ``AioS3FileSystem`` transaction, which a cancellation discards.
+
+    Once the write is cancelled, its file is dropped from the transaction,
+    also if it is opened only afterwards, so that the transaction cannot
+    commit it, and it is discarded once the write has ended: by the thread
+    of the write if the write was still running when it was cancelled, and
+    otherwise by the cancelled coroutine, after :meth:`cancel`.
+
+    Args:
+        transaction: The transaction of the filesystem.
+        write: The write, which calls the callable that it receives with
+            the file once it has opened it.
+    """
+
+    def __init__(
+        self, transaction: Transaction, write: Callable[[Callable[[S3File], None]], None]
+    ) -> None:
+        """Initialize the write.
+
+        Args:
+            transaction: The transaction of the filesystem.
+            write: The write, which calls the callable that it receives with
+                the file once it has opened it.
+        """
+        self._transaction = transaction
+        self._write = write
+        self._lock = threading.Lock()
+        self._files: list[S3File] = []
+        self._cancelled = False
+        self._ended = False
+
+    def run(self) -> None:
+        """Run the write, and discard its file if it was cancelled while running."""
+        try:
+            self._write(self._opened)
+        finally:
+            with self._lock:
+                self._ended = True
+                cancelled = self._cancelled
+            if cancelled:
+                self.discard()
+
+    def cancel(self) -> bool:
+        """Drop the file of the write from the transaction.
+
+        Returns:
+            Whether the write has ended, in which case the caller discards
+            the file with :meth:`discard`; otherwise the thread of the write
+            discards it when the write ends.
+        """
+        with self._lock:
+            self._cancelled = True
+            files = list(self._files)
+            ended = self._ended
+        for file in files:
+            self._detach(file)
+        return ended
+
+    def discard(self) -> None:
+        """Discard the file of the write.
+
+        A failed discard is logged instead of raised, so that the
+        cancellation is raised.
+        """
+        with self._lock:
+            files = list(self._files)
+        for file in files:
+            try:
+                file.discard()
+            except Exception:
+                _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")
+
+    def _opened(self, file: S3File) -> None:
+        """Record the opened file, and drop it from the transaction if the write was cancelled.
+
+        Args:
+            file: The file, which ``open()`` has added to the transaction.
+        """
+        with self._lock:
+            self._files.append(file)
+            cancelled = self._cancelled
+        if cancelled:
+            self._detach(file)
+
+    def _detach(self, file: S3File) -> None:
+        """Drop a file from the transaction, if it is there.
+
+        Args:
+            file: The file.
+        """
+        with contextlib.suppress(ValueError):
+            self._transaction.files.remove(file)
+
+
 class AioS3FileSystem(AsyncFileSystem):
     """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
 
@@ -290,49 +388,35 @@ async def _pipe_file(
             return
         await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
 
-    async def _write_in_transaction(self, write: Callable[[list[S3File]], None]) -> None:
+    async def _write_in_transaction(
+        self, write: Callable[[Callable[[S3File], None]], None]
+    ) -> None:
         """Run a write of a file of this filesystem's transaction in a thread of the filesystem.
 
-        A cancellation discards the file once the write has ended, and drops
-        it from the transaction, which would otherwise commit it: a write
-        whose last part was running when it was cancelled has nothing left
-        that its stop would skip.
+        A cancellation drops the file from the transaction, which would
+        otherwise commit it, and discards it once the write has ended: a
+        write whose last part was running when it was cancelled has nothing
+        left that its stop would skip. See :class:`_S3TransactionWrite`.
 
         Args:
-            write: The write, which appends the file that it opens to the
-                list that it receives.
+            write: The write, which calls the callable that it receives with
+                the file once it has opened it.
         """
-        files: list[S3File] = []
+        transaction_write = _S3TransactionWrite(self.transaction, write)
         try:
-            await self._thread_pool.run(write, files)
+            await self._thread_pool.run(transaction_write.run)
         except asyncio.CancelledError:
-            # A repeated cancellation does not stop the discard.
-            await self._thread_pool.run_shielded(self._discard_transaction_files, files)
+            if transaction_write.cancel():
+                # A repeated cancellation does not stop the discard.
+                await self._thread_pool.run_shielded(transaction_write.discard)
             raise
 
-    def _discard_transaction_files(self, files: list[S3File]) -> None:
-        """Drop files from this filesystem's transaction and discard them.
-
-        A failed discard is logged instead of raised, so that the
-        cancellation that discarded the files is raised.
-
-        Args:
-            files: The files.
-        """
-        for file in files:
-            with contextlib.suppress(ValueError):
-                self.transaction.files.remove(file)
-            try:
-                file.discard()
-            except Exception:
-                _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")
-
     def _pipe_file_in_transaction(
         self,
         path: str,
         value: bytes | bytearray | memoryview,
         mode: str,
-        files: list[S3File],
+        opened: Callable[[S3File], None],
         **kwargs,
     ) -> None:
         """Write bytes into the path as a file of this filesystem's transaction.
@@ -344,7 +428,7 @@ 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.
-            files: The list to which the opened file is appended.
+            opened: Called with the file once it is opened.
             **kwargs: Additional parameters passed to ``open()``, except
                 ``compression``, with which the value is compressed before
                 it is written, as in :meth:`S3FileSystem.pipe_file`.
@@ -363,7 +447,7 @@ def _pipe_file_in_transaction(
         # The size in bytes; the length of a memoryview counts its items.
         self.core.check_multipart_upload_size(memoryview(value).nbytes, block_size)
         file = self.open(path, "xb" if mode == "create" else "wb", **kwargs)
-        files.append(file)
+        opened(file)
         file._write_and_close(value)
 
     async def _put_file(
@@ -388,7 +472,13 @@ async def _put_file(
         )
 
     def _put_file_in_transaction(
-        self, lpath: str, rpath: str, callback, mode: str, files: list[S3File], **kwargs
+        self,
+        lpath: str,
+        rpath: str,
+        callback,
+        mode: str,
+        opened: Callable[[S3File], None],
+        **kwargs,
     ) -> None:
         """Upload a local file as a file of this filesystem's transaction.
 
@@ -403,7 +493,7 @@ def _put_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.
-            files: The list to which the opened file is appended.
+            opened: Called with the file once it is opened.
             **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
                 The ``block_size``, ``max_workers``, and ``s3_additional_kwargs``
                 parameters of ``open()`` are also accepted.
@@ -440,7 +530,7 @@ def _put_file_in_transaction(
                 max_workers=max_workers,
                 s3_additional_kwargs=s3_additional_kwargs,
             )
-            files.append(file)
+            opened(file)
             file._write_file_and_close(local, callback)
         self.invalidate_cache(rpath)
 
diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py
index aaaac7a41..4210067bb 100644
--- a/tests/pyathena/filesystem/test_s3.py
+++ b/tests/pyathena/filesystem/test_s3.py
@@ -1986,14 +1986,18 @@ def test_mv_versions_onto_their_key(self):
 
     @staticmethod
     def _moves(fs):
-        # The CopyObject requests, as (CopySource, Key), and the objects of
-        # the DeleteObjects requests that the filesystem sent.
+        # The CopyObject requests, as (CopySource, Key) sorted by Key, since
+        # the copies run in parallel, and the objects of the DeleteObjects
+        # requests that the filesystem sent.
         calls = fs._call.call_args_list
-        copies = [
-            (c.kwargs["CopySource"], c.kwargs["Key"])
-            for c in calls
-            if c.args[0] is fs._client.copy_object
-        ]
+        copies = sorted(
+            (
+                (c.kwargs["CopySource"], c.kwargs["Key"])
+                for c in calls
+                if c.args[0] is fs._client.copy_object
+            ),
+            key=lambda copy: copy[1],
+        )
         deletes = [
             c.kwargs["Delete"]["Objects"] for c in calls if c.args[0] is fs._client.delete_objects
         ]
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 669e262ea..4f09d3c03 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1851,17 +1851,25 @@ def call(*args, **kwargs):
 
         return call
 
-    @pytest.mark.parametrize("shutdown", [False, True])
     @pytest.mark.parametrize(
-        "method", ["_cat", "_cat_ranges", "_pipe", "_put", "_get", "_copy", "_mv"]
+        ("method", "cancellation"),
+        [
+            (method, cancellation)
+            for method in ["_cat", "_cat_ranges", "_pipe", "_put", "_get", "_copy", "_mv"]
+            for cancellation in ["batch", "calls_first", "shutdown"]
+            # asyncio.gather() cancels the copies of _mv() again.
+            if (method, cancellation) != ("_mv", "calls_first")
+        ],
     )
-    def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
+    def test_batch_cancelled_waits_for_calls(self, tmp_path, method, cancellation):
         # GH-1337: a cancelled batch coroutine cancels its per-object calls,
         # which make no further request, and raises once they have ended.
         # fsspec's batch coroutines left them running after the
-        # cancellation. When the shutdown of the event loop at the end of
-        # asyncio.run() cancels the batch and its calls together, a call is
-        # not cancelled again, which would stop its wait.
+        # cancellation. A call whose task is already being cancelled, as the
+        # shutdown of the event loop at the end of asyncio.run() cancels
+        # every task, possibly before the batch handles its own
+        # cancellation (calls_first), is not cancelled again, which would
+        # stop its wait.
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
         sync_fs = fs._sync_fs
         state = {"started": 0, "ended": 0, "stopped": 0}
@@ -1887,24 +1895,33 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, shutdown):
             "_copy": lambda: fs._copy(sources, destinations),
             "_mv": lambda: fs._mv(sources, destinations),
         }
-        returned = threading.Event()
-        returned_while_held = []
+        # The calls that had ended when the batch finished, recorded in the
+        # event loop.
+        ended_when_done = []
 
-        def release_once_stopped():
-            # The shutdown: asyncio.run() must not return while the stopped
-            # calls are held.
+        def release_once_stopped(task):
+            # The shutdown: the stopped calls are held until the batch has
+            # finished, which it must not before they end, or for 0.5 s.
             poll_until(lambda: state["stopped"] == 2, "the calls were not stopped")
-            returned_while_held.append(returned.wait(0.1))
+            deadline = time.monotonic() + 0.5
+            while not task.done() and time.monotonic() < deadline:
+                time.sleep(0.01)
             release.set()
 
         async def main():
             task = asyncio.ensure_future(batches[method]())
+            task.add_done_callback(lambda _: ended_when_done.append(state["ended"]))
             await wait_until(lambda: state["started"] == 2)
-            if shutdown:
+            if cancellation == "shutdown":
                 # asyncio.run() cancels the batch and its calls once main()
                 # has returned.
-                threading.Thread(target=release_once_stopped, daemon=True).start()
+                threading.Thread(target=release_once_stopped, args=(task,), daemon=True).start()
                 return
+            if cancellation == "calls_first":
+                current = asyncio.current_task()
+                for call in asyncio.all_tasks() - {current, task}:
+                    call.cancel()
+                await wait_until(lambda: state["stopped"] == 2)
             task.cancel()
             await wait_until(lambda: state["stopped"] == 2)
             done, _ = await asyncio.wait([task], timeout=0.1)
@@ -1912,15 +1929,13 @@ async def main():
             release.set()
             with pytest.raises(asyncio.CancelledError):
                 await await_bounded(task)
-            assert state["ended"] == 2
 
         try:
             run_bounded(main)
-            returned.set()
         finally:
             release.set()
         assert state == {"started": 2, "ended": 2, "stopped": 2}
-        assert returned_while_held == ([False] if shutdown else [])
+        assert ended_when_done == [2]
 
     @pytest.mark.parametrize("method", ["_pipe", "_put", "_get"])
     def test_batch_failure_waits_for_calls(self, tmp_path, method):
@@ -2016,6 +2031,99 @@ def call(method, **request):
             "abort_multipart_upload",
         ]
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("stage", ["before_open", "last_part", "discard_queued"])
+    async def test_transaction_put_file_cancelled_twice(self, tmp_path, stage):
+        # GH-1337: a second cancellation of a _put_file() in a transaction
+        # returns without a further wait, and the transaction does not commit
+        # the file, whether the write opens it only afterwards (before_open),
+        # still runs its last part (last_part), or has ended and its discard
+        # waits for a thread of the filesystem (discard_queued).
+        fs = AioS3FileSystem(
+            key="dummy",
+            secret="dummy",
+            region_name="us-east-1",
+            max_workers=1 if stage == "discard_queued" else None,
+            skip_instance_cache=True,
+        )
+        block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE
+        local = tmp_path / "local"
+        local.write_bytes(bytes(2 * block_size if stage == "last_part" else 1))
+        state = {"started": 0, "ended": 0, "stopped": 0}
+        release = threading.Event()
+        held = self._held_until_stopped(state, release)
+        requests = []
+
+        def call(method, **request):
+            requests.append(method.__name__)
+            if method.__name__ == "create_multipart_upload":
+                return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}
+            if method.__name__ == "upload_part" and request["PartNumber"] == 2:
+                held()
+            return {"ETag": '"e"'}
+
+        class HeldCallback(Callback):
+            def set_size(self, size):
+                if stage == "before_open":
+                    held()
+
+            def relative_update(self, inc=1):
+                if stage == "discard_queued":
+                    held()
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        blocker_started = threading.Event()
+        blocker_release = threading.Event()
+
+        def blocker(*args, **kwargs):
+            blocker_started.set()
+            blocker_release.wait(EVENT_TIMEOUT)
+            return {}
+
+        fs._sync_fs.info = blocker
+        try:
+            with fs.transaction:
+                task = asyncio.ensure_future(
+                    fs._put_file(
+                        str(local),
+                        "s3://bucket/key",
+                        callback=HeldCallback(),
+                        block_size=block_size,
+                        max_workers=1,
+                    )
+                )
+                await wait_until(lambda: state["started"] == 1)
+                if stage == "discard_queued":
+                    # Takes the only thread once the write has ended.
+                    info = asyncio.ensure_future(fs._info("s3://bucket/other"))
+                    await asyncio.sleep(0.1)
+                task.cancel()
+                await wait_until(lambda: state["stopped"] == 1)
+                if stage == "discard_queued":
+                    release.set()
+                    assert await asyncio.to_thread(blocker_started.wait, EVENT_TIMEOUT)
+                    # The file has been dropped and its discard queued.
+                    await wait_until(lambda: not fs.transaction.files)
+                task.cancel()
+                # While the write or the queued discard is still held.
+                with pytest.raises(asyncio.CancelledError):
+                    await await_bounded(task, timeout=2)
+                if stage != "discard_queued":
+                    # The write ends while the transaction is still open.
+                    release.set()
+                    await asyncio.to_thread(shutdown_bounded, fs._thread_pool)
+            blocker_release.set()
+            if stage == "discard_queued":
+                await info
+                await asyncio.to_thread(shutdown_bounded, fs._thread_pool)
+        finally:
+            release.set()
+            blocker_release.set()
+        assert "complete_multipart_upload" not in requests
+        assert "put_object" not in requests
+        if stage == "last_part":
+            assert requests[-1] == "abort_multipart_upload"
+
     @pytest.mark.asyncio
     async def test_put_pipe_return_values(self, tmp_path):
         # GH-1337: _put() and _pipe() return the values of their per-object

From 727a1dd612a34752ed12d6199122dc0807ad0c4e Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 11:43:14 +0900
Subject: [PATCH 5/7] Keep a cancelled transaction write out of the transaction
 (GH-1337)

The second follow-up review found that a write cancelled twice could
still join the transaction between fsspec's open() and the record of its
file, and that a second cancellation arriving just as the write ended
started a new wait for the discard.

_S3TransactionWrite now opens the file with _open(autocommit=False) and
adds it to the transaction itself, under the lock that the cancellation
takes, so a cancelled write never joins it. After a repeated
cancellation, the discard runs without being waited for.

The batch cancellation test releases the held calls only once the end
of the batch has been recorded in the event loop, and the transaction
test checks that the queued discard was still held when the task ended,
with a new case in which the write ends just before the second
cancellation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py            | 132 +++++++++++----------
 tests/pyathena/filesystem/test_s3_async.py |  67 +++++++----
 2 files changed, 113 insertions(+), 86 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 412a511d8..c8141a99b 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -48,8 +48,6 @@
 if TYPE_CHECKING:
     from datetime import datetime
 
-    from fsspec.transaction import Transaction
-
     from pyathena.connection import Connection
 
 _logger = logging.getLogger(__name__)
@@ -120,29 +118,30 @@ async def run(self, batch: Awaitable[T]) -> T:
 class _S3TransactionWrite:
     """A write of a file of an ``AioS3FileSystem`` transaction, which a cancellation discards.
 
-    Once the write is cancelled, its file is dropped from the transaction,
-    also if it is opened only afterwards, so that the transaction cannot
-    commit it, and it is discarded once the write has ended: by the thread
-    of the write if the write was still running when it was cancelled, and
-    otherwise by the cancelled coroutine, after :meth:`cancel`.
+    The write opens its file with :meth:`open`, which adds it to the
+    transaction unless the write has been cancelled. Once the write is
+    cancelled, its file is dropped from the transaction, so that the
+    transaction cannot commit it, and it is discarded once the write has
+    ended: by the thread of the write if the write was still running when
+    it was cancelled, and otherwise by the cancelled coroutine, after
+    :meth:`cancel`.
 
     Args:
-        transaction: The transaction of the filesystem.
-        write: The write, which calls the callable that it receives with
-            the file once it has opened it.
+        fs: The filesystem, which is in a transaction.
+        write: The write, which opens its file with the callable that it
+            receives, which takes the arguments of ``open()``.
     """
 
-    def __init__(
-        self, transaction: Transaction, write: Callable[[Callable[[S3File], None]], None]
-    ) -> None:
+    def __init__(self, fs: AioS3FileSystem, write: Callable[[Callable[..., S3File]], None]) -> None:
         """Initialize the write.
 
         Args:
-            transaction: The transaction of the filesystem.
-            write: The write, which calls the callable that it receives with
-                the file once it has opened it.
+            fs: The filesystem, which is in a transaction.
+            write: The write, which opens its file with the callable that it
+                receives, which takes the arguments of ``open()``.
         """
-        self._transaction = transaction
+        self._fs = fs
+        self._transaction = fs.transaction
         self._write = write
         self._lock = threading.Lock()
         self._files: list[S3File] = []
@@ -152,7 +151,7 @@ def __init__(
     def run(self) -> None:
         """Run the write, and discard its file if it was cancelled while running."""
         try:
-            self._write(self._opened)
+            self._write(self.open)
         finally:
             with self._lock:
                 self._ended = True
@@ -160,6 +159,30 @@ def run(self) -> None:
             if cancelled:
                 self.discard()
 
+    def open(self, path: str, mode: str, **kwargs) -> S3File:
+        """Open the file of the write, which joins the transaction unless the write was cancelled.
+
+        Args:
+            path: S3 path (s3://bucket/key) of the file.
+            mode: The binary write mode of the file.
+            **kwargs: The other arguments of ``open()``, except
+                ``compression``. With ``autocommit=True``, as with
+                ``open()``, the file commits when it is closed and does not
+                join the transaction.
+
+        Returns:
+            The file.
+        """
+        autocommit = kwargs.pop("autocommit", False)
+        file = self._fs._open(self._fs._strip_protocol(path), mode, autocommit=autocommit, **kwargs)
+        with self._lock:
+            self._files.append(file)
+            # Under the lock, so that a cancellation either finds the file
+            # in the transaction or keeps it out.
+            if not autocommit and not self._cancelled:
+                self._transaction.files.append(file)
+        return file
+
     def cancel(self) -> bool:
         """Drop the file of the write from the transaction.
 
@@ -170,11 +193,10 @@ def cancel(self) -> bool:
         """
         with self._lock:
             self._cancelled = True
-            files = list(self._files)
-            ended = self._ended
-        for file in files:
-            self._detach(file)
-        return ended
+            for file in self._files:
+                with contextlib.suppress(ValueError):
+                    self._transaction.files.remove(file)
+            return self._ended
 
     def discard(self) -> None:
         """Discard the file of the write.
@@ -190,27 +212,6 @@ def discard(self) -> None:
             except Exception:
                 _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")
 
-    def _opened(self, file: S3File) -> None:
-        """Record the opened file, and drop it from the transaction if the write was cancelled.
-
-        Args:
-            file: The file, which ``open()`` has added to the transaction.
-        """
-        with self._lock:
-            self._files.append(file)
-            cancelled = self._cancelled
-        if cancelled:
-            self._detach(file)
-
-    def _detach(self, file: S3File) -> None:
-        """Drop a file from the transaction, if it is there.
-
-        Args:
-            file: The file.
-        """
-        with contextlib.suppress(ValueError):
-            self._transaction.files.remove(file)
-
 
 class AioS3FileSystem(AsyncFileSystem):
     """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
@@ -388,27 +389,34 @@ async def _pipe_file(
             return
         await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
 
-    async def _write_in_transaction(
-        self, write: Callable[[Callable[[S3File], None]], None]
-    ) -> None:
+    async def _write_in_transaction(self, write: Callable[[Callable[..., S3File]], None]) -> None:
         """Run a write of a file of this filesystem's transaction in a thread of the filesystem.
 
-        A cancellation drops the file from the transaction, which would
+        A cancellation keeps the file out of the transaction, which would
         otherwise commit it, and discards it once the write has ended: a
         write whose last part was running when it was cancelled has nothing
         left that its stop would skip. See :class:`_S3TransactionWrite`.
+        After a repeated cancellation, the discard runs without being waited
+        for.
 
         Args:
-            write: The write, which calls the callable that it receives with
-                the file once it has opened it.
+            write: The write, which opens its file with the callable that it
+                receives, which takes the arguments of ``open()``.
         """
-        transaction_write = _S3TransactionWrite(self.transaction, write)
+        task = asyncio.current_task()
+        cancellations = task.cancelling() if task is not None else 0
+        transaction_write = _S3TransactionWrite(self, write)
         try:
             await self._thread_pool.run(transaction_write.run)
         except asyncio.CancelledError:
             if transaction_write.cancel():
-                # A repeated cancellation does not stop the discard.
-                await self._thread_pool.run_shielded(transaction_write.discard)
+                if task is not None and task.cancelling() > cancellations + 1:
+                    # A repeated cancellation, which stopped the wait for the
+                    # write, does not wait for the discard either.
+                    self._thread_pool.submit(transaction_write.discard)
+                else:
+                    # A repeated cancellation does not stop the discard.
+                    await self._thread_pool.run_shielded(transaction_write.discard)
             raise
 
     def _pipe_file_in_transaction(
@@ -416,7 +424,7 @@ def _pipe_file_in_transaction(
         path: str,
         value: bytes | bytearray | memoryview,
         mode: str,
-        opened: Callable[[S3File], None],
+        open_file: Callable[..., S3File],
         **kwargs,
     ) -> None:
         """Write bytes into the path as a file of this filesystem's transaction.
@@ -428,7 +436,7 @@ 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.
-            opened: Called with the file once it is opened.
+            open_file: Opens the file, with the arguments of ``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`.
@@ -446,9 +454,7 @@ def _pipe_file_in_transaction(
         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.core.check_multipart_upload_size(memoryview(value).nbytes, block_size)
-        file = self.open(path, "xb" if mode == "create" else "wb", **kwargs)
-        opened(file)
-        file._write_and_close(value)
+        open_file(path, "xb" if mode == "create" else "wb", **kwargs)._write_and_close(value)
 
     async def _put_file(
         self,
@@ -477,7 +483,7 @@ def _put_file_in_transaction(
         rpath: str,
         callback,
         mode: str,
-        opened: Callable[[S3File], None],
+        open_file: Callable[..., S3File],
         **kwargs,
     ) -> None:
         """Upload a local file as a file of this filesystem's transaction.
@@ -493,7 +499,7 @@ def _put_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.
-            opened: Called with the file once it is opened.
+            open_file: Opens the file, with the arguments of ``open()``.
             **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
                 The ``block_size``, ``max_workers``, and ``s3_additional_kwargs``
                 parameters of ``open()`` are also accepted.
@@ -523,15 +529,13 @@ def _put_file_in_transaction(
 
         # See S3FileSystem.put_file.
         with open(lpath, "rb") as local:
-            file = self.open(
+            open_file(
                 rpath,
                 "xb" if mode == "create" else "wb",
                 block_size=block_size,
                 max_workers=max_workers,
                 s3_additional_kwargs=s3_additional_kwargs,
-            )
-            opened(file)
-            file._write_file_and_close(local, callback)
+            )._write_file_and_close(local, callback)
         self.invalidate_cache(rpath)
 
     async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 4f09d3c03..6088b993b 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1667,7 +1667,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs):
         fs._sync_fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3
         fs._sync_fs.default_block_size = 4
         fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock()
-        fs.open = mock.MagicMock()
+        fs._open = mock.MagicMock()
         local = tmp_path / "local"
         local.write_bytes(b"a" * 13)
 
@@ -1676,7 +1676,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs):
                 fs.pipe_file("s3://bucket/k1", b"a" * 13, **kwargs)
             with pytest.raises(ValueError, match="in 3 parts with a block size of 4 bytes"):
                 fs.put_file(str(local), "s3://bucket/k2", **kwargs)
-        fs.open.assert_not_called()
+        fs._open.assert_not_called()
         fs._sync_fs._call.assert_not_called()
 
     @pytest.mark.asyncio
@@ -1899,12 +1899,13 @@ def test_batch_cancelled_waits_for_calls(self, tmp_path, method, cancellation):
         # event loop.
         ended_when_done = []
 
-        def release_once_stopped(task):
-            # The shutdown: the stopped calls are held until the batch has
-            # finished, which it must not before they end, or for 0.5 s.
+        def release_once_stopped():
+            # The shutdown: the stopped calls are held until the end of the
+            # batch has been recorded, which must not happen before they end,
+            # or for 0.5 s.
             poll_until(lambda: state["stopped"] == 2, "the calls were not stopped")
             deadline = time.monotonic() + 0.5
-            while not task.done() and time.monotonic() < deadline:
+            while not ended_when_done and time.monotonic() < deadline:
                 time.sleep(0.01)
             release.set()
 
@@ -1915,7 +1916,7 @@ async def main():
             if cancellation == "shutdown":
                 # asyncio.run() cancels the batch and its calls once main()
                 # has returned.
-                threading.Thread(target=release_once_stopped, args=(task,), daemon=True).start()
+                threading.Thread(target=release_once_stopped, daemon=True).start()
                 return
             if cancellation == "calls_first":
                 current = asyncio.current_task()
@@ -2032,18 +2033,21 @@ def call(method, **request):
         ]
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("stage", ["before_open", "last_part", "discard_queued"])
+    @pytest.mark.parametrize("stage", ["before_open", "last_part", "discard_queued", "write_ended"])
     async def test_transaction_put_file_cancelled_twice(self, tmp_path, stage):
         # GH-1337: a second cancellation of a _put_file() in a transaction
         # returns without a further wait, and the transaction does not commit
         # the file, whether the write opens it only afterwards (before_open),
-        # still runs its last part (last_part), or has ended and its discard
-        # waits for a thread of the filesystem (discard_queued).
+        # still runs its last part (last_part), has ended and its discard
+        # waits for a thread of the filesystem (discard_queued), or ends just
+        # before the second cancellation, which then interrupts the wait for
+        # the write (write_ended).
+        queued = stage in ("discard_queued", "write_ended")
         fs = AioS3FileSystem(
             key="dummy",
             secret="dummy",
             region_name="us-east-1",
-            max_workers=1 if stage == "discard_queued" else None,
+            max_workers=1 if queued else None,
             skip_instance_cache=True,
         )
         block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE
@@ -2068,16 +2072,18 @@ def set_size(self, size):
                     held()
 
             def relative_update(self, inc=1):
-                if stage == "discard_queued":
+                if queued:
                     held()
 
         fs._sync_fs._call = fs._sync_fs._core.call = call
         blocker_started = threading.Event()
         blocker_release = threading.Event()
+        blocker_ended = threading.Event()
 
         def blocker(*args, **kwargs):
             blocker_started.set()
             blocker_release.wait(EVENT_TIMEOUT)
+            blocker_ended.set()
             return {}
 
         fs._sync_fs.info = blocker
@@ -2093,7 +2099,7 @@ def blocker(*args, **kwargs):
                     )
                 )
                 await wait_until(lambda: state["started"] == 1)
-                if stage == "discard_queued":
+                if queued:
                     # Takes the only thread once the write has ended.
                     info = asyncio.ensure_future(fs._info("s3://bucket/other"))
                     await asyncio.sleep(0.1)
@@ -2103,17 +2109,32 @@ def blocker(*args, **kwargs):
                     release.set()
                     assert await asyncio.to_thread(blocker_started.wait, EVENT_TIMEOUT)
                     # The file has been dropped and its discard queued.
-                    await wait_until(lambda: not fs.transaction.files)
-                task.cancel()
+                    await wait_until(lambda: not fs.transaction.files, timeout=2)
+                    task.cancel()
+                elif stage == "write_ended":
+
+                    def end_write_then_cancel():
+                        # In one callback of the event loop, so that the task
+                        # is cancelled again before it sees the end of the
+                        # write, once the only thread has taken the next call.
+                        release.set()
+                        assert blocker_started.wait(EVENT_TIMEOUT)
+                        task.cancel()
+
+                    asyncio.get_running_loop().call_soon(end_write_then_cancel)
+                else:
+                    task.cancel()
                 # While the write or the queued discard is still held.
                 with pytest.raises(asyncio.CancelledError):
                     await await_bounded(task, timeout=2)
-                if stage != "discard_queued":
+                # The queued discard was not what ended the task.
+                assert not blocker_ended.is_set()
+                if not queued:
                     # The write ends while the transaction is still open.
                     release.set()
                     await asyncio.to_thread(shutdown_bounded, fs._thread_pool)
             blocker_release.set()
-            if stage == "discard_queued":
+            if queued:
                 await info
                 await asyncio.to_thread(shutdown_bounded, fs._thread_pool)
         finally:
@@ -2217,10 +2238,10 @@ async def test_cp_file_fsspec_arguments(self):
     @pytest.mark.parametrize(("mode", "open_mode"), [("overwrite", "wb"), ("create", "xb")])
     def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode):
         # GH-969: in a transaction, the open() parameters of put_file() go to
-        # open(), and the other parameters, also in s3_additional_kwargs, to
+        # the file, and the other parameters, also in s3_additional_kwargs, to
         # S3. GH-972: fsspec's mode argument selects the mode of the file.
         fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True)
-        fs.open = mock.MagicMock()
+        fs._open = mock.MagicMock()
         lpath = tmp_path / "data.csv"
         lpath.write_bytes(b"a")
 
@@ -2236,17 +2257,19 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode
             # Without them, the file gets the defaults of the filesystem.
             fs.put_file(str(lpath), "s3://bucket/key", mode=mode)
 
-        assert fs.open.call_args_list == [
+        assert fs._open.call_args_list == [
             mock.call(
-                "s3://bucket/key",
+                "bucket/key",
                 open_mode,
+                autocommit=False,
                 block_size=S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE,
                 max_workers=2,
                 s3_additional_kwargs={"StorageClass": "STANDARD_IA", "ContentType": "text/csv"},
             ),
             mock.call(
-                "s3://bucket/key",
+                "bucket/key",
                 open_mode,
+                autocommit=False,
                 block_size=fs._sync_fs.default_block_size,
                 max_workers=fs._sync_fs.max_workers,
                 s3_additional_kwargs={"ContentType": "text/csv"},

From 6894974af56a995593b9a392daa65f7beebe44e9 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 12:01:46 +0900
Subject: [PATCH 6/7] Discard a stopped transaction write in its own thread
 (GH-1337)

The third follow-up review found that the count of cancellations at the
start of the write misses a cancellation still pending then, and that
dropping a file from the transaction could fail while another write
appended to it, since the search compares files with their __eq__().

A write that ends stopped now drops and discards its file itself, so
that the wait of a single cancellation for the write includes the
discard, and the coroutine never waits for a discard: it drops the file
of a write that ended before it was stopped and discards it in a thread.
The removal from the transaction is retried when another write changed
it during the search. pyathena.util.retries_stopped() reports whether a
stop scope has stopped.

The transaction tests cover a cancellation after the write returned and
hold a write at the moment its file would join the transaction.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py            |  81 +++++++++------
 pyathena/util.py                           |  12 ++-
 tests/pyathena/filesystem/test_s3_async.py | 113 ++++++++++++++++-----
 3 files changed, 146 insertions(+), 60 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index c8141a99b..900782582 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -44,6 +44,7 @@
 from pyathena.filesystem.s3_path import S3Path
 from pyathena.filesystem.s3_path_pairing import S3PathPairing
 from pyathena.filesystem.s3_writer import _S3WriteSession, _S3WriteState
+from pyathena.util import retries_stopped
 
 if TYPE_CHECKING:
     from datetime import datetime
@@ -118,13 +119,16 @@ async def run(self, batch: Awaitable[T]) -> T:
 class _S3TransactionWrite:
     """A write of a file of an ``AioS3FileSystem`` transaction, which a cancellation discards.
 
-    The write opens its file with :meth:`open`, which adds it to the
-    transaction unless the write has been cancelled. Once the write is
-    cancelled, its file is dropped from the transaction, so that the
-    transaction cannot commit it, and it is discarded once the write has
-    ended: by the thread of the write if the write was still running when
-    it was cancelled, and otherwise by the cancelled coroutine, after
-    :meth:`cancel`.
+    The write runs with :meth:`run` in a thread of the filesystem, through
+    :meth:`S3Executor.run() <pyathena.filesystem.s3_executor.S3Executor.run>`,
+    whose cancellation stops the stop scope of the write, and opens its file
+    with :meth:`open`, which adds the file to the transaction unless the
+    write has been stopped. A write that ends stopped drops its file from
+    the transaction and discards it, so that the wait of the cancelled
+    coroutine for the write includes the discard. The coroutine calls
+    :meth:`cancel` once the cancellation has been raised, which drops the
+    file from the transaction also when the write had ended before it was
+    stopped, and then discards it.
 
     Args:
         fs: The filesystem, which is in a transaction.
@@ -149,18 +153,23 @@ def __init__(self, fs: AioS3FileSystem, write: Callable[[Callable[..., S3File]],
         self._ended = False
 
     def run(self) -> None:
-        """Run the write, and discard its file if it was cancelled while running."""
+        """Run the write, and drop and discard its file if it ends stopped."""
         try:
             self._write(self.open)
         finally:
             with self._lock:
                 self._ended = True
-                cancelled = self._cancelled
-            if cancelled:
+                stopped = self._cancelled or retries_stopped()
+                if stopped:
+                    self._cancelled = True
+                    # Also here, for a write that a scope other than the
+                    # cancellation of the coroutine has stopped.
+                    self._detach_all()
+            if stopped:
                 self.discard()
 
     def open(self, path: str, mode: str, **kwargs) -> S3File:
-        """Open the file of the write, which joins the transaction unless the write was cancelled.
+        """Open the file of the write, which joins the transaction unless the write is stopped.
 
         Args:
             path: S3 path (s3://bucket/key) of the file.
@@ -179,7 +188,7 @@ def open(self, path: str, mode: str, **kwargs) -> S3File:
             self._files.append(file)
             # Under the lock, so that a cancellation either finds the file
             # in the transaction or keeps it out.
-            if not autocommit and not self._cancelled:
+            if not autocommit and not self._cancelled and not retries_stopped():
                 self._transaction.files.append(file)
         return file
 
@@ -187,16 +196,15 @@ def cancel(self) -> bool:
         """Drop the file of the write from the transaction.
 
         Returns:
-            Whether the write has ended, in which case the caller discards
-            the file with :meth:`discard`; otherwise the thread of the write
-            discards it when the write ends.
+            Whether the write ended before it was stopped, in which case
+            the caller discards the file with :meth:`discard`; otherwise the
+            thread of the write discards it when the write ends, or has.
         """
         with self._lock:
+            ended_unstopped = self._ended and not self._cancelled
             self._cancelled = True
-            for file in self._files:
-                with contextlib.suppress(ValueError):
-                    self._transaction.files.remove(file)
-            return self._ended
+            self._detach_all()
+            return ended_unstopped
 
     def discard(self) -> None:
         """Discard the file of the write.
@@ -212,6 +220,20 @@ def discard(self) -> None:
             except Exception:
                 _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")
 
+    def _detach_all(self) -> None:
+        """Drop the files of the write from the transaction; called under the lock."""
+        for file in self._files:
+            while True:
+                try:
+                    self._transaction.files.remove(file)
+                except ValueError:
+                    break
+                except (IndexError, RuntimeError):
+                    # Another write changed the transaction during the search,
+                    # which compares the files with their __eq__().
+                    continue
+                break
+
 
 class AioS3FileSystem(AsyncFileSystem):
     """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
@@ -393,30 +415,23 @@ async def _write_in_transaction(self, write: Callable[[Callable[..., S3File]], N
         """Run a write of a file of this filesystem's transaction in a thread of the filesystem.
 
         A cancellation keeps the file out of the transaction, which would
-        otherwise commit it, and discards it once the write has ended: a
-        write whose last part was running when it was cancelled has nothing
-        left that its stop would skip. See :class:`_S3TransactionWrite`.
-        After a repeated cancellation, the discard runs without being waited
-        for.
+        otherwise commit it, and discards it: a write whose last part was
+        running when it was cancelled has nothing left that its stop would
+        skip. A write that the cancellation stops discards its file before
+        it ends, which the cancellation waits for; the file of a write that
+        ended before is discarded in a thread of the filesystem without
+        being waited for. See :class:`_S3TransactionWrite`.
 
         Args:
             write: The write, which opens its file with the callable that it
                 receives, which takes the arguments of ``open()``.
         """
-        task = asyncio.current_task()
-        cancellations = task.cancelling() if task is not None else 0
         transaction_write = _S3TransactionWrite(self, write)
         try:
             await self._thread_pool.run(transaction_write.run)
         except asyncio.CancelledError:
             if transaction_write.cancel():
-                if task is not None and task.cancelling() > cancellations + 1:
-                    # A repeated cancellation, which stopped the wait for the
-                    # write, does not wait for the discard either.
-                    self._thread_pool.submit(transaction_write.discard)
-                else:
-                    # A repeated cancellation does not stop the discard.
-                    await self._thread_pool.run_shielded(transaction_write.discard)
+                self._thread_pool.submit(transaction_write.discard)
             raise
 
     def _pipe_file_in_transaction(
diff --git a/pyathena/util.py b/pyathena/util.py
index b70e78e77..0097e41ea 100644
--- a/pyathena/util.py
+++ b/pyathena/util.py
@@ -328,6 +328,16 @@ def unstoppable_retries() -> Iterator[None]:
         _RETRY_STOP_EVENTS.reset(token)
 
 
+def retries_stopped() -> bool:
+    """Return whether an event of the enclosing stop scopes is set.
+
+    Returns:
+        Whether an event of the enclosing :func:`stop_retries_on` scopes,
+        apart from those that :func:`unstoppable_retries` clears, is set.
+    """
+    return any(event.is_set() for event in _RETRY_STOP_EVENTS.get())
+
+
 def raise_if_stopped(operation: str) -> None:
     """Raise if an event of the enclosing stop scopes is set.
 
@@ -346,7 +356,7 @@ def raise_if_stopped(operation: str) -> None:
         concurrent.futures.CancelledError: If an event of the enclosing
             :func:`stop_retries_on` scopes is set.
     """
-    if any(event.is_set() for event in _RETRY_STOP_EVENTS.get()):
+    if retries_stopped():
         raise CancelledError(f"Stopped before it could {operation}.")
 
 
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 6088b993b..00891c8f8 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1,5 +1,6 @@
 import array
 import asyncio
+import collections
 import contextlib
 import gc
 import gzip
@@ -2033,16 +2034,16 @@ def call(method, **request):
         ]
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("stage", ["before_open", "last_part", "discard_queued", "write_ended"])
-    async def test_transaction_put_file_cancelled_twice(self, tmp_path, stage):
-        # GH-1337: a second cancellation of a _put_file() in a transaction
-        # returns without a further wait, and the transaction does not commit
-        # the file, whether the write opens it only afterwards (before_open),
-        # still runs its last part (last_part), has ended and its discard
-        # waits for a thread of the filesystem (discard_queued), or ends just
-        # before the second cancellation, which then interrupts the wait for
-        # the write (write_ended).
-        queued = stage in ("discard_queued", "write_ended")
+    @pytest.mark.parametrize("stage", ["before_open", "last_part", "write_ended", "returned"])
+    async def test_transaction_put_file_cancelled_late(self, tmp_path, stage):
+        # GH-1337: the transaction does not commit the file of a _put_file()
+        # whose task is cancelled, and a cancellation adds no wait for a
+        # discard behind other calls of the filesystem. The cases: cancelled
+        # twice before the write opens its file (before_open), while it runs
+        # its last part (last_part), or just after the write has ended, the
+        # second time (write_ended); and cancelled once after the write has
+        # returned (returned), whose file the coroutine discards.
+        queued = stage in ("write_ended", "returned")
         fs = AioS3FileSystem(
             key="dummy",
             secret="dummy",
@@ -2052,7 +2053,8 @@ async def test_transaction_put_file_cancelled_twice(self, tmp_path, stage):
         )
         block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE
         local = tmp_path / "local"
-        local.write_bytes(bytes(2 * block_size if stage == "last_part" else 1))
+        multipart = stage in ("last_part", "returned")
+        local.write_bytes(bytes(2 * block_size if multipart else 1))
         state = {"started": 0, "ended": 0, "stopped": 0}
         release = threading.Event()
         held = self._held_until_stopped(state, release)
@@ -2062,7 +2064,11 @@ def call(method, **request):
             requests.append(method.__name__)
             if method.__name__ == "create_multipart_upload":
                 return {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"}
-            if method.__name__ == "upload_part" and request["PartNumber"] == 2:
+            if (
+                stage == "last_part"
+                and method.__name__ == "upload_part"
+                and request["PartNumber"] == 2
+            ):
                 held()
             return {"ETag": '"e"'}
 
@@ -2072,8 +2078,12 @@ def set_size(self, size):
                     held()
 
             def relative_update(self, inc=1):
-                if queued:
+                if stage == "write_ended":
                     held()
+                elif stage == "returned" and state["started"] == 0:
+                    # Held without a stop, until the test releases it.
+                    state["started"] = 1
+                    release.wait(EVENT_TIMEOUT)
 
         fs._sync_fs._call = fs._sync_fs._core.call = call
         blocker_started = threading.Event()
@@ -2103,20 +2113,15 @@ def blocker(*args, **kwargs):
                     # Takes the only thread once the write has ended.
                     info = asyncio.ensure_future(fs._info("s3://bucket/other"))
                     await asyncio.sleep(0.1)
-                task.cancel()
-                await wait_until(lambda: state["stopped"] == 1)
-                if stage == "discard_queued":
-                    release.set()
-                    assert await asyncio.to_thread(blocker_started.wait, EVENT_TIMEOUT)
-                    # The file has been dropped and its discard queued.
-                    await wait_until(lambda: not fs.transaction.files, timeout=2)
+                if stage != "returned":
                     task.cancel()
-                elif stage == "write_ended":
+                    await wait_until(lambda: state["stopped"] == 1)
+                if queued:
 
                     def end_write_then_cancel():
                         # In one callback of the event loop, so that the task
-                        # is cancelled again before it sees the end of the
-                        # write, once the only thread has taken the next call.
+                        # is cancelled before it sees the end of the write,
+                        # once the only thread has taken the next call.
                         release.set()
                         assert blocker_started.wait(EVENT_TIMEOUT)
                         task.cancel()
@@ -2124,10 +2129,9 @@ def end_write_then_cancel():
                     asyncio.get_running_loop().call_soon(end_write_then_cancel)
                 else:
                     task.cancel()
-                # While the write or the queued discard is still held.
+                # While the write is held, or the next call holds the thread.
                 with pytest.raises(asyncio.CancelledError):
                     await await_bounded(task, timeout=2)
-                # The queued discard was not what ended the task.
                 assert not blocker_ended.is_set()
                 if not queued:
                     # The write ends while the transaction is still open.
@@ -2142,9 +2146,66 @@ def end_write_then_cancel():
             blocker_release.set()
         assert "complete_multipart_upload" not in requests
         assert "put_object" not in requests
-        if stage == "last_part":
+        if multipart:
             assert requests[-1] == "abort_multipart_upload"
 
+    @pytest.mark.asyncio
+    async def test_transaction_write_cancelled_before_it_joins(self, tmp_path):
+        # GH-1337: a _put_file() in a transaction that is cancelled twice
+        # before it opens its file never adds the file to the transaction,
+        # whose exit at any moment therefore commits nothing. The file used
+        # to join the transaction in open(), and to be dropped from it only
+        # afterwards.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        local = tmp_path / "local"
+        local.write_bytes(b"a")
+        state = {"started": 0, "ended": 0, "stopped": 0}
+        release = threading.Event()
+        held = self._held_until_stopped(state, release)
+        requests = []
+        joined = threading.Event()
+        proceed = threading.Event()
+
+        class HeldFiles(collections.deque):
+            def append(self, file):
+                # Holds the write once its file has joined the transaction.
+                super().append(file)
+                joined.set()
+                proceed.wait(EVENT_TIMEOUT)
+
+        class HeldCallback(Callback):
+            def set_size(self, size):
+                held()
+
+        def call(method, **request):
+            requests.append(method.__name__)
+            return {}
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        try:
+            with fs.transaction:
+                fs.transaction.files = HeldFiles()
+                task = asyncio.ensure_future(
+                    fs._put_file(str(local), "s3://bucket/key", callback=HeldCallback())
+                )
+                await wait_until(lambda: state["started"] == 1)
+                task.cancel()
+                await wait_until(lambda: state["stopped"] == 1)
+                task.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await await_bounded(task, timeout=2)
+                release.set()
+                # Exits once the file has joined, or after 1 s.
+                await asyncio.to_thread(joined.wait, 1)
+        finally:
+            proceed.set()
+            release.set()
+        await asyncio.to_thread(shutdown_bounded, fs._thread_pool)
+        assert not joined.is_set()
+        assert requests == []
+
     @pytest.mark.asyncio
     async def test_put_pipe_return_values(self, tmp_path):
         # GH-1337: _put() and _pipe() return the values of their per-object

From 440b65267648754a4501d7b61395b16a1b56d935 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Thu, 12 Nov 2026 12:26:47 +0900
Subject: [PATCH 7/7] Raise for a transaction write that a stop scope discarded
 (GH-1337)

The fourth follow-up review found that a transaction write stopped by an
enclosing stop scope, rather than by a cancellation of its coroutine,
discarded its file and still returned. It now raises
concurrent.futures.CancelledError after the discard, which reaches the
coroutine as asyncio.CancelledError.

The test of the registration gap keeps the transaction open until the
file joins it or the write has ended, instead of for at most one second.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py            | 11 ++++-
 tests/pyathena/filesystem/test_s3_async.py | 52 +++++++++++++++++++---
 2 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 900782582..6c9880b1d 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -19,7 +19,7 @@
 import threading
 import warnings
 from collections.abc import Awaitable, Callable, Iterable
-from concurrent.futures import Future, wait
+from concurrent.futures import CancelledError, Future, wait
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
@@ -153,7 +153,12 @@ def __init__(self, fs: AioS3FileSystem, write: Callable[[Callable[..., S3File]],
         self._ended = False
 
     def run(self) -> None:
-        """Run the write, and drop and discard its file if it ends stopped."""
+        """Run the write, and drop and discard its file if it ends stopped.
+
+        Raises:
+            concurrent.futures.CancelledError: If the write was stopped, also
+                when the write itself succeeded, since its file is discarded.
+        """
         try:
             self._write(self.open)
         finally:
@@ -167,6 +172,8 @@ def run(self) -> None:
                     self._detach_all()
             if stopped:
                 self.discard()
+        if stopped:
+            raise CancelledError("Stopped before the transaction could commit the write.")
 
     def open(self, path: str, mode: str, **kwargs) -> S3File:
         """Open the file of the write, which joins the transaction unless the write is stopped.
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 00891c8f8..60bf10e1e 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -23,7 +23,12 @@
 from fsspec.callbacks import Callback
 
 from pyathena.filesystem.s3 import S3FileSystem
-from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile
+from pyathena.filesystem.s3_async import (
+    AioS3File,
+    AioS3FileSystem,
+    AioS3StreamedFile,
+    _S3TransactionWrite,
+)
 from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan
 from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope
 from pyathena.filesystem.s3_object import (
@@ -36,7 +41,7 @@
 )
 from pyathena.filesystem.s3_path import S3Path
 from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession
-from pyathena.util import RetryConfig, raise_if_stopped
+from pyathena.util import RetryConfig, raise_if_stopped, stop_retries_on
 from tests import ENV
 from tests.pyathena.conftest import connect
 from tests.pyathena.util import (
@@ -2184,8 +2189,20 @@ def call(method, **request):
             return {}
 
         fs._sync_fs._call = fs._sync_fs._core.call = call
+        write_ended = threading.Event()
+        run = _S3TransactionWrite.run
+
+        def run_and_signal(transaction_write):
+            try:
+                run(transaction_write)
+            finally:
+                write_ended.set()
+
         try:
-            with fs.transaction:
+            with (
+                mock.patch.object(_S3TransactionWrite, "run", run_and_signal),
+                fs.transaction,
+            ):
                 fs.transaction.files = HeldFiles()
                 task = asyncio.ensure_future(
                     fs._put_file(str(local), "s3://bucket/key", callback=HeldCallback())
@@ -2197,8 +2214,8 @@ def call(method, **request):
                 with pytest.raises(asyncio.CancelledError):
                     await await_bounded(task, timeout=2)
                 release.set()
-                # Exits once the file has joined, or after 1 s.
-                await asyncio.to_thread(joined.wait, 1)
+                # Exits once the file has joined, or once the write has ended.
+                await wait_until(lambda: joined.is_set() or write_ended.is_set())
         finally:
             proceed.set()
             release.set()
@@ -2206,6 +2223,31 @@ def call(method, **request):
         assert not joined.is_set()
         assert requests == []
 
+    @pytest.mark.asyncio
+    async def test_transaction_write_stopped_by_enclosing_scope(self):
+        # GH-1337: a transaction write that an enclosing stop scope, not a
+        # cancellation, has stopped raises instead of returning, since its
+        # file is discarded and left out of the transaction. The
+        # concurrent.futures.CancelledError of the stop reaches the coroutine
+        # as asyncio.CancelledError, which asyncio.wrap_future() converts it to.
+        fs = AioS3FileSystem(
+            key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True
+        )
+        requests = []
+
+        def call(method, **request):
+            requests.append(method.__name__)
+            return {}
+
+        fs._sync_fs._call = fs._sync_fs._core.call = call
+        stop = threading.Event()
+        stop.set()
+        with fs.transaction:
+            with stop_retries_on(stop), pytest.raises(asyncio.CancelledError):
+                await fs._pipe_file("s3://bucket/key", b"a")
+            assert not fs.transaction.files
+        assert requests == []
+
     @pytest.mark.asyncio
     async def test_put_pipe_return_values(self, tmp_path):
         # GH-1337: _put() and _pipe() return the values of their per-object