From 9ae2bd6814d2be6c8e6ef59fff68bad7c73cdfe6 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:00:57 +0900
Subject: [PATCH 1/7] Align AioS3FileSystem with S3FileSystem in transactions,
 open() and touch()

- pipe_file() and put_file() inside a transaction of an AioS3FileSystem
  now defer the commit to that transaction. They used to write through
  the internal S3FileSystem, which is not in the transaction, so the
  objects were written immediately and not rolled back.
- S3AioExecutor runs at most max_workers functions at once, so
  AioS3FileSystem.open(max_workers=N) bounds the parallel part uploads
  and range reads as S3FileSystem.open() does. max_workers is a new
  required keyword argument of S3AioExecutor.
- An AioS3FileSystem created with asynchronous=True has no event loop
  of its own, so its files use an S3ThreadPoolExecutor instead of an
  S3AioExecutor that always raised RuntimeError.
- touch() delegates to S3FileSystem.touch(): it passes the PutObject
  parameters, returns the response and raises ValueError for an
  existing object with truncate=False. _touch() returns the response.

Closes #977
Closes #954

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/aio.md                                |   2 +
 pyathena/filesystem/s3_async.py            | 129 +++++++++++++++++++-
 pyathena/filesystem/s3_executor.py         |  30 ++++-
 tests/pyathena/filesystem/test_s3_async.py | 134 ++++++++++++++++++++-
 4 files changed, 287 insertions(+), 8 deletions(-)

diff --git a/docs/aio.md b/docs/aio.md
index 492d50186..76d917d02 100644
--- a/docs/aio.md
+++ b/docs/aio.md
@@ -280,6 +280,8 @@ parallel operations. Two implementations are provided:
 `AioS3FileSystem` automatically uses `S3AioExecutor` for file handles, so multipart
 uploads and parallel range reads are dispatched through the event loop with
 `asyncio.to_thread()` instead of a separate `ThreadPoolExecutor` per file.
+At most `max_workers` of them run at once.
+An instance created with `asynchronous=True` has no event loop of its own, so its file handles use `S3ThreadPoolExecutor`.
 
 ### Usage with AioS3FSCursor
 
diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 8d1177a6c..52049f401 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -11,6 +11,8 @@
 
 import asyncio
 import logging
+import mimetypes
+import os
 from multiprocessing import cpu_count
 from typing import TYPE_CHECKING, Any, cast
 
@@ -18,7 +20,7 @@
 from fsspec.callbacks import _DEFAULT_CALLBACK
 
 from pyathena.filesystem.s3 import S3File, S3FileSystem
-from pyathena.filesystem.s3_executor import S3AioExecutor
+from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor
 from pyathena.filesystem.s3_object import (
     S3Metadata,
     S3MultipartUpload,
@@ -48,6 +50,8 @@ class AioS3FileSystem(AsyncFileSystem):
     File handles created by ``_open`` use ``S3AioExecutor`` so that parallel
     operations (range reads, multipart uploads) are dispatched through the event
     loop with ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file.
+    An instance created with ``asynchronous=True`` has no event loop of its own,
+    so its file handles use a ``ThreadPoolExecutor``.
 
     Attributes:
         _sync_fs: The internal synchronous S3FileSystem instance.
@@ -161,11 +165,93 @@ 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:
+        if self._intrans:
+            # The transaction belongs to this filesystem, not to the internal
+            # S3FileSystem, so defer the commit to it.
+            await asyncio.to_thread(self._pipe_file_in_transaction, path, value, mode, **kwargs)
+            return
         await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
 
+    def _open_in_transaction(self, path: str, **kwargs) -> S3File:
+        """Open a file to write that commits with this filesystem's transaction.
+
+        The file belongs to the internal ``S3FileSystem`` and uploads its
+        parts with its own thread pool. A file of this filesystem would
+        dispatch them to the event loop's default executor, whose threads
+        the callers of this method occupy while they wait.
+
+        Args:
+            path: S3 path (s3://bucket/key) to write to.
+            **kwargs: Additional parameters passed to ``S3FileSystem._open``.
+
+        Returns:
+            The file, registered with the transaction of this filesystem.
+        """
+        f = self._sync_fs._open(self._strip_protocol(path), "wb", autocommit=False, **kwargs)
+        self.transaction.files.append(f)
+        return f
+
+    def _pipe_file_in_transaction(
+        self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs
+    ) -> None:
+        """Write bytes into the path as a file of this filesystem's transaction.
+
+        Args:
+            path: S3 path (s3://bucket/key) to write to.
+            value: The bytes to write.
+            mode: "overwrite" or "create". With "create", raise
+                FileExistsError when the object already exists.
+            **kwargs: Additional parameters passed to ``S3FileSystem._open``.
+
+        Raises:
+            FileExistsError: If the mode is "create" and the path already
+                exists.
+        """
+        if mode == "create" and self._sync_fs.exists(path):
+            raise FileExistsError(path)
+        with self._open_in_transaction(path, **kwargs) as f:
+            f.write(value)
+
     async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
+        if self._intrans:
+            # See _pipe_file.
+            await asyncio.to_thread(self._put_file_in_transaction, lpath, rpath, callback, **kwargs)
+            return
         await asyncio.to_thread(self._sync_fs.put_file, lpath, rpath, callback=callback, **kwargs)
 
+    def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -> None:
+        """Upload a local file as a file of this filesystem's transaction.
+
+        Mirrors :meth:`S3FileSystem.put_file`, but defers the commit to the
+        transaction of this filesystem.
+
+        Args:
+            lpath: Local file path to upload.
+            rpath: S3 destination path (s3://bucket/key).
+            callback: Progress callback for tracking upload progress.
+            **kwargs: Additional S3 parameters (e.g., ContentType, StorageClass).
+        """
+        if os.path.isdir(lpath):
+            return
+        _, key, _ = self.parse_path(rpath)
+        if not key:
+            return
+
+        callback.set_size(os.path.getsize(lpath))
+        if "ContentType" not in kwargs:
+            content_type, _ = mimetypes.guess_type(lpath)
+            if content_type is not None:
+                kwargs["ContentType"] = content_type
+
+        with (
+            self._open_in_transaction(rpath, s3_additional_kwargs=kwargs) as remote,
+            open(lpath, "rb") as local,
+        ):
+            while data := local.read(remote.blocksize):
+                remote.write(data)
+                callback.relative_update(len(data))
+        self.invalidate_cache(rpath)
+
     async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
         await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs)
 
@@ -344,6 +430,23 @@ async def _find(
             return {f.name: f for f in files}
         return [f.name for f in files]
 
+    def _create_executor(self, max_workers: int) -> S3Executor:
+        """Create the executor for the parallel operations of a file.
+
+        An instance created with ``asynchronous=True`` has no event loop of
+        its own, so its files run the operations in a thread pool.
+
+        Args:
+            max_workers: The maximum number of operations that run at once.
+
+        Returns:
+            An ``S3AioExecutor`` on the event loop of this filesystem, or an
+            ``S3ThreadPoolExecutor`` if it has none.
+        """
+        if self._loop is None:
+            return S3ThreadPoolExecutor(max_workers=max_workers)
+        return S3AioExecutor(loop=self._loop, max_workers=max_workers)
+
     def _open(
         self,
         path: str,
@@ -367,7 +470,7 @@ def _open(
             path,
             mode,
             max_workers=max_workers,
-            executor=S3AioExecutor(loop=self._loop),
+            executor=self._create_executor(max_workers=max_workers),
             block_size=block_size,
             cache_type=cache_type,
             autocommit=autocommit,
@@ -576,8 +679,24 @@ def invalidate_cache(self, path: str | None = None) -> None:
         """
         self._sync_fs.invalidate_cache(path)
 
-    async def _touch(self, path: str, truncate: bool = True, **kwargs) -> None:
-        await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs)
+    async def _touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]:
+        return await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs)
+
+    def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]:
+        """Create an empty object with PutObject.
+
+        See :meth:`S3FileSystem.touch`.
+
+        Args:
+            path: S3 path (s3://bucket/key) of the object.
+            truncate: If True, replace an existing object with an empty one;
+                if False, raise if the object exists.
+            **kwargs: Additional parameters passed to the PutObject API.
+
+        Returns:
+            The PutObject response as a dictionary.
+        """
+        return self._sync_fs.touch(path, truncate=truncate, **kwargs)
 
 
 class AioS3File(S3File):
@@ -589,4 +708,6 @@ class AioS3File(S3File):
     through the ``S3Executor`` interface — the ``S3AioExecutor``
     provided by ``AioS3FileSystem`` dispatches them through the event loop with
     ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file.
+    For an ``AioS3FileSystem`` created with ``asynchronous=True``, it is an
+    ``S3ThreadPoolExecutor``.
     """
diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py
index 475e21387..d40660fed 100644
--- a/pyathena/filesystem/s3_executor.py
+++ b/pyathena/filesystem/s3_executor.py
@@ -80,26 +80,50 @@ class S3AioExecutor(S3Executor):
     ``concurrent.futures.Future`` objects that are compatible with
     ``as_completed()``, ``wait()`` and ``Future.cancel()``. As with
     ``ThreadPoolExecutor``, a future cannot be cancelled once its function has
-    started.
+    started. At most ``max_workers`` of the submitted functions run at once.
 
     This avoids thread-in-thread nesting when ``S3File`` is used from within
     ``asyncio.to_thread()`` calls (the pattern used by ``AioS3FileSystem``).
 
     Args:
         loop: A running asyncio event loop.
+        max_workers: The maximum number of submitted functions that run at once.
 
     Raises:
         RuntimeError: If the event loop is not running when ``submit`` is called.
     """
 
-    def __init__(self, loop: asyncio.AbstractEventLoop | None = None) -> None:
+    def __init__(self, loop: asyncio.AbstractEventLoop | None = None, *, max_workers: int) -> None:
         """Initialize the executor with the event loop to schedule work on.
 
         Args:
             loop: The asyncio event loop. ``submit`` raises ``RuntimeError``
                 if it is None or not running.
+            max_workers: The maximum number of submitted functions that run
+                at once.
+
+        Raises:
+            ValueError: If ``max_workers`` is not positive.
         """
+        if max_workers <= 0:
+            # As ThreadPoolExecutor does; a semaphore of 0 would never run anything.
+            raise ValueError("max_workers must be greater than 0")
         self._loop = loop
+        self._semaphore = asyncio.Semaphore(max_workers)
+
+    async def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
+        """Run the function in a thread once fewer than ``max_workers`` run.
+
+        Args:
+            fn: The blocking function to run.
+            *args: Positional arguments passed to the function.
+            **kwargs: Keyword arguments passed to the function.
+
+        Returns:
+            The return value of the function.
+        """
+        async with self._semaphore:
+            return await asyncio.to_thread(fn, *args, **kwargs)
 
     @override
     def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
@@ -142,7 +166,7 @@ def settle(task: Future[None]) -> None:
                 elif future.set_running_or_notify_cancel():
                     future.set_exception(task.exception())
 
-            task = asyncio.run_coroutine_threadsafe(asyncio.to_thread(run), self._loop)
+            task = asyncio.run_coroutine_threadsafe(self._run(run), self._loop)
             task.add_done_callback(settle)
             return future
         raise RuntimeError(
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index e85438d4e..2c97d03b6 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1,5 +1,7 @@
+import asyncio
 import os
 import tempfile
+import threading
 import time
 import urllib.parse
 import urllib.request
@@ -16,7 +18,12 @@
 
 from pyathena.filesystem.s3 import S3File, S3FileSystem
 from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem
-from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass
+from pyathena.filesystem.s3_object import (
+    S3MultipartUploadPart,
+    S3Object,
+    S3ObjectType,
+    S3StorageClass,
+)
 from tests import ENV
 from tests.pyathena.conftest import connect
 
@@ -193,6 +200,60 @@ async def test_copy_object_with_multipart_upload_invalid_block_size(self, block_
             )
         fs._sync_fs._call.assert_not_called()
 
+    @pytest.mark.parametrize("commit", [True, False])
+    def test_transaction_pipe_put_file(self, tmp_path, commit):
+        # GH-977: pipe_file() and put_file() join the transaction of this
+        # filesystem; they used to write through the internal S3FileSystem,
+        # which is not in the transaction, and were not rolled back.
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        put_object = fs._sync_fs._put_object = mock.MagicMock()
+        local = tmp_path / "local.txt"
+        local.write_bytes(b"local")
+
+        def write():
+            with fs.transaction:
+                fs.pipe_file("s3://bucket/k1", b"data")
+                fs.put_file(str(local), "s3://bucket/k2")
+                put_object.assert_not_called()
+                if not commit:
+                    raise RuntimeError("rollback")
+
+        if commit:
+            write()
+            assert [
+                (c.kwargs["key"], c.kwargs["body"], c.kwargs.get("ContentType"))
+                for c in put_object.call_args_list
+            ] == [("k1", b"data", None), ("k2", b"local", "text/plain")]
+        else:
+            with pytest.raises(RuntimeError, match="rollback"):
+                write()
+            put_object.assert_not_called()
+
+    def test_transaction_pipe_file_create_existing(self):
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        fs._sync_fs.exists = mock.MagicMock(return_value=True)
+        with fs.transaction, pytest.raises(FileExistsError):
+            fs.pipe_file("s3://bucket/key", b"data", mode="create")
+        fs._sync_fs.exists.assert_called_once_with("s3://bucket/key")
+
+    def test_touch_sync_wrapper(self):
+        # GH-977: touch() used to be fsspec's open()-based default, which
+        # dropped the PutObject parameters and returned None.
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        fs._sync_fs._call = mock.MagicMock(return_value={"ETag": '"e"'})
+
+        actual = fs.touch("s3://bucket/key", ContentType="text/plain")
+        assert isinstance(actual, dict)
+        assert fs._sync_fs._call.call_args.kwargs == {
+            "Bucket": "bucket",
+            "Key": "key",
+            "ContentType": "text/plain",
+        }
+
+        fs._sync_fs.exists = mock.MagicMock(return_value=True)
+        with pytest.raises(ValueError, match="Cannot touch the existing file"):
+            fs.touch("s3://bucket/key", truncate=False)
+
     @pytest.fixture(scope="class")
     def fs(self, request):
         if not hasattr(request, "param"):
@@ -968,6 +1029,77 @@ def test_open_max_workers(self):
             assert isinstance(f, AioS3File)
             assert f.max_workers == 2
 
+    @pytest.mark.parametrize("asynchronous", [False, True])
+    @pytest.mark.asyncio
+    async def test_open_parallel_requests(self, asynchronous):
+        # GH-954: max_workers bounds the parallel part uploads and range
+        # reads. GH-977: a filesystem created with asynchronous=True has no
+        # event loop of its own and used to fail to run them.
+        fs = AioS3FileSystem(
+            connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True
+        )
+        block_size = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE
+        size = block_size * 4
+        lock = threading.Lock()
+        state = {"active": 0, "peak": 0}
+
+        def track(result):
+            def call(**kwargs):
+                with lock:
+                    state["active"] += 1
+                    state["peak"] = max(state["peak"], state["active"])
+                time.sleep(0.1)
+                with lock:
+                    state["active"] -= 1
+                return result(**kwargs)
+
+            return mock.MagicMock(side_effect=call)
+
+        sync_fs = fs._sync_fs
+        sync_fs._create_multipart_upload = mock.MagicMock(
+            return_value=SimpleNamespace(upload_id="uploadid")
+        )
+        sync_fs._upload_part = track(
+            lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"e"'})
+        )
+        sync_fs._complete_multipart_upload = mock.MagicMock()
+        sync_fs._get_object = track(
+            lambda **kw: (kw["ranges"][0], b"a" * (kw["ranges"][1] - kw["ranges"][0]))
+        )
+        sync_fs.info = mock.MagicMock(
+            return_value=S3Object(
+                init={"Key": "key"},
+                type=S3ObjectType.S3_OBJECT_TYPE_FILE,
+                bucket="bucket",
+                key="key",
+            )
+        )
+        sync_fs.info.return_value.size = size
+
+        def write():
+            with fs.open("s3://bucket/key", "wb", block_size=block_size, max_workers=2) as f:
+                f.write(b"a" * size)
+
+        await asyncio.to_thread(write)
+        assert sync_fs._upload_part.call_count == 4
+        assert state["peak"] == 2
+
+        def read():
+            with fs.open(
+                "s3://bucket/key", "rb", block_size=block_size, cache_type="none", max_workers=2
+            ) as f:
+                return f.read()
+
+        state["peak"] = 0
+        assert await asyncio.to_thread(read) == b"a" * size
+        assert sync_fs._get_object.call_count == 4
+        assert state["peak"] == 2
+
+    def test_open_invalid_max_workers(self):
+        fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
+        with pytest.raises(ValueError, match="max_workers must be greater than 0"):
+            fs.open("s3://bucket/key", "wb", max_workers=0)
+
     def test_open_version_id(self):
         fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True)
         fs._sync_fs.info = mock.MagicMock(

From 2f65a46ac0175f22b5f0a801546958f74d2d2c0a Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:03:41 +0900
Subject: [PATCH 2/7] Write transaction files through open() of AioS3FileSystem

The parts of an AioS3File are only waited on at commit, which a
transaction runs in the caller's thread, so writing through open() does
not occupy the event loop's default executor and the internal-file
detour is unnecessary.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_async.py | 31 ++++++-------------------------
 1 file changed, 6 insertions(+), 25 deletions(-)

diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py
index 52049f401..7006c6e47 100644
--- a/pyathena/filesystem/s3_async.py
+++ b/pyathena/filesystem/s3_async.py
@@ -167,30 +167,11 @@ async def _pipe_file(
     ) -> None:
         if self._intrans:
             # The transaction belongs to this filesystem, not to the internal
-            # S3FileSystem, so defer the commit to it.
+            # S3FileSystem, so write through open() to defer the commit to it.
             await asyncio.to_thread(self._pipe_file_in_transaction, path, value, mode, **kwargs)
             return
         await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)
 
-    def _open_in_transaction(self, path: str, **kwargs) -> S3File:
-        """Open a file to write that commits with this filesystem's transaction.
-
-        The file belongs to the internal ``S3FileSystem`` and uploads its
-        parts with its own thread pool. A file of this filesystem would
-        dispatch them to the event loop's default executor, whose threads
-        the callers of this method occupy while they wait.
-
-        Args:
-            path: S3 path (s3://bucket/key) to write to.
-            **kwargs: Additional parameters passed to ``S3FileSystem._open``.
-
-        Returns:
-            The file, registered with the transaction of this filesystem.
-        """
-        f = self._sync_fs._open(self._strip_protocol(path), "wb", autocommit=False, **kwargs)
-        self.transaction.files.append(f)
-        return f
-
     def _pipe_file_in_transaction(
         self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs
     ) -> None:
@@ -201,7 +182,7 @@ def _pipe_file_in_transaction(
             value: The bytes to write.
             mode: "overwrite" or "create". With "create", raise
                 FileExistsError when the object already exists.
-            **kwargs: Additional parameters passed to ``S3FileSystem._open``.
+            **kwargs: Additional parameters passed to ``open()``.
 
         Raises:
             FileExistsError: If the mode is "create" and the path already
@@ -209,7 +190,7 @@ def _pipe_file_in_transaction(
         """
         if mode == "create" and self._sync_fs.exists(path):
             raise FileExistsError(path)
-        with self._open_in_transaction(path, **kwargs) as f:
+        with self.open(path, "wb", **kwargs) as f:
             f.write(value)
 
     async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
@@ -222,8 +203,8 @@ async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **
     def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -> None:
         """Upload a local file as a file of this filesystem's transaction.
 
-        Mirrors :meth:`S3FileSystem.put_file`, but defers the commit to the
-        transaction of this filesystem.
+        Mirrors :meth:`S3FileSystem.put_file`, but writes through ``open()``
+        of this filesystem.
 
         Args:
             lpath: Local file path to upload.
@@ -244,7 +225,7 @@ def _put_file_in_transaction(self, lpath: str, rpath: str, callback, **kwargs) -
                 kwargs["ContentType"] = content_type
 
         with (
-            self._open_in_transaction(rpath, s3_additional_kwargs=kwargs) as remote,
+            self.open(rpath, "wb", s3_additional_kwargs=kwargs) as remote,
             open(lpath, "rb") as local,
         ):
             while data := local.read(remote.blocksize):

From 4e95bc51a9f473909e39388d056d7c7f308b4f5b Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:16:10 +0900
Subject: [PATCH 3/7] Keep S3AioExecutor compatible and its limit through
 cancellation

- max_workers defaults to the S3File default, so S3AioExecutor(loop=...)
  keeps working.
- A cancelled function keeps running in its thread, so its permit is
  released when it returns, not when the coroutine is cancelled.
- The parallel request test waits for overlapping calls instead of
  relying on a sleep window.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_executor.py            | 27 +++++++++++++++--
 tests/pyathena/filesystem/test_s3_async.py    | 12 +++++---
 tests/pyathena/filesystem/test_s3_executor.py | 29 +++++++++++++++++++
 3 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py
index d40660fed..1e7b94129 100644
--- a/pyathena/filesystem/s3_executor.py
+++ b/pyathena/filesystem/s3_executor.py
@@ -15,6 +15,7 @@
 from collections.abc import Callable
 from concurrent.futures import Future
 from concurrent.futures.thread import ThreadPoolExecutor
+from multiprocessing import cpu_count
 from typing import Any, TypeVar
 
 from pyathena.util import override
@@ -93,7 +94,11 @@ class S3AioExecutor(S3Executor):
         RuntimeError: If the event loop is not running when ``submit`` is called.
     """
 
-    def __init__(self, loop: asyncio.AbstractEventLoop | None = None, *, max_workers: int) -> None:
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop | None = None,
+        max_workers: int = (cpu_count() or 1) * 5,
+    ) -> None:
         """Initialize the executor with the event loop to schedule work on.
 
         Args:
@@ -122,8 +127,24 @@ async def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
         Returns:
             The return value of the function.
         """
-        async with self._semaphore:
-            return await asyncio.to_thread(fn, *args, **kwargs)
+        await self._semaphore.acquire()
+        task = asyncio.ensure_future(asyncio.to_thread(fn, *args, **kwargs))
+        # Cancelling cannot stop the thread, so keep the permit until the
+        # function returns instead of until this coroutine is cancelled.
+        task.add_done_callback(self._release)
+        return await asyncio.shield(task)
+
+    def _release(self, task: asyncio.Future[Any]) -> None:
+        """Release the permit of a finished function.
+
+        Args:
+            task: The finished task that ran the function.
+        """
+        self._semaphore.release()
+        if not task.cancelled():
+            # Mark the exception as retrieved; the caller may have stopped
+            # waiting for it after a cancellation.
+            task.exception()
 
     @override
     def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 2c97d03b6..85de308a7 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1040,16 +1040,20 @@ async def test_open_parallel_requests(self, asynchronous):
         )
         block_size = S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE
         size = block_size * 4
-        lock = threading.Lock()
+        condition = threading.Condition()
         state = {"active": 0, "peak": 0}
 
         def track(result):
             def call(**kwargs):
-                with lock:
+                with condition:
                     state["active"] += 1
                     state["peak"] = max(state["peak"], state["active"])
-                time.sleep(0.1)
-                with lock:
+                    condition.notify_all()
+                    # Wait for a second call so that the calls overlap
+                    # whatever the scheduling, then leave time for more.
+                    condition.wait_for(lambda: state["active"] >= 2, timeout=5)
+                time.sleep(0.05)
+                with condition:
                     state["active"] -= 1
                 return result(**kwargs)
 
diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py
index 4168c04e9..5f28557c8 100644
--- a/tests/pyathena/filesystem/test_s3_executor.py
+++ b/tests/pyathena/filesystem/test_s3_executor.py
@@ -16,6 +16,12 @@
 
 
 class TestS3AioExecutor:
+    def test_init(self):
+        # max_workers is optional, as before it was added.
+        S3AioExecutor(loop=None)
+        with pytest.raises(ValueError, match="max_workers must be greater than 0"):
+            S3AioExecutor(loop=None, max_workers=0)
+
     def test_submit(self):
         async def main():
             executor = S3AioExecutor(loop=asyncio.get_running_loop())
@@ -122,3 +128,26 @@ async def main():
         future = asyncio.run(main())
 
         assert not future.cancelled()
+
+    @pytest.mark.asyncio
+    async def test_cancel_keeps_permit_until_function_returns(self):
+        # Cancelling a running function cannot stop its thread, so the next
+        # function must not start until it returns.
+        executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1)
+        started = threading.Event()
+        release = threading.Event()
+
+        def blocking():
+            started.set()
+            release.wait(5)
+            return "first"
+
+        first = executor.submit(blocking)
+        assert await asyncio.to_thread(started.wait, 5)
+        second = executor.submit(lambda: "second")
+        assert first.cancel()
+        await asyncio.sleep(0.1)
+        assert not second.done()
+
+        release.set()
+        assert await asyncio.wrap_future(second) == "second"

From cbe4e960530e12bac272c3abe5fb674fdb1e4c8e Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:50:43 +0900
Subject: [PATCH 4/7] Keep queued S3AioExecutor work cancellable

Shielding the thread kept the permit of a cancelled running function but
also let functions still queued for a thread run after S3File.discard()
aborted the upload. Since a future can no longer be cancelled once its
function has started, the plain semaphore already holds the permit of a
running function until it returns.

The parallel request test now holds each call until a third one runs or
a timeout passes, so an unbounded executor fails it whatever the
scheduling.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyathena/filesystem/s3_executor.py            | 20 ++---------------
 tests/pyathena/filesystem/test_s3_async.py    |  8 +++----
 tests/pyathena/filesystem/test_s3_executor.py | 22 -------------------
 3 files changed, 6 insertions(+), 44 deletions(-)

diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py
index 1e7b94129..c0f6adf2d 100644
--- a/pyathena/filesystem/s3_executor.py
+++ b/pyathena/filesystem/s3_executor.py
@@ -127,24 +127,8 @@ async def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
         Returns:
             The return value of the function.
         """
-        await self._semaphore.acquire()
-        task = asyncio.ensure_future(asyncio.to_thread(fn, *args, **kwargs))
-        # Cancelling cannot stop the thread, so keep the permit until the
-        # function returns instead of until this coroutine is cancelled.
-        task.add_done_callback(self._release)
-        return await asyncio.shield(task)
-
-    def _release(self, task: asyncio.Future[Any]) -> None:
-        """Release the permit of a finished function.
-
-        Args:
-            task: The finished task that ran the function.
-        """
-        self._semaphore.release()
-        if not task.cancelled():
-            # Mark the exception as retrieved; the caller may have stopped
-            # waiting for it after a cancellation.
-            task.exception()
+        async with self._semaphore:
+            return await asyncio.to_thread(fn, *args, **kwargs)
 
     @override
     def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index 85de308a7..f86ca73e2 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1049,10 +1049,10 @@ def call(**kwargs):
                     state["active"] += 1
                     state["peak"] = max(state["peak"], state["active"])
                     condition.notify_all()
-                    # Wait for a second call so that the calls overlap
-                    # whatever the scheduling, then leave time for more.
-                    condition.wait_for(lambda: state["active"] >= 2, timeout=5)
-                time.sleep(0.05)
+                    # Hold the call until a third one runs alongside it, so
+                    # that the calls overlap and an unbounded executor
+                    # exceeds max_workers=2 whatever the scheduling.
+                    condition.wait_for(lambda: state["active"] >= 3, timeout=0.5)
                 with condition:
                     state["active"] -= 1
                 return result(**kwargs)
diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py
index 5f28557c8..3df469d80 100644
--- a/tests/pyathena/filesystem/test_s3_executor.py
+++ b/tests/pyathena/filesystem/test_s3_executor.py
@@ -129,25 +129,3 @@ async def main():
 
         assert not future.cancelled()
 
-    @pytest.mark.asyncio
-    async def test_cancel_keeps_permit_until_function_returns(self):
-        # Cancelling a running function cannot stop its thread, so the next
-        # function must not start until it returns.
-        executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1)
-        started = threading.Event()
-        release = threading.Event()
-
-        def blocking():
-            started.set()
-            release.wait(5)
-            return "first"
-
-        first = executor.submit(blocking)
-        assert await asyncio.to_thread(started.wait, 5)
-        second = executor.submit(lambda: "second")
-        assert first.cancel()
-        await asyncio.sleep(0.1)
-        assert not second.done()
-
-        release.set()
-        assert await asyncio.wrap_future(second) == "second"

From 71a24b6d46b8d201d832e4b9503165ac735194aa Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:50:49 +0900
Subject: [PATCH 5/7] Settle the parallel request test

The parallel request test waits up to 5 s for two calls to overlap and
then 0.2 s for a third one, instead of relying on the scheduling within
a single short wait.

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

diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py
index f86ca73e2..9ba1e3cb7 100644
--- a/tests/pyathena/filesystem/test_s3_async.py
+++ b/tests/pyathena/filesystem/test_s3_async.py
@@ -1049,10 +1049,11 @@ def call(**kwargs):
                     state["active"] += 1
                     state["peak"] = max(state["peak"], state["active"])
                     condition.notify_all()
-                    # Hold the call until a third one runs alongside it, so
-                    # that the calls overlap and an unbounded executor
-                    # exceeds max_workers=2 whatever the scheduling.
-                    condition.wait_for(lambda: state["active"] >= 3, timeout=0.5)
+                    # Hold the call until another one overlaps it, then
+                    # briefly for a third one, which only an executor
+                    # without the limit runs.
+                    condition.wait_for(lambda: state["active"] >= 2, timeout=5)
+                    condition.wait_for(lambda: state["active"] >= 3, timeout=0.2)
                 with condition:
                     state["active"] -= 1
                 return result(**kwargs)

From dec6e128aca85eafdcb30204b2b1e0c8a8ac5c50 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:51:19 +0900
Subject: [PATCH 6/7] Drop a trailing blank line left by the rebase

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/pyathena/filesystem/test_s3_executor.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py
index 3df469d80..44d7c7b91 100644
--- a/tests/pyathena/filesystem/test_s3_executor.py
+++ b/tests/pyathena/filesystem/test_s3_executor.py
@@ -128,4 +128,3 @@ async def main():
         future = asyncio.run(main())
 
         assert not future.cancelled()
-

From 4ddec3715e640c2b4395076550aef23db6947bd5 Mon Sep 17 00:00:00 2001
From: laughingman7743 <laughingman7743@gmail.com>
Date: Sat, 3 Oct 2026 17:52:07 +0900
Subject: [PATCH 7/7] Cover the loop shutdown while a function waits for a
 permit

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/pyathena/filesystem/test_s3_executor.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py
index 44d7c7b91..9b1d005f6 100644
--- a/tests/pyathena/filesystem/test_s3_executor.py
+++ b/tests/pyathena/filesystem/test_s3_executor.py
@@ -79,7 +79,14 @@ async def main():
         assert not running.cancelled()
         assert pending.cancelled()
 
-    def test_loop_shutdown(self):
+    @pytest.mark.parametrize(
+        ("threads", "max_workers"),
+        [
+            (1, 5),  # The pending function waits for a thread.
+            (2, 1),  # The pending function waits for a permit.
+        ],
+    )
+    def test_loop_shutdown(self, threads, max_workers):
         # A function that has not started when the event loop shuts down is
         # never run, and its future is cancelled and settled, so that wait()
         # returns, instead of left pending.
@@ -96,8 +103,8 @@ def work():
 
         async def main():
             loop = asyncio.get_running_loop()
-            loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
-            executor = S3AioExecutor(loop=loop)
+            loop.set_default_executor(ThreadPoolExecutor(max_workers=threads))
+            executor = S3AioExecutor(loop=loop, max_workers=max_workers)
             running = executor.submit(work)
             pending = executor.submit(events.append, "pending finished")
             pending.add_done_callback(lambda _: settled.set())