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..4f25c0488 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,20 @@ 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, 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 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 through the buffered file path. Inside an diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index da6dfcb60..f02ecded5 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3147,8 +3147,20 @@ 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 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 + 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 +3170,18 @@ 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. 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, 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): @@ -3176,17 +3197,44 @@ 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 or "", kwargs.pop("compression", 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 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 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: 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 +4851,73 @@ 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, 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 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, + 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) + # 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 * block + 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 + block, wave_end)) + for start in range(wave_start, wave_end, block) + ], + ) + 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..33d810eff 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,44 @@ 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 + # 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), + 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..14296663a 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,299 @@ 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"] + + @pytest.mark.parametrize("compressed", [False, True]) + @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 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", + 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 + # 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, + # One worker, so that a whole-object read is one request. + max_workers=1, + cache_type="all", + ) + 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. + 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": + # 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) + + 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: dict[int, weakref.ref] = {} + held = [] + + class Body(bytearray): + # Unlike bytes, weakly referenceable. + pass + + get_object = fs.core.get_object + + def read(path, range_, **kwargs): + 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 + + 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() + 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): + # 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..1b92ddc0b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2251,6 +2251,62 @@ 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, + ) + serve_s3_keys(fs._sync_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 + @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 + ) + serve_s3_keys(fs._sync_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