From 000c8e62131cb74782eeb5a688a8239914a46be7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 00:54:22 +0900 Subject: [PATCH 1/6] Download in parallel ranges in get_file() (GH-1323) get_file() read the object one block at a time after the first, so a download ran one range request at a time whatever max_workers was. It now reads waves of ranges of the block size with S3Executor.map(), at most max_workers at once, capped at the request pool, and writes the bodies of each wave in order, dropping each once written, before it reads the next, so that at most one wave of bodies is held. The file is opened without a data cache, which open() constructed and which read the whole object with cache_type="all", and the reads keep the version, the IfMatch condition, and the parameters of open(). A failure while a wave is written drops its bodies, also from the completed frames of the traceback and of the chained errors, through the new S3MapResult.release(). With compression, the object is decompressed into the destination one block at a time, which raised AttributeError. AioS3FileSystem._get() pairs its paths as S3FileSystem.get() does and gives each download max(1, max_workers // n) range reads, n being the number of files or a smaller positive batch_size. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 4 +- docs/filesystem.md | 15 +- pyathena/filesystem/s3.py | 110 ++++++++- pyathena/filesystem/s3_async.py | 52 ++-- pyathena/filesystem/s3_executor.py | 24 ++ tests/pyathena/filesystem/test_s3.py | 261 ++++++++++++++++++++- tests/pyathena/filesystem/test_s3_async.py | 41 ++++ 7 files changed, 467 insertions(+), 40 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index ca8f3a90c..49f82b6d5 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -347,7 +347,9 @@ coroutines of their objects as tasks, `batch_size` of them at a time as in fsspe `_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 +Each download of `_get()` gets `max(1, max_workers // n)` range reads, as with +`S3FileSystem.get`, `n` being the number of files, or `batch_size` when it is positive and +smaller. 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`. diff --git a/docs/filesystem.md b/docs/filesystem.md index 979d0cd64..36e0519c5 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -134,8 +134,7 @@ and the error is raised once they have ended. A running multipart upload or copy no further part and is aborted instead of being completed, and a running download reads and writes no further block. A request already sent, such as the PutObject of a small object, runs to its end; a multipart copy that has been completed still copies its -annotations. With `cache_type="background"`, a block that fsspec's cache prefetches in -its own thread may still be read after the stop. The transfers that ended before the failure stay done, and their number is +annotations. The transfers that ended before the failure stay done, and their number is not bounded by `batch_size`. `mv` deletes the sources only once every copy has succeeded. `on_error` of `cat`, `cat_ranges`, and `cp` keeps its fsspec meaning; with `on_error="omit"` or `"return"` of `cat`, `"return"` of `cat_ranges`, or `"ignore"` of @@ -154,6 +153,18 @@ into it one after another, in order. `rm` deletes up to 1,000 objects per DeleteObjects request, `batch_size` requests at a time. Its `max_workers` argument is a deprecated alias of `batch_size`. +`get` and `get_file` download each object in waves of parallel range reads of the block +size: at most `max_workers` ranges at a time, and no more than the request pool runs at +once. The ranges of a wave are written in order before the next wave is read, so a +download holds at most one wave of data, `max_workers` × the block size. No data cache +reads the object, so `cache_type` and `cache_options` do not apply. The range reads are +conditional on the ETag, and the version, if any, of the object looked up when the +download starts: an object replaced during the download fails it with `OSError`, unless +the download reads a version other than `null`, which a replacement does not change. A +failed download leaves the part written so far in the local file, as fsspec's +`get_file` does. With `compression`, the object is decompressed into the destination, +one block at a time. + Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the block size (5 MiB by default); larger data is uploaded as a parallel multipart upload through the buffered file path. Inside an diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index da6dfcb60..23923fc82 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3147,8 +3147,18 @@ def get_file( ): """Download an S3 file to local filesystem. - Downloads a file from S3 to the local filesystem with progress tracking. - Reads the file in chunks to handle large files efficiently. + The object is read in waves of parallel range requests of the block + size, at most ``max_workers`` of them at once, but no more than the + request pool of the filesystem runs, and the bodies of each wave are + written in order before the next wave is read, so that at most one + wave of bodies is held (see :meth:`S3File._download`). The reads are + conditional on the version and the ETag of the object looked up when + the download starts, as with ``open()``, so that an object replaced + during the download fails it instead of mixing two objects, unless + the download reads a version other than ``null``, which a + replacement does not change. A failed + download leaves the part written so far in a local file ``lpath``, + as fsspec's ``AbstractFileSystem.get_file()`` does. As with fsspec's ``AbstractFileSystem.get_file()``, a directory ``rpath`` creates the local directory ``lpath``, and the parent @@ -3158,9 +3168,16 @@ def get_file( rpath: S3 source path (s3://bucket/key). lpath: Local destination path, or a file-like object to write to. Not needed when ``outfile`` is given. - callback: Progress callback for tracking download progress. + callback: Progress callback for tracking download progress, with + the size of the object. outfile: A file-like object to write to instead of ``lpath``. - **kwargs: Additional S3 parameters passed to open(). + **kwargs: The arguments of ``open()``, such as ``block_size``, + ``max_workers``, ``version_id``, and ``s3_additional_kwargs``, + and additional S3 parameters. ``cache_type`` and + ``cache_options`` are not used, since no data cache reads the + object. With ``compression``, the object is decompressed into + the destination, read one block at a time through ``open()``, + and the size of the callback is unknown. """ path_version_id = S3Path.parse(self._strip_protocol(rpath)).version_id if outfile is None and isfilelike(lpath): @@ -3176,17 +3193,37 @@ def get_file( os.makedirs(lpath, exist_ok=True) return + # The codec, if any, as fsspec's open() resolves it, with "infer" + # from the key without the version. + codec = get_compression( + S3Path.parse(self._strip_protocol(rpath)).key, kwargs.pop("compression", None) + ) + if codec is None: + # No data cache: the object is read in ranges, and a cache + # constructed when the file is opened, such as "all", would read + # the whole object. + kwargs.pop("cache_options", None) + kwargs["cache_type"] = "none" # The remote file is opened first so that no local file is created # when open() finds no object at the path. with contextlib.ExitStack() as stack: remote = stack.enter_context(self.open(rpath, "rb", **kwargs)) + reader: Any = None + if codec is not None: + # As fsspec's open() wraps the file, but closed with it. + reader = stack.enter_context(compr[codec](remote, mode="r")) if outfile is None: # Not abspath(), which would resolve ".." before symlinks. if parent := os.path.dirname(lpath): os.makedirs(parent, exist_ok=True) outfile = stack.enter_context(open(lpath, "wb")) - callback.set_size(remote.size) - while data := remote.read(remote.blocksize): + if reader is None: + callback.set_size(remote.size) + remote._download(outfile, callback) + return + # The size of the decompressed data is unknown. + callback.set_size(None) + while data := reader.read(remote.blocksize): # A block read when the stop came is not written. raise_if_stopped(f"write the contents of {rpath}") outfile.write(data) @@ -4803,6 +4840,67 @@ def _get_range(self, path: S3Path, range_: tuple[int, int], **kwargs: Any) -> by raise_if_stopped(f"read {path.uri}") return self.fs.core.get_object(path, range_, **kwargs) + def _download(self, outfile: Any, callback: Callback) -> None: + """Write the whole object into a writable file with parallel range reads. + + The object is read in waves of ranges of the block size, one wave of + at most ``max_workers`` ranges at a time, but no more than the request + pool of the filesystem runs at once, since the others would only wait + with their bodies held. Each wave is read with :meth:`S3Executor.map() + `, and its bodies are + written in order, each dropped once written, before the next wave is + read, so that at most one wave of bodies is held. The reads carry the + version and the conditions of the file, and each checks the stop + scopes when it starts; a body read when a stop came is not written. + + A failure or an interrupt while a wave is read is raised as + :meth:`S3Executor.map()` raises it. One while its bodies are written, + including a failure of the write or of the callback, drops the bodies + first, also from the completed frames of its traceback (see + :meth:`S3MapResult.release() + `). + + Args: + outfile: The writable file-like object. + callback: Progress callback, updated with the size of each body + once it is written. + + Raises: + concurrent.futures.CancelledError: If an enclosing stop scope has + been stopped. + """ + size = cast(int, self.size) + window = min(self.max_workers, self.fs._request_pool.max_workers) + path = S3Path(self.bucket, self.key, self.version_id) + read = functools.partial(self._get_range, path, **self._get_request_kwargs("get_object")) + wave = window * self.blocksize + result: S3MapResult[bytes] | None = None + try: + for wave_start in range(0, size, wave): + wave_end = min(size, wave_start + wave) + # Assigned within the try block, so that an interrupt right + # after the wave was read still drops its bodies. + result = self._executor.map( + read, + [ + (start, min(start + self.blocksize, wave_end)) + for start in range(wave_start, wave_end, self.blocksize) + ], + ) + bodies = result.outcomes + for i in range(len(bodies)): + # A body read when the stop came is not written. + raise_if_stopped(f"write the contents of {self.path}") + length = len(cast(bytes, bodies[i])) + outfile.write(bodies[i]) + # Not kept while the next one is written or read. + bodies[i] = b"" + callback.relative_update(length) + except BaseException as e: + if result is not None: + result.release(e) + raise + @staticmethod def _get_ranges( start: int, end: int, max_workers: int, worker_block_size: int diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 31bfc44d4..fdb1ccbf4 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -861,9 +861,13 @@ async def _get( ) -> None: """Copy files from S3 to the local filesystem. - See :meth:`S3FileSystem.get`. The downloads run as in fsspec's - ``_get()``, and a cancellation cancels them and waits for them, as - in :meth:`_copy`. + See :meth:`S3FileSystem.get`. The paths are paired as there, and the + downloads run as in fsspec's ``_get()``; a cancellation cancels them + and waits for them, as in :meth:`_copy`. As in + :meth:`S3FileSystem.get`, each download gets ``max(1, max_workers // + n)`` parallel range reads, or the caller's smaller ``max_workers``, + where ``n`` is the number of files, or ``batch_size`` when it is + positive and smaller. Args: rpath: Source S3 path, glob pattern, or list of them. @@ -872,28 +876,38 @@ async def _get( recursive: Whether to copy the directories with their contents. callback: Progress callback. maxdepth: Maximum depth of a recursive copy. - **kwargs: Additional parameters passed to ``_get_file()``. + **kwargs: Additional parameters passed to ``_get_file()``, and + ``batch_size``, the number of downloads that fsspec runs at + once. Raises: - ValueError: If a source with a version ID is paired, and a - destination lies outside ``lpath``. + ValueError: If a destination named after a source lies outside + ``lpath``. """ rpath = self._sync_fs._stringify_paths(rpath) - sources = [rpath] if isinstance(rpath, str) else rpath - if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): - root = make_path_posix(lpath) - pairs = await self._thread_pool.run( - self._sync_fs._copy_pairs, - S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), - isdir=LocalFileSystem().isdir, - ) - rpath, lpath = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] - check_contained(root, lpath) - if not rpath: - return + root = make_path_posix(lpath) if isinstance(lpath, (str, os.PathLike)) else lpath + pairs = await self._thread_pool.run( + self._sync_fs._copy_pairs, + S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), + isdir=LocalFileSystem().isdir, + ) + if isinstance(root, str): + check_contained(root, [lp for _, lp in pairs]) + if not pairs: + return + # fsspec's default batch size, when none is given, is not counted: + # at most max_workers downloads run at once in the threads of the + # filesystem, and more than max_workers get one range each anyway. + concurrency = len(pairs) + batch_size = kwargs.get("batch_size", self.batch_size) + if batch_size is not None and batch_size > 0: + concurrency = min(concurrency, batch_size) await _S3BatchTasks().run( super()._get( - rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs + [rp for rp, _ in pairs], + [lp for _, lp in pairs], + callback=callback, + **self._sync_fs._batch_kwargs(concurrency, kwargs), ) ) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 95e8760f1..8b4797171 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -162,6 +162,30 @@ class S3MapResult(Generic[T]): outcomes: list[T | BaseException] failure: BaseException | None + def release(self, error: BaseException) -> None: + """Drop the outcomes when the processing of the result fails, such as the bodies of reads. + + The outcomes are cleared, and so are the completed frames of the + traceback of the error and of the errors chained to it + (``__cause__`` and ``__context__``), which can keep an outcome, such + as the argument of a write that wrapped its own failure; + :func:`traceback.clear_frames` skips the frames that are executing. + The caller drops its own references to the outcomes. + + Args: + error: The error that the processing raises. + """ + self.outcomes.clear() + seen: set[int] = set() + errors: list[BaseException] = [error] + while errors: + current = errors.pop() + if id(current) in seen: + continue + seen.add(id(current)) + traceback.clear_frames(current.__traceback__) + errors.extend(e for e in (current.__cause__, current.__context__) if e is not None) + class _S3Batch(Generic[ItemT, T]): """The calls of one map of :class:`S3Executor`, and the state that they share. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index f653496e8..e6f3949df 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -2921,24 +2921,21 @@ def test_rm_batch_size(self, kwargs): def test_get_file_stopped_writes_no_further_block(self): # GH-1322 (B2): a download whose batch has stopped writes no further - # block, including the one read when the stop came. - fs = self._make_fs() + # block, including the ones read when the stop came (GH-1323: those + # of the wave). + fs, _ = self._make_object_fs(b"abcdefgh") stop = threading.Event() - remote = mock.MagicMock(size=8, blocksize=4) - remote.__enter__.return_value = remote - - blocks = [b"abcd", b"efgh", b""] + get_object = fs._client.get_object.side_effect - def read(size): - # The batch stops while the first block is being read. + def stopping(**request): + # The batch stops while the first wave, of one range, is read. stop.set() - return blocks.pop(0) + return get_object(**request) - remote.read.side_effect = read - fs.open = mock.MagicMock(return_value=remote) + fs._client.get_object.side_effect = stopping outfile = io.BytesIO() with stop_retries_on(stop), pytest.raises(CancelledError): - fs.get_file("s3://bucket/key?versionId=v1", outfile=outfile) + fs.get_file("s3://bucket/key", outfile=outfile, block_size=4, max_workers=1) assert outfile.getvalue() == b"" def test_multipart_copy_annotations_retry_after_stop(self): @@ -4807,6 +4804,246 @@ def test_get_file_file_like(self, tmp_path): fs.get_file("s3://bucket/key", outfile=outfile) assert outfile.getvalue() == b"data" + @pytest.mark.parametrize( + ("max_workers", "pool", "waves"), + [ + # Waves of max_workers ranges of the block size, each written + # before the next is read. + (2, 4, [["0-3", "4-7"], ["8-9"]]), + # One range at a time with one worker, not the whole object. + (1, 4, [["0-3"], ["4-7"], ["8-9"]]), + # No more ranges at once than the request pool runs. + (8, 2, [["0-3", "4-7"], ["8-9"]]), + (8, 4, [["0-3", "4-7", "8-9"]]), + ], + ) + def test_get_file_reads_waves_of_ranges(self, max_workers, pool, waves): + # GH-1323: get_file() read one block at a time after the first, so a + # download ran one range request at a time whatever max_workers was. + fs, _ = self._make_object_fs(b"0123456789") + fs._request_pool = S3ThreadPool(pool) + events = [] + get_object = fs._client.get_object.side_effect + + def record(**request): + events.append(("read", request["Range"].removeprefix("bytes="))) + return get_object(**request) + + fs._client.get_object.side_effect = record + + class Outfile(io.BytesIO): + def write(self, data): + events.append(("write", None)) + return super().write(data) + + outfile = Outfile() + callback = Callback() + fs.get_file( + "s3://bucket/key", + outfile=outfile, + callback=callback, + block_size=4, + max_workers=max_workers, + ) + assert outfile.getvalue() == b"0123456789" + observed = [] + for kind, group in itertools.groupby(events, key=lambda event: event[0]): + ranges = [range_ for _, range_ in group] + observed.append(sorted(ranges) if kind == "read" else len(ranges)) + assert observed == [part for wave in waves for part in (wave, len(wave))] + assert callback.size == callback.value == 10 + + @pytest.mark.parametrize("data", [b"", b"012", b"01234567", b"0123456789"]) + def test_get_file_contents(self, tmp_path, data): + # GH-1323: an empty object, less than a block, an exact multiple of + # the block size, and a remainder, with the progress of each body. + fs, ranges = self._make_object_fs(data) + lpath = tmp_path / "key" + updates = [] + + class Recording(Callback): + def relative_update(self, inc=1): + updates.append(inc) + super().relative_update(inc) + + fs.get_file( + "s3://bucket/key", str(lpath), callback=Recording(), block_size=4, max_workers=2 + ) + assert lpath.read_bytes() == data + assert updates == [min(4, len(data) - start) for start in range(0, len(data), 4)] + assert len(ranges) == len(updates) + + def test_get_file_constructs_no_data_cache(self): + # GH-1323: the download reads each range once; a data cache, which + # open() constructed, such as "all", read the whole object first. + fs, ranges = self._make_object_fs(b"0123456789") + outfile = io.BytesIO() + fs.get_file( + "s3://bucket/key", + outfile=outfile, + block_size=4, + max_workers=4, + cache_type="all", + cache_options={"unused": True}, + ) + assert outfile.getvalue() == b"0123456789" + assert sorted(ranges) == ["bytes=0-3", "bytes=4-7", "bytes=8-9"] + + def test_get_file_failed_range_reads_no_further_wave(self): + # GH-1323: a failed range stops the download: no further wave is + # read, and nothing of the failed wave is written. + fs, _ = self._make_object_fs(b"0123456789") + get_object = fs._client.get_object.side_effect + ranges = [] + + def failing(**request): + ranges.append(request["Range"]) + if request["Range"] == "bytes=4-7": + raise OSError("failed") + return get_object(**request) + + fs._client.get_object.side_effect = failing + outfile = io.BytesIO() + with pytest.raises(OSError, match="failed"): + fs.get_file("s3://bucket/key", outfile=outfile, block_size=4, max_workers=2) + assert outfile.getvalue() == b"" + assert sorted(ranges) == ["bytes=0-3", "bytes=4-7"] + + def test_get_file_holds_one_wave(self): + # GH-1323 (Q2): the bodies of a wave are dropped once written, so + # that they are not kept while the next wave is read. + fs, _ = self._make_object_fs(b"0123456789") + refs = {} + held = [] + + class Body(bytearray): + # Unlike bytes, weakly referenceable. + pass + + get_object = fs.core.get_object + + def read(path, range_, **kwargs): + held.append(sorted(start for start, ref in refs.items() if ref() is not None)) + body = Body(get_object(path, range_, **kwargs)) + refs[range_[0]] = weakref.ref(body) + return body + + fs.core.get_object = read + gc.disable() + try: + fs.get_file("s3://bucket/key", outfile=io.BytesIO(), block_size=4, max_workers=2) + finally: + gc.enable() + # The read of the second wave finds no body of the first. + assert held[-1] == [] + + @pytest.mark.parametrize("failure", ["write", "wrapped_write", "callback", "interrupt"]) + def test_get_file_failure_while_writing_keeps_no_bodies(self, failure): + # GH-1323 (F5): a failure while the bodies of a wave are written, + # including one that a write wraps in another error, keeps none of + # them in the traceback. + fs, _ = self._make_object_fs(b"0123456789") + refs = [] + + class Body(bytearray): + # Unlike bytes, weakly referenceable. + pass + + get_object = fs.core.get_object + + def read(path, range_, **kwargs): + body = Body(get_object(path, range_, **kwargs)) + refs.append(weakref.ref(body)) + return body + + fs.core.get_object = read + + class Outfile(io.BytesIO): + def write(self, data): + if failure == "write": + raise OSError("failed") + if failure == "wrapped_write": + try: + self._write(data) + except ValueError as e: + raise OSError("failed") from e + if failure == "interrupt": + raise KeyboardInterrupt + return super().write(data) + + def _write(self, data): + raise ValueError("inner") + + class Failing(Callback): + def relative_update(self, inc=1): + if failure == "callback": + raise OSError("failed") + + gc.disable() + try: + with pytest.raises((OSError, KeyboardInterrupt)) as info: + fs.get_file( + "s3://bucket/key", + outfile=Outfile(), + callback=Failing(), + block_size=4, + max_workers=2, + ) + assert info.value.__traceback__ is not None + assert len(refs) == 2 + assert all(ref() is None for ref in refs) + finally: + gc.enable() + + def test_get_file_request_parameters(self): + # GH-1323: every range of a download carries the version, the + # condition on the ETag looked up when it starts, and the + # parameters of the call, as the reads of open() do. + fs, _ = self._make_object_fs(b"0123456789") + fs.info.return_value = S3Object( + init={"Key": "key", "ETag": '"e"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + fs.info.return_value.size = 10 + + fs.get_file( + "s3://bucket/key?versionId=v1", + outfile=io.BytesIO(), + block_size=4, + max_workers=2, + RequestPayer="requester", + ) + requests = fs._client.get_object.call_args_list + assert len(requests) == 3 + assert { + (c.kwargs["VersionId"], c.kwargs["IfMatch"], c.kwargs["RequestPayer"]) for c in requests + } == {("v1", '"e"', "requester")} + + @pytest.mark.parametrize( + ("rpath", "compression", "compressed"), + [ + ("s3://bucket/key", "gzip", True), + # Inferred from the key, without the version. + ("s3://bucket/key.gz?versionId=v1", "infer", True), + # Nothing to infer: read in parallel ranges, with the size. + ("s3://bucket/key", "infer", False), + ], + ) + def test_get_file_compression(self, rpath, compression, compressed): + # GH-1323: with compression, the object is decompressed into the + # destination; the size of the decompressed data is unknown. It + # raised AttributeError, since the decompressor has no size. + data = b"0123456789" * 100 + fs, _ = self._make_object_fs(gzip.compress(data) if compressed else data) + outfile = io.BytesIO() + callback = Callback() + fs.get_file(rpath, outfile=outfile, callback=callback, compression=compression) + assert outfile.getvalue() == data + assert callback.size == (None if compressed else len(data)) + assert callback.value == len(data) + def test_cat_ranges_range(self): fs, ranges = self._make_object_fs(b"0123456789") diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 3c1a9263e..ae45d7db2 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2251,6 +2251,47 @@ async def test_put_pipe_return_values(self, tmp_path): assert await fs._put(str(local), PurePosixPath("bucket/b")) == [None] assert store == {"a", "b"} + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("files", "kwargs", "expected"), + [ + # A single download keeps the max_workers of the filesystem. + (1, {}, None), + (2, {}, 2), + # More downloads than max_workers (4) get one range each. + (8, {}, 1), + (2, {"batch_size": 1}, None), + (2, {"batch_size": -1}, 2), + # A smaller max_workers of the caller is kept. + (2, {"max_workers": 1}, 1), + ], + ) + async def test_get_window(self, tmp_path, files, kwargs, expected): + # GH-1323: as in S3FileSystem.get(), each download gets max(1, + # max_workers // n) parallel range reads, n being the downloads + # that run at once, so that a batch holds about max_workers bodies. + # Each download used to get max_workers of them. + fs = AioS3FileSystem( + key="dummy", + secret="dummy", + region_name="us-east-1", + max_workers=4, + skip_instance_cache=True, + ) + self._serve_keys(fs, {f"dir/{i}" for i in range(files)}) + windows = [] + + def get_file(rpath, lpath, **get_kwargs): + windows.append(get_kwargs.get("max_workers")) + + fs._sync_fs.get_file = get_file + await fs._get( + [f"s3://bucket/dir/{i}" for i in range(files)], + [str(tmp_path / str(i)) for i in range(files)], + **kwargs, + ) + assert windows == [expected] * files + @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 From d45a2d7ff6fba2ee6254c88a6036c91a208fe73e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 00:56:23 +0900 Subject: [PATCH 2/6] Infer no codec for a bucket path in get_file() (GH-1323) get_compression() infers the codec from the extension of the key, and a bucket path has none; it raised AttributeError for None. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 23923fc82..354f95fff 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3196,7 +3196,7 @@ def get_file( # The codec, if any, as fsspec's open() resolves it, with "infer" # from the key without the version. codec = get_compression( - S3Path.parse(self._strip_protocol(rpath)).key, kwargs.pop("compression", None) + S3Path.parse(self._strip_protocol(rpath)).key or "", kwargs.pop("compression", None) ) if codec is None: # No data cache: the object is read in ranges, and a cache From 23791eb4ef9ec9aac00e54f682c513a8a0e435bc Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 01:14:35 +0900 Subject: [PATCH 3/6] Address the independent review of parallel downloads (GH-1323) - A block size that is not positive reads the object in one range, as read() reads the rest; the waves read nothing for -1 and failed for 0. - The decompressing path of get_file() reads through its own cache of one block too, so that the caller's cache_type, such as "all", which read the whole object when the file was opened, or "background", which reads outside the stop scopes, does not apply there either. - AioS3FileSystem._get() makes the local paths of a destination list absolute, as fsspec's _get() does for one source, so that fsspec can create their parent directories. - The docstring of S3File._download() states that a pool thread that has just returned a body drops it when its call returns; the test of the one-wave bound waits for that, and the failed-range test lets the other range of the wave be requested first. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 6 +-- pyathena/filesystem/s3.py | 45 ++++++++++++--------- pyathena/filesystem/s3_async.py | 8 +++- tests/pyathena/filesystem/test_s3.py | 47 ++++++++++++++++++++-- tests/pyathena/filesystem/test_s3_async.py | 15 +++++++ 5 files changed, 94 insertions(+), 27 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 36e0519c5..2a64f589e 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -156,14 +156,14 @@ time. Its `max_workers` argument is a deprecated alias of `batch_size`. `get` and `get_file` download each object in waves of parallel range reads of the block size: at most `max_workers` ranges at a time, and no more than the request pool runs at once. The ranges of a wave are written in order before the next wave is read, so a -download holds at most one wave of data, `max_workers` × the block size. No data cache -reads the object, so `cache_type` and `cache_options` do not apply. The range reads are +download holds at most one wave of data, `max_workers` × the block size. `cache_type` +and `cache_options` do not apply: no data cache reads the object. The range reads are conditional on the ETag, and the version, if any, of the object looked up when the download starts: an object replaced during the download fails it with `OSError`, unless the download reads a version other than `null`, which a replacement does not change. A failed download leaves the part written so far in the local file, as fsspec's `get_file` does. With `compression`, the object is decompressed into the destination, -one block at a time. +read one block at a time through a cache of one block. Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the block size (5 MiB by default); larger data is uploaded as a parallel multipart upload diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 354f95fff..8ed13815b 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3174,10 +3174,10 @@ def get_file( **kwargs: The arguments of ``open()``, such as ``block_size``, ``max_workers``, ``version_id``, and ``s3_additional_kwargs``, and additional S3 parameters. ``cache_type`` and - ``cache_options`` are not used, since no data cache reads the - object. With ``compression``, the object is decompressed into - the destination, read one block at a time through ``open()``, - and the size of the callback is unknown. + ``cache_options`` are not used. With ``compression``, the + object is decompressed into the destination, read one block + at a time through ``open()``, and the size of the callback is + unknown. """ path_version_id = S3Path.parse(self._strip_protocol(rpath)).version_id if outfile is None and isfilelike(lpath): @@ -3198,12 +3198,13 @@ def get_file( codec = get_compression( S3Path.parse(self._strip_protocol(rpath)).key or "", kwargs.pop("compression", None) ) - if codec is None: - # No data cache: the object is read in ranges, and a cache - # constructed when the file is opened, such as "all", would read - # the whole object. - kwargs.pop("cache_options", None) - kwargs["cache_type"] = "none" + # Not the caller's cache: one such as "all" would read the whole + # object when the file is opened, and "background" reads outside + # the stop scopes of the caller. Without a codec, the object is read + # in ranges without a cache; with one, the decompressor's small reads + # go through a cache of one block at a time. + kwargs.pop("cache_options", None) + kwargs["cache_type"] = "none" if codec is None else "bytes" # The remote file is opened first so that no local file is created # when open() finds no object at the path. with contextlib.ExitStack() as stack: @@ -4843,15 +4844,18 @@ def _get_range(self, path: S3Path, range_: tuple[int, int], **kwargs: Any) -> by def _download(self, outfile: Any, callback: Callback) -> None: """Write the whole object into a writable file with parallel range reads. - The object is read in waves of ranges of the block size, one wave of - at most ``max_workers`` ranges at a time, but no more than the request - pool of the filesystem runs at once, since the others would only wait + The object is read in waves of ranges of the block size, or in one + range if the block size is not positive, one wave of at most + ``max_workers`` ranges at a time, but no more than the request pool + of the filesystem runs at once, since the others would only wait with their bodies held. Each wave is read with :meth:`S3Executor.map() `, and its bodies are written in order, each dropped once written, before the next wave is - read, so that at most one wave of bodies is held. The reads carry the - version and the conditions of the file, and each checks the stop - scopes when it starts; a body read when a stop came is not written. + read, so that the download holds at most one wave of bodies; a thread + of the pool that has just returned a body drops it when its call + returns. The reads carry the version and the conditions of the file, + and each checks the stop scopes when it starts; a body read when a + stop came is not written. A failure or an interrupt while a wave is read is raised as :meth:`S3Executor.map()` raises it. One while its bodies are written, @@ -4870,10 +4874,13 @@ def _download(self, outfile: Any, callback: Callback) -> None: been stopped. """ size = cast(int, self.size) + # A block size that is not positive reads the rest of the object, + # as read() does. + block = self.blocksize if self.blocksize > 0 else max(size, 1) window = min(self.max_workers, self.fs._request_pool.max_workers) path = S3Path(self.bucket, self.key, self.version_id) read = functools.partial(self._get_range, path, **self._get_request_kwargs("get_object")) - wave = window * self.blocksize + wave = window * block result: S3MapResult[bytes] | None = None try: for wave_start in range(0, size, wave): @@ -4883,8 +4890,8 @@ def _download(self, outfile: Any, callback: Callback) -> None: result = self._executor.map( read, [ - (start, min(start + self.blocksize, wave_end)) - for start in range(wave_start, wave_end, self.blocksize) + (start, min(start + block, wave_end)) + for start in range(wave_start, wave_end, block) ], ) bodies = result.outcomes diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index fdb1ccbf4..33d810eff 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -885,7 +885,13 @@ async def _get( ``lpath``. """ rpath = self._sync_fs._stringify_paths(rpath) - root = make_path_posix(lpath) if isinstance(lpath, (str, os.PathLike)) else lpath + # Local paths are made absolute, as fsspec's _get() does, so that + # it can create their parent directories. + root = ( + make_path_posix(lpath) + if isinstance(lpath, (str, os.PathLike)) + else [make_path_posix(p) if isinstance(p, (str, os.PathLike)) else p for p in lpath] + ) pairs = await self._thread_pool.run( self._sync_fs._copy_pairs, S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index e6f3949df..50c4a991b 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -4889,6 +4889,35 @@ def test_get_file_constructs_no_data_cache(self): assert outfile.getvalue() == b"0123456789" assert sorted(ranges) == ["bytes=0-3", "bytes=4-7", "bytes=8-9"] + @pytest.mark.parametrize("block_size", [0, -1]) + def test_get_file_non_positive_block_size(self, block_size): + # GH-1323: a block size that is not positive reads the object in one + # range, as read() reads the rest; the waves would read nothing. + fs, ranges = self._make_object_fs(b"0123456789") + outfile = io.BytesIO() + fs.get_file("s3://bucket/key", outfile=outfile, block_size=block_size, max_workers=1) + assert outfile.getvalue() == b"0123456789" + assert ranges == ["bytes=0-9"] + + def test_get_file_compression_ignores_cache_type(self): + # GH-1323: the decompressed download reads one block at a time + # through its own cache; the caller's "all" read the whole object + # when the file was opened. + data = bytes(range(256)) * 64 + compressed = gzip.compress(data, compresslevel=0) + fs, ranges = self._make_object_fs(compressed) + outfile = io.BytesIO() + fs.get_file( + "s3://bucket/key", + outfile=outfile, + compression="gzip", + block_size=1024, + cache_type="all", + cache_options={"unused": True}, + ) + assert outfile.getvalue() == data + assert f"bytes=0-{len(compressed) - 1}" not in ranges + def test_get_file_failed_range_reads_no_further_wave(self): # GH-1323: a failed range stops the download: no further wave is # read, and nothing of the failed wave is written. @@ -4899,6 +4928,9 @@ def test_get_file_failed_range_reads_no_further_wave(self): def failing(**request): ranges.append(request["Range"]) if request["Range"] == "bytes=4-7": + # Once the other range of the wave has been requested, which + # would otherwise find the batch stopped. + poll_until(lambda: "bytes=0-3" in ranges, "the first range was not requested") raise OSError("failed") return get_object(**request) @@ -4913,7 +4945,7 @@ def test_get_file_holds_one_wave(self): # GH-1323 (Q2): the bodies of a wave are dropped once written, so # that they are not kept while the next wave is read. fs, _ = self._make_object_fs(b"0123456789") - refs = {} + refs: dict[int, weakref.ref] = {} held = [] class Body(bytearray): @@ -4923,7 +4955,15 @@ class Body(bytearray): get_object = fs.core.get_object def read(path, range_, **kwargs): - held.append(sorted(start for start, ref in refs.items() if ref() is not None)) + if range_[0] >= 8: + # The second wave: no body of the first is alive, once the + # threads that returned them have returned themselves. + poll_until( + lambda: all(ref() is None for ref in refs.values()), + "a body of the first wave is held", + timeout=2, + ) + held.append(range_) body = Body(get_object(path, range_, **kwargs)) refs[range_[0]] = weakref.ref(body) return body @@ -4934,8 +4974,7 @@ def read(path, range_, **kwargs): fs.get_file("s3://bucket/key", outfile=io.BytesIO(), block_size=4, max_workers=2) finally: gc.enable() - # The read of the second wave finds no body of the first. - assert held[-1] == [] + assert held == [(8, 10)] @pytest.mark.parametrize("failure", ["write", "wrapped_write", "callback", "interrupt"]) def test_get_file_failure_while_writing_keeps_no_bodies(self, failure): diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index ae45d7db2..4cb861336 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2292,6 +2292,21 @@ def get_file(rpath, lpath, **get_kwargs): ) assert windows == [expected] * files + @pytest.mark.asyncio + @pytest.mark.parametrize("rpath", ["s3://bucket/key", ["s3://bucket/key"]]) + async def test_get_relative_destination_list(self, tmp_path, monkeypatch, rpath): + # GH-1323: relative local paths in a list are made absolute, as + # fsspec's _get() does for one source, so that it can create their + # parent directories. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + self._serve_keys(fs, {"key"}) + monkeypatch.chdir(tmp_path) + + await fs._get(rpath, ["download.txt"]) + assert (tmp_path / "download.txt").read_bytes() == b"aa" + @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 From a44cfc3b54796a9140211074a1dad29df13de73f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 01:25:54 +0900 Subject: [PATCH 4/6] Normalize the block size of decompressed downloads (GH-1323) The follow-up review found that the block cache of the decompressing path failed with a block size that is not positive, which now takes the filesystem's, that its read-ahead was described as one block, and that the public descriptions of the one-wave bound omitted the body that a pool thread keeps until its call returns. The test of the ignored cache_type reads with one worker, so that a whole-object read is one request that it can see. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 6 ++++-- pyathena/filesystem/s3.py | 19 +++++++++++++------ tests/pyathena/filesystem/test_s3.py | 26 +++++++++++++++++++------- 3 files changed, 36 insertions(+), 15 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 2a64f589e..4f25c0488 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -156,14 +156,16 @@ time. Its `max_workers` argument is a deprecated alias of `batch_size`. `get` and `get_file` download each object in waves of parallel range reads of the block size: at most `max_workers` ranges at a time, and no more than the request pool runs at once. The ranges of a wave are written in order before the next wave is read, so a -download holds at most one wave of data, `max_workers` × the block size. `cache_type` +download holds at most one wave of data, `max_workers` × the block size, apart from a +range that a thread of the request pool has just read until its call returns. `cache_type` and `cache_options` do not apply: no data cache reads the object. The range reads are conditional on the ETag, and the version, if any, of the object looked up when the download starts: an object replaced during the download fails it with `OSError`, unless the download reads a version other than `null`, which a replacement does not change. A failed download leaves the part written so far in the local file, as fsspec's `get_file` does. With `compression`, the object is decompressed into the destination, -read one block at a time through a cache of one block. +read in order through a block cache, which reads ahead of the decompressor; a block size +that is not positive takes the filesystem's. Writes with `pipe`/`pipe_file` issue a single PutObject request for data up to the block size (5 MiB by default); larger data is uploaded as a parallel multipart upload diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 8ed13815b..9f3ad6c5d 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3150,8 +3150,10 @@ def get_file( The object is read in waves of parallel range requests of the block size, at most ``max_workers`` of them at once, but no more than the request pool of the filesystem runs, and the bodies of each wave are - written in order before the next wave is read, so that at most one - wave of bodies is held (see :meth:`S3File._download`). The reads are + written in order before the next wave is read, so that the download + holds at most one wave of bodies; a thread of the pool that has just + read a range drops its body when its call returns (see + :meth:`S3File._download`). The reads are conditional on the version and the ETag of the object looked up when the download starts, as with ``open()``, so that an object replaced during the download fails it instead of mixing two objects, unless @@ -3175,9 +3177,10 @@ def get_file( ``max_workers``, ``version_id``, and ``s3_additional_kwargs``, and additional S3 parameters. ``cache_type`` and ``cache_options`` are not used. With ``compression``, the - object is decompressed into the destination, read one block - at a time through ``open()``, and the size of the callback is - unknown. + object is decompressed into the destination, read in order + through ``open()`` with a block cache, whose block size, if + not positive, is that of the filesystem, and the size of the + callback is unknown. """ path_version_id = S3Path.parse(self._strip_protocol(rpath)).version_id if outfile is None and isfilelike(lpath): @@ -3202,9 +3205,13 @@ def get_file( # object when the file is opened, and "background" reads outside # the stop scopes of the caller. Without a codec, the object is read # in ranges without a cache; with one, the decompressor's small reads - # go through a cache of one block at a time. + # go through a block cache, which needs a positive block size. kwargs.pop("cache_options", None) kwargs["cache_type"] = "none" if codec is None else "bytes" + if codec is not None and not ( + isinstance(kwargs.get("block_size"), int) and kwargs["block_size"] > 0 + ): + kwargs["block_size"] = self.default_block_size # The remote file is opened first so that no local file is created # when open() finds no object at the path. with contextlib.ExitStack() as stack: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 50c4a991b..5a48ff21e 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -4889,15 +4889,26 @@ def test_get_file_constructs_no_data_cache(self): assert outfile.getvalue() == b"0123456789" assert sorted(ranges) == ["bytes=0-3", "bytes=4-7", "bytes=8-9"] + @pytest.mark.parametrize("compressed", [False, True]) @pytest.mark.parametrize("block_size", [0, -1]) - def test_get_file_non_positive_block_size(self, block_size): + def test_get_file_non_positive_block_size(self, block_size, compressed): # GH-1323: a block size that is not positive reads the object in one - # range, as read() reads the rest; the waves would read nothing. - fs, ranges = self._make_object_fs(b"0123456789") + # range, as read() reads the rest, and with compression takes the + # filesystem's; the waves read nothing, and the block cache of the + # decompressing path failed. + data = b"0123456789" + fs, ranges = self._make_object_fs(gzip.compress(data) if compressed else data) outfile = io.BytesIO() - fs.get_file("s3://bucket/key", outfile=outfile, block_size=block_size, max_workers=1) - assert outfile.getvalue() == b"0123456789" - assert ranges == ["bytes=0-9"] + fs.get_file( + "s3://bucket/key", + outfile=outfile, + block_size=block_size, + max_workers=1, + compression="gzip" if compressed else None, + ) + assert outfile.getvalue() == data + if not compressed: + assert ranges == ["bytes=0-9"] def test_get_file_compression_ignores_cache_type(self): # GH-1323: the decompressed download reads one block at a time @@ -4912,8 +4923,9 @@ def test_get_file_compression_ignores_cache_type(self): outfile=outfile, compression="gzip", block_size=1024, + # One worker, so that a whole-object read is one request. + max_workers=1, cache_type="all", - cache_options={"unused": True}, ) assert outfile.getvalue() == data assert f"bytes=0-{len(compressed) - 1}" not in ranges From 9ff12b211ee498fa05d8f241f5f47800fede4fcc Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 01:31:37 +0900 Subject: [PATCH 5/6] Fall back to a positive block size for decompressed downloads (GH-1323) The second follow-up review found that a filesystem whose default_block_size is not positive, which the constructor accepts, gave the block cache of the decompressing path a negative block size again. DEFAULT_BLOCK_SIZE applies then. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 9 ++++++--- tests/pyathena/filesystem/test_s3.py | 8 +++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 9f3ad6c5d..f02ecded5 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3179,8 +3179,9 @@ def get_file( ``cache_options`` are not used. With ``compression``, the object is decompressed into the destination, read in order through ``open()`` with a block cache, whose block size, if - not positive, is that of the filesystem, and the size of the - callback is unknown. + not positive, is that of the filesystem, or + ``DEFAULT_BLOCK_SIZE`` if that is not positive either, and the + size of the callback is unknown. """ path_version_id = S3Path.parse(self._strip_protocol(rpath)).version_id if outfile is None and isfilelike(lpath): @@ -3211,7 +3212,9 @@ def get_file( if codec is not None and not ( isinstance(kwargs.get("block_size"), int) and kwargs["block_size"] > 0 ): - kwargs["block_size"] = self.default_block_size + kwargs["block_size"] = ( + self.default_block_size if self.default_block_size > 0 else self.DEFAULT_BLOCK_SIZE + ) # The remote file is opened first so that no local file is created # when open() finds no object at the path. with contextlib.ExitStack() as stack: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 5a48ff21e..14296663a 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -4893,11 +4893,13 @@ def test_get_file_constructs_no_data_cache(self): @pytest.mark.parametrize("block_size", [0, -1]) def test_get_file_non_positive_block_size(self, block_size, compressed): # GH-1323: a block size that is not positive reads the object in one - # range, as read() reads the rest, and with compression takes the - # filesystem's; the waves read nothing, and the block cache of the - # decompressing path failed. + # range, as read() reads the rest, and with compression takes a + # positive default; the waves read nothing, and the block cache of + # the decompressing path failed. data = b"0123456789" fs, ranges = self._make_object_fs(gzip.compress(data) if compressed else data) + # Accepted by the constructor; DEFAULT_BLOCK_SIZE then applies. + fs.default_block_size = -1 outfile = io.BytesIO() fs.get_file( "s3://bucket/key", From 05eb56be4586bb1e74faa5c7e84381df50f5d970 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 02:05:41 +0900 Subject: [PATCH 6/6] Serve the keys of the aio download tests with serve_s3_keys() (GH-1323) The rebase onto #1394 moved the helper of TestAioS3FileSystem to tests.pyathena.util.serve_s3_keys(); the two tests added here still called the removed method. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 4cb861336..1b92ddc0b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2278,7 +2278,7 @@ async def test_get_window(self, tmp_path, files, kwargs, expected): max_workers=4, skip_instance_cache=True, ) - self._serve_keys(fs, {f"dir/{i}" for i in range(files)}) + serve_s3_keys(fs._sync_fs, {f"dir/{i}" for i in range(files)}) windows = [] def get_file(rpath, lpath, **get_kwargs): @@ -2301,7 +2301,7 @@ async def test_get_relative_destination_list(self, tmp_path, monkeypatch, rpath) fs = AioS3FileSystem( key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True ) - self._serve_keys(fs, {"key"}) + serve_s3_keys(fs._sync_fs, {"key"}) monkeypatch.chdir(tmp_path) await fs._get(rpath, ["download.txt"])