Download in parallel ranges in get_file() (GH-1323)

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`.
 

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

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()
+        <pyathena.filesystem.s3_executor.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()
+        <pyathena.filesystem.s3_executor.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

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),
             )
         )
 

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.