From b09d4cbb9d131e07d635aa772255d736f75e6d42 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 17:29:00 +0900 Subject: [PATCH 01/10] Move the path pairing into S3PathPairing and add S3Path.target Add the public S3PathPairing (pyathena/filesystem/s3_path_pairing.py), built once by S3FileSystem and exposed as S3FileSystem.pairing and AioS3FileSystem.pairing. Its copy_pairs(), move_pairs() and delete_paths() replace the private _copy_paths(), _move_paths() and _expand_delete_paths(); copy_pairs() and move_pairs() both return (source, destination) pairs. S3Path.target, the object that a write to the path replaces, replaces _move_target(). The pairing results are unchanged. Closes #1063. Co-Authored-By: Claude Opus 5.5 --- docs/api/filesystem.rst | 3 + docs/filesystem.md | 15 ++ pyathena/filesystem/s3.py | 198 ++-------------- pyathena/filesystem/s3_async.py | 22 +- pyathena/filesystem/s3_path.py | 10 + pyathena/filesystem/s3_path_pairing.py | 215 ++++++++++++++++++ tests/pyathena/filesystem/test_s3.py | 2 + tests/pyathena/filesystem/test_s3_path.py | 14 ++ .../filesystem/test_s3_path_pairing.py | 117 ++++++++++ 9 files changed, 411 insertions(+), 185 deletions(-) create mode 100644 pyathena/filesystem/s3_path_pairing.py create mode 100644 tests/pyathena/filesystem/test_s3_path_pairing.py diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index 170e1c2f..b7c78f56 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -49,6 +49,9 @@ S3 Paths .. autoclass:: pyathena.filesystem.s3_path.S3Path :members: +.. autoclass:: pyathena.filesystem.s3_path_pairing.S3PathPairing + :members: + S3 Core ------- diff --git a/docs/filesystem.md b/docs/filesystem.md index f3f62fbd..893d8ee1 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -348,6 +348,21 @@ instead. `copy_object_annotation()` copies one annotation onto the destination a the upload completes, with GetObjectAnnotation and PutObjectAnnotation. The filesystems' `cp_file()`, `copy()` and `mv()` run these plans. +`S3FileSystem.pairing` is an `S3PathPairing`, the expansion and pairing of the paths +that `copy()`, `get()`, `mv()` and `rm()` operate on, as fsspec pairs them; +`AioS3FileSystem.pairing` is the same object. `copy_pairs()` and `move_pairs()` return +the `(source, destination)` pairs, and `delete_paths()` the paths that `rm()` deletes. A +path with a version ID names that version, and its destination is named after its key. +The pairing lists and looks up paths through the filesystem and its cache, and writes +nothing. + +```python +for source, destination in fs.pairing.copy_pairs( + "s3://YOUR_S3_BUCKET/src/", "s3://YOUR_S3_BUCKET/dst/", recursive=True +): + print(source, "->", destination) +``` + ## Async filesystem `AioS3FileSystem` provides the same functionality on top of fsspec's diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 2a0d0912..79d97774 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -9,12 +9,10 @@ import mimetypes import os.path import time -from collections import Counter from collections.abc import Callable, Mapping from concurrent.futures import Future, as_completed, wait from copy import deepcopy from datetime import datetime -from glob import has_magic from io import BytesIO from multiprocessing import cpu_count from re import Pattern @@ -29,9 +27,9 @@ from fsspec.callbacks import _DEFAULT_CALLBACK, Callback from fsspec.compression import compr from fsspec.core import get_compression -from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep +from fsspec.implementations.local import LocalFileSystem, make_path_posix from fsspec.spec import AbstractBufferedFile -from fsspec.utils import check_contained, isfilelike, other_paths, tokenize +from fsspec.utils import check_contained, isfilelike, tokenize import pyathena from pyathena.connection import Connection @@ -50,6 +48,7 @@ S3StorageClass, ) from pyathena.filesystem.s3_path import S3Path +from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig, override _logger = logging.getLogger(__name__) @@ -258,12 +257,18 @@ def __init__( retry_config=retry_config, request_kwargs={"RequestPayer": "requester"} if requester_pays else None, ) + self._pairing = S3PathPairing(self) @property def core(self) -> S3Core: """The typed S3 operations that the filesystem sends its requests with.""" return self._core + @property + def pairing(self) -> S3PathPairing: + """The pairing of the paths that the filesystem copies, moves and deletes.""" + return self._pairing + @property def _client(self) -> BaseClient: """The boto3 S3 client of the filesystem (``core.client``).""" @@ -1100,46 +1105,9 @@ def rm(self, path, recursive=False, maxdepth=None, **kwargs) -> None: ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = self._expand_delete_paths(path, recursive=recursive, maxdepth=maxdepth) + paths = self.pairing.delete_paths(path, recursive=recursive, maxdepth=maxdepth) self._delete_objects(paths, **kwargs) - def _expand_delete_paths( - self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None - ) -> list[str]: - """Expand the paths that ``rm`` deletes. - - Args: - path: S3 path or list of paths. - recursive: Whether to include all objects below the paths. - maxdepth: Maximum depth to expand when ``recursive`` is True. - - Returns: - The paths with a version ID as given, followed by the expansion - of the other paths by ``expand_path``. - - Raises: - ValueError: If a path is a bucket. - """ - paths = [path] if isinstance(path, str) else list(path) - versioned_paths, unversioned_paths = [], [] - for p in paths: - s3_path = S3Path.parse(p) - # expand_path strips the slashes of "bucket//" to the bucket. - if s3_path.is_bucket: - raise ValueError("Cannot delete the bucket.") - if s3_path.version_id: - versioned_paths.append(p) - else: - unversioned_paths.append(p) - - if unversioned_paths: - # Versioned paths are deleted as given, without the lookup that - # expand_path makes for them with recursive. - unversioned_paths = self.expand_path( - unversioned_paths, recursive=recursive, maxdepth=maxdepth - ) - return versioned_paths + unversioned_paths - def _create_executor(self, max_workers: int) -> S3Executor: """Create an executor strategy for parallel operations. @@ -1449,129 +1417,19 @@ def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: return copied = [ p1 - for p1, p2 in self._move_paths(path1, path2, recursive=recursive, maxdepth=maxdepth) + for p1, p2 in self.pairing.move_pairs( + path1, path2, recursive=recursive, maxdepth=maxdepth + ) if self._copy_file(p1, p2, **kwargs) ] self._delete_objects(copied) - def _move_paths( - self, path1, path2, recursive: bool = False, maxdepth: int | None = None - ) -> list[tuple[str, str]]: - """Pair the sources and destinations of a move as fsspec's ``copy()`` does. - - Args: - path1: Source S3 path, glob pattern, or list of paths. - path2: Destination S3 path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. - - Returns: - The source and destination paths, except the sources whose - destination is the source itself or, for a ``null`` version, the - key of the source. - - Raises: - ValueError: If two sources have the same destination, or a - destination is another source, including one left in place, - except for a directory with no object at its key, which is not - copied. - """ - paths1, paths2 = self._copy_paths(path1, path2, recursive=recursive, maxdepth=maxdepth) - # The paths are copied as given, and compared by what they name. - named = [ - (p1, p2, self._move_target(p1), self._move_target(p2)) - for p1, p2 in zip(paths1, paths2, strict=False) - ] - pairs = [(p1, p2) for p1, p2, source, dest in named if source != dest] - moved = [(p1, source, dest) for p1, _, source, dest in named if source != dest] - # The sources left in place count too; a copy onto one overwrites it. - sources = {source for _, _, source, _ in named} - counts = Counter(dest for _, _, dest in moved) - # A source with another source below it may be a directory. - directories: set[str] = set() - for source in sources: - parent = source.rpartition("/")[0] - while parent and parent not in directories: - directories.add(parent) - parent = parent.rpartition("/")[0] - # A directory without an object at its key is not copied, so it - # writes no destination and is left out of the checks. A path with a - # version always names an object. - writers = [ - (source, dest) - for p1, source, dest in moved - if not ( - (counts[dest] > 1 or dest in sources) - and source in directories - and not S3Path.parse(p1).version_id - and self._head_object(source) is None - ) - ] - counts = Counter(dest for _, dest in writers) - for _, dest in writers: - if counts[dest] > 1: - raise ValueError("Cannot move several paths to the same destination.") - if dest in sources: - raise ValueError("Cannot move a path onto another path that is moved.") - return pairs - - def _copy_paths( - self, - path1: str | list[str], - path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, - isdir: Callable[[str], bool] | None = None, - ) -> tuple[list[str], list[str]]: - """Pair the sources of a copy with their destinations as fsspec's ``copy()`` does. - - A source with a version ID names that version of an object: it is not - a glob pattern, and its destination is named after its key without - the version. - - Args: - path1: Source S3 path, glob pattern, or list of them. - path2: Destination path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. - isdir: Whether a destination path is a directory, by default - ``self.isdir``; ``get()`` passes the local filesystem's. - - Returns: - The sources and their destinations. Both are empty if ``path1`` - is a string that matches only directories without ``recursive``. - """ - if isinstance(path1, list) and isinstance(path2, list): - return path1, path2 - source_is_str = isinstance(path1, str) - paths1 = self.expand_path(path1, recursive=recursive, maxdepth=maxdepth) - if source_is_str and (not recursive or maxdepth is not None): - # Non-recursive glob does not copy directories. - paths1 = [p for p in paths1 if not (trailing_sep(p) or self.isdir(p))] - if not paths1: - return [], [] - glob = isinstance(path1, str) and has_magic(path1) and not S3Path.has_version_id(path1) - # The destination is looked up only when it decides the mapping. - exists = source_is_str and ( - (glob and len(paths1) == 1) - or ( - not glob - and not trailing_sep(path1) - and isinstance(path2, str) - and (trailing_sep(path2) or (isdir or self.isdir)(path2)) - ) - ) - names = [S3Path.split_version_id(p)[0] for p in paths1] - return paths1, other_paths(names, path2, exists=exists, flatten=not source_is_str) - def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kwargs) -> None: """Copy files within S3. As fsspec's ``copy()``, except that a source with a version ID copies that version of the object to a destination named after its key, as - ``_copy_paths`` pairs them. + :meth:`S3PathPairing.copy_pairs` pairs them. Args: path1: Source S3 path, glob pattern, or list of them. @@ -1586,9 +1444,10 @@ def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kw """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - path1, path2 = self._copy_paths(path1, path2, recursive=recursive, maxdepth=maxdepth) - if not path1: + pairs = self.pairing.copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + if not pairs: return + path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] super().copy( path1, path2, recursive=recursive, maxdepth=maxdepth, on_error=on_error, **kwargs ) @@ -1600,7 +1459,7 @@ def get( As fsspec's ``get()``, except that a source with a version ID downloads that version of the object to a local path named after its - key, as ``_copy_paths`` pairs them. Those destinations are checked to + key, as :meth:`S3PathPairing.copy_pairs` pairs them. Those destinations are checked to lie under ``lpath``. Args: @@ -1619,9 +1478,10 @@ def get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): root = make_path_posix(lpath) - rpath, lpath = self._copy_paths( + pairs = self.pairing.copy_pairs( 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 @@ -1629,24 +1489,6 @@ def get( rpath, lpath, recursive=recursive, callback=callback, maxdepth=maxdepth, **kwargs ) - def _move_target(self, path: str) -> str: - """Return what a path of a move names, for comparing the paths. - - A write to a key replaces its ``null`` version, which the objects of a - bucket without versioning have, so that version names the key itself. - - Args: - path: S3 path, possibly with a version ID. - - Returns: - The path in ``bucket/key`` form, with the version ID unless it is - ``null``. - """ - s3_path = S3Path.parse(path) - if s3_path.version_id == "null": - return s3_path.name - return str(s3_path) - def cp_file( self, path1: str, path2: str, recursive=False, maxdepth=None, on_error=None, **kwargs ): diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 9c84b516..5a776395 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -35,6 +35,7 @@ S3ObjectVersion, ) from pyathena.filesystem.s3_path import S3Path +from pyathena.filesystem.s3_path_pairing import S3PathPairing if TYPE_CHECKING: from datetime import datetime @@ -142,6 +143,11 @@ def core(self) -> S3Core: """The typed S3 operations of the wrapped ``S3FileSystem``.""" return self._sync_fs.core + @property + def pairing(self) -> S3PathPairing: + """The pairing of the paths of the wrapped ``S3FileSystem``.""" + return self._sync_fs.pairing + @staticmethod def parse_path(path: str) -> tuple[str, str | None, str | None]: """Parse an S3 path into its bucket, key and version ID. @@ -324,7 +330,7 @@ async def _rm( OSError: If S3 could not delete some of the objects. """ paths = await asyncio.to_thread( - self._sync_fs._expand_delete_paths, path, recursive=recursive, maxdepth=maxdepth + self.pairing.delete_paths, path, recursive=recursive, maxdepth=maxdepth ) await self._delete_objects(paths, **kwargs) @@ -372,7 +378,7 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N if path1 == path2: return pairs = await asyncio.to_thread( - self._sync_fs._move_paths, path1, path2, recursive=recursive, maxdepth=maxdepth + self.pairing.move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth ) # Every copy finishes before a failure is raised, as in fsspec's # _copy(), and nothing is deleted after a failure. @@ -429,11 +435,12 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - path1, path2 = await asyncio.to_thread( - self._sync_fs._copy_paths, path1, path2, recursive=recursive, maxdepth=maxdepth + pairs = await asyncio.to_thread( + self.pairing.copy_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth ) - if not path1: + if not pairs: return + path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] await super()._copy( path1, path2, @@ -468,14 +475,15 @@ async def _get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): root = make_path_posix(lpath) - rpath, lpath = await asyncio.to_thread( - self._sync_fs._copy_paths, + pairs = await asyncio.to_thread( + self.pairing.copy_pairs, 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 diff --git a/pyathena/filesystem/s3_path.py b/pyathena/filesystem/s3_path.py index 6b2af357..427f495f 100644 --- a/pyathena/filesystem/s3_path.py +++ b/pyathena/filesystem/s3_path.py @@ -133,6 +133,16 @@ def uri(self) -> str: """The path as an ``s3://`` URI, with its version ID query, if any.""" return f"s3://{self}" + @property + def target(self) -> S3Path: + """The object that a write to the path replaces, for comparing paths. + + A write to a key replaces its ``null`` version, which the objects of a + bucket without versioning have, so a ``null`` version names the key + itself. Any other path is its own target. + """ + return self.with_version_id(None) if self.version_id == "null" else self + def with_version_id(self, version_id: str | None) -> S3Path: """Return the path with another version ID. diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py new file mode 100644 index 00000000..bca0eed8 --- /dev/null +++ b/pyathena/filesystem/s3_path_pairing.py @@ -0,0 +1,215 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +"""The pairing of the paths that the S3 filesystem copies, moves and deletes.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable +from glob import has_magic +from typing import TYPE_CHECKING + +from fsspec.implementations.local import trailing_sep +from fsspec.utils import other_paths + +from pyathena.filesystem.s3_path import S3Path + +if TYPE_CHECKING: + from pyathena.filesystem.s3 import S3FileSystem + + +class S3PathPairing: + """The paths that ``copy()``, ``get()``, ``mv()`` and ``rm()`` operate on. + + Expands the paths, glob patterns and directories as fsspec's ``copy()`` + and ``rm()`` do, and pairs the sources with their destinations, except + that a path with a version ID names that version of an object: it is not + a glob pattern, nothing is expanded below it, and its destination is named + after its key without the version. The paths are compared by what they + name (see :attr:`S3Path.target`). + + ``S3FileSystem`` builds one and exposes it as ``S3FileSystem.pairing``; + ``AioS3FileSystem.pairing`` is the same object. The lookups go through the + filesystem, with its cache. + + Example: + >>> for source, destination in fs.pairing.copy_pairs( + ... "s3://bucket/dir/", "s3://bucket/copy/", recursive=True + ... ): + ... print(source, "->", destination) + """ + + def __init__(self, fs: S3FileSystem) -> None: + """Create the pairing of a filesystem. + + Args: + fs: The filesystem that expands and looks up the paths. + """ + self._fs = fs + + def copy_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + isdir: Callable[[str], bool] | None = None, + ) -> list[tuple[str, str]]: + """Pair the sources of a copy with their destinations as fsspec's ``copy()`` does. + + Args: + path1: Source S3 path, glob pattern, or list of them. + path2: Destination path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + isdir: Whether a destination path is a directory, by default the + filesystem's ``isdir``; ``get()`` passes the local + filesystem's. + + Returns: + The sources and their destinations. When both ``path1`` and + ``path2`` are lists, they are paired as given. Empty if ``path1`` + is a string that matches only directories without ``recursive``. + + Raises: + ValueError: If ``maxdepth`` is less than 1. + FileNotFoundError: If the expansion of ``path1`` matches nothing: + a glob pattern without matches, or a missing path with + ``recursive``. + """ + if isinstance(path1, list) and isinstance(path2, list): + return list(zip(path1, path2, strict=False)) + fs = self._fs + source_is_str = isinstance(path1, str) + paths1 = fs.expand_path(path1, recursive=recursive, maxdepth=maxdepth) + if source_is_str and (not recursive or maxdepth is not None): + # Non-recursive glob does not copy directories. + paths1 = [p for p in paths1 if not (trailing_sep(p) or fs.isdir(p))] + if not paths1: + return [] + glob = isinstance(path1, str) and has_magic(path1) and not S3Path.has_version_id(path1) + # The destination is looked up only when it decides the mapping. + exists = source_is_str and ( + (glob and len(paths1) == 1) + or ( + not glob + and not trailing_sep(path1) + and isinstance(path2, str) + and (trailing_sep(path2) or (isdir or fs.isdir)(path2)) + ) + ) + names = [S3Path.split_version_id(p)[0] for p in paths1] + paths2 = other_paths(names, path2, exists=exists, flatten=not source_is_str) + return list(zip(paths1, paths2, strict=True)) + + def move_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + ) -> list[tuple[str, str]]: + """Pair the sources of a move with their destinations as :meth:`copy_pairs` does. + + Args: + path1: Source S3 path, glob pattern, or list of paths. + path2: Destination S3 path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + + Returns: + The sources and their destinations, except the sources whose + destination is the source itself or, for a ``null`` version, the + key of the source. + + Raises: + ValueError: If two sources have the same destination, or a + destination is another source, including one left in place, + except for a directory with no object at its key, which is not + copied. Also if ``maxdepth`` is less than 1. + FileNotFoundError: If the expansion of ``path1`` matches nothing, + as for :meth:`copy_pairs`. + """ + # The paths are moved as given, and compared by what they name. + named = [ + (p1, p2, str(S3Path.parse(p1).target), str(S3Path.parse(p2).target)) + for p1, p2 in self.copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + ] + pairs = [(p1, p2) for p1, p2, source, dest in named if source != dest] + moved = [(p1, source, dest) for p1, _, source, dest in named if source != dest] + # The sources left in place count too; a copy onto one overwrites it. + sources = {source for _, _, source, _ in named} + counts = Counter(dest for _, _, dest in moved) + # A source with another source below it may be a directory. + directories: set[str] = set() + for source in sources: + parent = source.rpartition("/")[0] + while parent and parent not in directories: + directories.add(parent) + parent = parent.rpartition("/")[0] + # A directory without an object at its key is not copied, so it + # writes no destination and is left out of the checks. A path with a + # version always names an object. + writers = [ + (source, dest) + for p1, source, dest in moved + if not ( + (counts[dest] > 1 or dest in sources) + and source in directories + and not S3Path.parse(p1).version_id + and self._fs._head_object(source) is None + ) + ] + counts = Counter(dest for _, dest in writers) + for _, dest in writers: + if counts[dest] > 1: + raise ValueError("Cannot move several paths to the same destination.") + if dest in sources: + raise ValueError("Cannot move a path onto another path that is moved.") + return pairs + + def delete_paths( + self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None + ) -> list[str]: + """Expand the paths that ``rm()`` deletes. + + Args: + path: S3 path or list of paths. + recursive: Whether to include all objects below the paths. + maxdepth: Maximum depth to expand when ``recursive`` is True. + + Returns: + The paths with a version ID as given, followed by the expansion + of the other paths by the filesystem's ``expand_path``. + + Raises: + ValueError: If a path is a bucket, or ``maxdepth`` is less than 1. + FileNotFoundError: If the expansion of the paths without a version + ID matches nothing, as for :meth:`copy_pairs`. + """ + paths = [path] if isinstance(path, str) else list(path) + versioned_paths, unversioned_paths = [], [] + for p in paths: + s3_path = S3Path.parse(p) + # expand_path strips the slashes of "bucket//" to the bucket. + if s3_path.is_bucket: + raise ValueError("Cannot delete the bucket.") + if s3_path.version_id: + versioned_paths.append(p) + else: + unversioned_paths.append(p) + + if unversioned_paths: + # Versioned paths are deleted as given, without the lookup that + # expand_path makes for them with recursive. + unversioned_paths = self._fs.expand_path( + unversioned_paths, recursive=recursive, maxdepth=maxdepth + ) + return versioned_paths + unversioned_paths diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 17bbb858..c50b1d7f 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -41,6 +41,7 @@ from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path +from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -174,6 +175,7 @@ def _make_fs(): fs._core = S3Core(client, retry_config=RetryConfig()) # The requests of the core and of the filesystem go to one mock. fs._call = fs._core.call = mock.MagicMock() + fs._pairing = S3PathPairing(fs) fs.max_workers = 4 fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE fs.allow_bucket_creation = False diff --git a/tests/pyathena/filesystem/test_s3_path.py b/tests/pyathena/filesystem/test_s3_path.py index cf143f33..0dda470d 100644 --- a/tests/pyathena/filesystem/test_s3_path.py +++ b/tests/pyathena/filesystem/test_s3_path.py @@ -123,3 +123,17 @@ def test_with_version_id(self): assert path == S3Path("bucket", "key", "v1") with pytest.raises(AttributeError): path.key = "other" # type: ignore[misc] + + @pytest.mark.parametrize( + ("path", "expected"), + [ + # A write to the key replaces its "null" version. + (S3Path("bucket", "key", "null"), S3Path("bucket", "key")), + # Any other path is its own target. + (S3Path("bucket", "key", "v1"), S3Path("bucket", "key", "v1")), + (S3Path("bucket", "key"), S3Path("bucket", "key")), + (S3Path("bucket"), S3Path("bucket")), + ], + ) + def test_target(self, path, expected): + assert path.target == expected diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py new file mode 100644 index 00000000..4a16fee6 --- /dev/null +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -0,0 +1,117 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +from unittest import mock + +import pytest + +from pyathena.filesystem.s3_async import AioS3FileSystem +from pyathena.filesystem.s3_path_pairing import S3PathPairing +from tests.pyathena.filesystem import test_s3 + + +def _make_fs(keys): + # An S3FileSystem whose requests are answered from the keys of "bucket". + fs = test_s3.TestS3FileSystem._make_fs() + test_s3.TestS3FileSystem._serve_keys(fs, keys) + return fs + + +class TestS3PathPairing: + def test_filesystems_share_one_pairing(self): + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + assert isinstance(fs.pairing, S3PathPairing) + assert fs.pairing is fs._sync_fs.pairing + assert fs._sync_fs.pairing is fs._sync_fs.pairing + + @pytest.mark.parametrize( + ("path1", "path2", "kwargs", "expected"), + [ + # The directory and the objects below it, as fsspec's copy() pairs + # them. + ( + "s3://bucket/d", + "s3://bucket/out", + {"recursive": True}, + [ + ("bucket/d", "s3://bucket/out"), + ("bucket/d/a", "s3://bucket/out/a"), + ("bucket/d/b", "s3://bucket/out/b"), + ], + ), + # A version is copied to a destination named after its key. + ( + "s3://bucket/b?versionId=v1", + "s3://bucket/d/", + {}, + [("bucket/b?versionId=v1", "s3://bucket/d/b")], + ), + # Without recursive, a directory is not copied. + ("s3://bucket/d", "s3://bucket/out", {}, []), + # Lists are paired as given. + ( + ["s3://bucket/b", "s3://bucket/d/a"], + ["s3://bucket/x", "s3://bucket/y"], + {}, + [("s3://bucket/b", "s3://bucket/x"), ("s3://bucket/d/a", "s3://bucket/y")], + ), + ], + ) + def test_copy_pairs(self, path1, path2, kwargs, expected): + fs = _make_fs({"d/a", "d/b", "b"}) + assert fs.pairing.copy_pairs(path1, path2, **kwargs) == expected + + def test_copy_pairs_isdir(self): + # The given isdir, such as the local filesystem's for get(), decides + # whether the destination is a directory. + fs = _make_fs({"b"}) + isdir = mock.MagicMock(return_value=True) + + pairs = fs.pairing.copy_pairs("s3://bucket/b?versionId=v1", "/tmp/out", isdir=isdir) + + assert pairs == [("bucket/b?versionId=v1", "/tmp/out/b")] + isdir.assert_called_once_with("/tmp/out") + + def test_move_pairs(self): + # The "null" version of a key moved onto the key stays in place. + fs = _make_fs({"b", "d/a"}) + + pairs = fs.pairing.move_pairs( + ["s3://bucket/b?versionId=null", "s3://bucket/d/a"], + ["s3://bucket/b", "s3://bucket/z"], + ) + + assert pairs == [("s3://bucket/d/a", "s3://bucket/z")] + + @pytest.mark.parametrize( + ("path2", "match"), + [ + (["s3://bucket/o", "s3://bucket/o"], "same destination"), + (["s3://bucket/z", "s3://bucket/n"], "another path that is moved"), + ], + ) + def test_move_pairs_conflicts(self, path2, match): + fs = _make_fs({"b", "z"}) + with pytest.raises(ValueError, match=match): + fs.pairing.move_pairs(["s3://bucket/b", "s3://bucket/z"], path2) + + def test_delete_paths(self): + # The versions are deleted as given, before the expanded paths. + fs = _make_fs({"d/a", "d/b"}) + + paths = fs.pairing.delete_paths( + ["s3://bucket/d", "s3://bucket/b?versionId=v1"], recursive=True + ) + + assert paths == ["s3://bucket/b?versionId=v1", "bucket/d", "bucket/d/a", "bucket/d/b"] + + @pytest.mark.parametrize("path", ["s3://bucket", "s3://bucket/", "s3://bucket//"]) + def test_delete_paths_bucket(self, path): + fs = _make_fs(set()) + with pytest.raises(ValueError, match="Cannot delete the bucket"): + fs.pairing.delete_paths(path) + fs._call.assert_not_called() From 422c3ad046dd5e8c9c58c6be4c289d4f4aee8062 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:06:29 +0900 Subject: [PATCH 02/10] Make S3PathPairing a pure planner and keep the lookups in the adapters S3PathPairing no longer holds the filesystem, and the filesystems no longer store it, so S3FileSystem is freed by reference counting again. Its static rules take the expanded paths and the results of the lookups that they ask for (skips_directories, looks_up_destination, conflict_candidates) and raise ValueError when a needed lookup is not passed. S3FileSystem and AioS3FileSystem each expand the paths and make the lookups with their own requests and cache, in _copy_pairs, _move_pairs and _delete_paths; aio uses its async _expand_path and _isdir instead of running the sync pairing in a thread. The pairing results are unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 23 +- pyathena/filesystem/s3.py | 111 ++++++- pyathena/filesystem/s3_async.py | 116 +++++-- pyathena/filesystem/s3_path.py | 8 +- pyathena/filesystem/s3_path_pairing.py | 309 +++++++++++------- tests/pyathena/filesystem/test_s3.py | 54 ++- tests/pyathena/filesystem/test_s3_async.py | 41 ++- .../filesystem/test_s3_path_pairing.py | 188 +++++++---- 8 files changed, 619 insertions(+), 231 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 893d8ee1..96e1a190 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -348,17 +348,22 @@ instead. `copy_object_annotation()` copies one annotation onto the destination a the upload completes, with GetObjectAnnotation and PutObjectAnnotation. The filesystems' `cp_file()`, `copy()` and `mv()` run these plans. -`S3FileSystem.pairing` is an `S3PathPairing`, the expansion and pairing of the paths -that `copy()`, `get()`, `mv()` and `rm()` operate on, as fsspec pairs them; -`AioS3FileSystem.pairing` is the same object. `copy_pairs()` and `move_pairs()` return -the `(source, destination)` pairs, and `delete_paths()` the paths that `rm()` deletes. A -path with a version ID names that version, and its destination is named after its key. -The pairing lists and looks up paths through the filesystem and its cache, and writes -nothing. +## Path pairing + +`S3PathPairing` holds the rules by which `copy()`, `get()`, `mv()` and `rm()` pair +their paths, as fsspec pairs them, except that a path with a version ID names that +version, and its destination is named after its key. The rules are pure functions: +the filesystems expand the paths, make the lookups that the rules ask for +(`skips_directories()`, `looks_up_destination()` and `conflict_candidates()`), and +pass the results to `copy_pairs()`, `move_pairs()` and `delete_paths()`. A rule that +needs a lookup that is not passed raises `ValueError`. ```python -for source, destination in fs.pairing.copy_pairs( - "s3://YOUR_S3_BUCKET/src/", "s3://YOUR_S3_BUCKET/dst/", recursive=True +from pyathena.filesystem.s3_path_pairing import S3PathPairing + +sources = fs.expand_path("s3://YOUR_S3_BUCKET/src/", recursive=True) +for source, destination in S3PathPairing.copy_pairs( + "s3://YOUR_S3_BUCKET/src/", "s3://YOUR_S3_BUCKET/dst/", sources ): print(source, "->", destination) ``` diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 79d97774..59ff6b5e 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -27,7 +27,7 @@ from fsspec.callbacks import _DEFAULT_CALLBACK, Callback from fsspec.compression import compr from fsspec.core import get_compression -from fsspec.implementations.local import LocalFileSystem, make_path_posix +from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep from fsspec.spec import AbstractBufferedFile from fsspec.utils import check_contained, isfilelike, tokenize @@ -257,18 +257,12 @@ def __init__( retry_config=retry_config, request_kwargs={"RequestPayer": "requester"} if requester_pays else None, ) - self._pairing = S3PathPairing(self) @property def core(self) -> S3Core: """The typed S3 operations that the filesystem sends its requests with.""" return self._core - @property - def pairing(self) -> S3PathPairing: - """The pairing of the paths that the filesystem copies, moves and deletes.""" - return self._pairing - @property def _client(self) -> BaseClient: """The boto3 S3 client of the filesystem (``core.client``).""" @@ -1105,9 +1099,35 @@ def rm(self, path, recursive=False, maxdepth=None, **kwargs) -> None: ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = self.pairing.delete_paths(path, recursive=recursive, maxdepth=maxdepth) + paths = self._delete_paths(path, recursive=recursive, maxdepth=maxdepth) self._delete_objects(paths, **kwargs) + def _delete_paths( + self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None + ) -> list[str]: + """Expand the paths that ``rm()`` deletes (see :meth:`S3PathPairing.delete_paths`). + + Args: + path: S3 path or list of paths. + recursive: Whether to include all objects below the paths. + maxdepth: Maximum depth to expand when ``recursive`` is True. + + Returns: + The paths with a version ID as given, followed by the expansion + of the other paths by ``expand_path``. + + Raises: + ValueError: If a path is a bucket. + """ + versioned_paths, unversioned_paths = S3PathPairing.delete_paths(path) + if not unversioned_paths: + # Versioned paths are deleted as given, without the lookup that + # expand_path makes for them with recursive. + return versioned_paths + return versioned_paths + self.expand_path( + unversioned_paths, recursive=recursive, maxdepth=maxdepth + ) + def _create_executor(self, max_workers: int) -> S3Executor: """Create an executor strategy for parallel operations. @@ -1417,13 +1437,78 @@ def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: return copied = [ p1 - for p1, p2 in self.pairing.move_pairs( - path1, path2, recursive=recursive, maxdepth=maxdepth - ) + for p1, p2 in self._move_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) if self._copy_file(p1, p2, **kwargs) ] self._delete_objects(copied) + def _move_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + ) -> list[tuple[str, str]]: + """Pair and check the paths of a move (see :meth:`S3PathPairing.move_pairs`). + + Args: + path1: Source S3 path, glob pattern, or list of paths. + path2: Destination S3 path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + + Returns: + The sources and destinations that are moved. + + Raises: + ValueError: If the move has conflicting paths. + """ + pairs = self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + missing = { + source + for source in S3PathPairing.conflict_candidates(pairs) + if self._head_object(source) is None + } + return S3PathPairing.move_pairs(pairs, missing=missing) + + def _copy_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + isdir: Callable[[str], bool] | None = None, + ) -> list[tuple[str, str]]: + """Expand and pair the paths of a copy (see :meth:`S3PathPairing.copy_pairs`). + + The destination is looked up only when it decides the pairing. + + Args: + path1: Source S3 path, glob pattern, or list of them. + path2: Destination path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + isdir: Whether the destination is a directory, by default + ``self.isdir``; ``get()`` passes the local filesystem's. + + Returns: + The sources and their destinations. + """ + if not S3PathPairing.expands(path1, path2): + return S3PathPairing.copy_pairs(path1, path2) + sources = self.expand_path(path1, recursive=recursive, maxdepth=maxdepth) + if S3PathPairing.skips_directories(path1, recursive, maxdepth): + sources = [p for p in sources if not (trailing_sep(p) or self.isdir(p))] + destination_is_dir = ( + # A string, as looks_up_destination() checks. + (isdir or self.isdir)(cast(str, path2)) + if sources and S3PathPairing.looks_up_destination(path1, path2) + else None + ) + return S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) + def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kwargs) -> None: """Copy files within S3. @@ -1444,7 +1529,7 @@ def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kw """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = self.pairing.copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) if not pairs: return path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] @@ -1478,7 +1563,7 @@ def get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) 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 = self.pairing.copy_pairs( + pairs = self._copy_pairs( rpath, root, recursive=recursive, maxdepth=maxdepth, isdir=LocalFileSystem().isdir ) rpath, lpath = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 5a776395..00224aee 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -13,13 +13,14 @@ import logging import mimetypes import os +from collections.abc import Callable from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast from fsspec.asyn import AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK from fsspec.core import get_compression -from fsspec.implementations.local import LocalFileSystem, make_path_posix +from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep from fsspec.utils import check_contained from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem @@ -143,11 +144,6 @@ def core(self) -> S3Core: """The typed S3 operations of the wrapped ``S3FileSystem``.""" return self._sync_fs.core - @property - def pairing(self) -> S3PathPairing: - """The pairing of the paths of the wrapped ``S3FileSystem``.""" - return self._sync_fs.pairing - @staticmethod def parse_path(path: str) -> tuple[str, str | None, str | None]: """Parse an S3 path into its bucket, key and version ID. @@ -329,11 +325,33 @@ async def _rm( ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = await asyncio.to_thread( - self.pairing.delete_paths, path, recursive=recursive, maxdepth=maxdepth - ) + paths = await self._delete_paths(path, recursive=recursive, maxdepth=maxdepth) await self._delete_objects(paths, **kwargs) + async def _delete_paths( + self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None + ) -> list[str]: + """Expand the paths that ``_rm()`` deletes, as ``S3FileSystem._delete_paths`` does. + + Args: + path: S3 path or list of paths. + recursive: Whether to include all objects below the paths. + maxdepth: Maximum depth to expand when ``recursive`` is True. + + Returns: + The paths with a version ID as given, followed by the expansion + of the other paths by ``_expand_path``. + + Raises: + ValueError: If a path is a bucket. + """ + versioned_paths, unversioned_paths = S3PathPairing.delete_paths(path) + if not unversioned_paths: + return versioned_paths + return versioned_paths + await self._expand_path( + unversioned_paths, recursive=recursive, maxdepth=maxdepth + ) + async def _delete_objects(self, paths: list[str], **kwargs) -> None: """Delete objects with DeleteObjects requests run with ``asyncio.gather``. @@ -377,9 +395,7 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N """ if path1 == path2: return - pairs = await asyncio.to_thread( - self.pairing.move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth - ) + pairs = await self._move_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) # Every copy finishes before a failure is raised, as in fsspec's # _copy(), and nothing is deleted after a failure. results = await asyncio.gather( @@ -392,6 +408,69 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N [p1 for (p1, _), copied in zip(pairs, results, strict=True) if copied] ) + async def _move_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + ) -> list[tuple[str, str]]: + """Pair and check the paths of a move, as ``S3FileSystem._move_pairs`` does. + + Args: + path1: Source S3 path, glob pattern, or list of paths. + path2: Destination S3 path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + + Returns: + The sources and destinations that are moved. + + Raises: + ValueError: If the move has conflicting paths. + """ + pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + missing = set() + for source in S3PathPairing.conflict_candidates(pairs): + if await asyncio.to_thread(self._sync_fs._head_object, source) is None: + missing.add(source) + return S3PathPairing.move_pairs(pairs, missing=missing) + + async def _copy_pairs( + self, + path1: str | list[str], + path2: str | list[str], + recursive: bool = False, + maxdepth: int | None = None, + isdir: Callable[[str], bool] | None = None, + ) -> list[tuple[str, str]]: + """Expand and pair the paths of a copy, as ``S3FileSystem._copy_pairs`` does. + + Args: + path1: Source S3 path, glob pattern, or list of them. + path2: Destination path, or list of paths when ``path1`` is a + list. + recursive: Whether to include the contents of the directories. + maxdepth: Maximum depth of the expansion. + isdir: Whether the destination is a directory, by default + ``self._isdir``; ``_get()`` passes the local filesystem's. + + Returns: + The sources and their destinations. + """ + if not S3PathPairing.expands(path1, path2): + return S3PathPairing.copy_pairs(path1, path2) + sources = await self._expand_path(path1, recursive=recursive, maxdepth=maxdepth) + if S3PathPairing.skips_directories(path1, recursive, maxdepth): + sources = [p for p in sources if not (trailing_sep(p) or await self._isdir(p))] + destination_is_dir = None + if sources and S3PathPairing.looks_up_destination(path1, path2): + # A string, as looks_up_destination() checks. + destination = cast(str, path2) + destination_is_dir = isdir(destination) if isdir else await self._isdir(destination) + return S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) + def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. @@ -435,9 +514,7 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = await asyncio.to_thread( - self.pairing.copy_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth - ) + pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) if not pairs: return path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] @@ -475,13 +552,8 @@ async def _get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) 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 asyncio.to_thread( - self.pairing.copy_pairs, - rpath, - root, - recursive=recursive, - maxdepth=maxdepth, - isdir=LocalFileSystem().isdir, + pairs = await self._copy_pairs( + 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) diff --git a/pyathena/filesystem/s3_path.py b/pyathena/filesystem/s3_path.py index 427f495f..1620bbca 100644 --- a/pyathena/filesystem/s3_path.py +++ b/pyathena/filesystem/s3_path.py @@ -137,9 +137,11 @@ def uri(self) -> str: def target(self) -> S3Path: """The object that a write to the path replaces, for comparing paths. - A write to a key replaces its ``null`` version, which the objects of a - bucket without versioning have, so a ``null`` version names the key - itself. Any other path is its own target. + A ``null`` version names the key itself: in a bucket without + versioning, or with versioning suspended, a write to the key replaces + its ``null`` version. With versioning enabled, a write adds a new + version instead, but the ``null`` version is still taken to name the + key. Any other path is its own target. """ return self.with_version_id(None) if self.version_id == "null" else self diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index bca0eed8..6ed46c0b 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -10,161 +10,199 @@ from __future__ import annotations from collections import Counter -from collections.abc import Callable +from collections.abc import Collection, Sequence from glob import has_magic -from typing import TYPE_CHECKING from fsspec.implementations.local import trailing_sep from fsspec.utils import other_paths from pyathena.filesystem.s3_path import S3Path -if TYPE_CHECKING: - from pyathena.filesystem.s3 import S3FileSystem - class S3PathPairing: - """The paths that ``copy()``, ``get()``, ``mv()`` and ``rm()`` operate on. + """The rules that pair the paths of ``copy()``, ``get()``, ``mv()`` and ``rm()``. - Expands the paths, glob patterns and directories as fsspec's ``copy()`` - and ``rm()`` do, and pairs the sources with their destinations, except - that a path with a version ID names that version of an object: it is not - a glob pattern, nothing is expanded below it, and its destination is named - after its key without the version. The paths are compared by what they - name (see :attr:`S3Path.target`). + The sources are paired with their destinations as fsspec's ``copy()`` + pairs them, except that a path with a version ID names that version of an + object: it is not a glob pattern, and its destination is named after its + key without the version. A move compares the paths by what they name + (see :attr:`S3Path.target`). - ``S3FileSystem`` builds one and exposes it as ``S3FileSystem.pairing``; - ``AioS3FileSystem.pairing`` is the same object. The lookups go through the - filesystem, with its cache. + The rules are pure functions of the paths and of the lookups that they + need, which the caller makes: the filesystems expand the paths and look + them up with their own requests and cache, and pass the results in. A + lookup that a rule needs and that is not passed raises ``ValueError``. Example: - >>> for source, destination in fs.pairing.copy_pairs( - ... "s3://bucket/dir/", "s3://bucket/copy/", recursive=True - ... ): - ... print(source, "->", destination) + >>> sources = fs.expand_path("s3://bucket/dir", recursive=True) + >>> S3PathPairing.copy_pairs("s3://bucket/dir", "s3://bucket/copy/", sources) """ - def __init__(self, fs: S3FileSystem) -> None: - """Create the pairing of a filesystem. + @staticmethod + def expands(path1: str | list[str], path2: str | list[str]) -> bool: + """Return whether a copy expands its sources. + + Args: + path1: Source path, glob pattern, or list of them. + path2: Destination path, or list of paths. + + Returns: + False if both are lists, which are paired as given, without + expansion or lookups. + """ + return not (isinstance(path1, list) and isinstance(path2, list)) + + @staticmethod + def skips_directories(path1: str | list[str], recursive: bool, maxdepth: int | None) -> bool: + """Return whether a copy leaves out the directories among its expanded sources. + + A directory is a source that ends with a slash or that the filesystem + reports as a directory. + + Args: + path1: Source path, glob pattern, or list of them. + recursive: Whether the copy includes the contents of directories. + maxdepth: Maximum depth of the expansion. + + Returns: + True for a string source copied without ``recursive``, or with a + ``maxdepth``. + """ + return isinstance(path1, str) and (not recursive or maxdepth is not None) + + @staticmethod + def looks_up_destination(path1: str | list[str], path2: str | list[str]) -> bool: + """Return whether the pairing of a copy depends on its destination being a directory. Args: - fs: The filesystem that expands and looks up the paths. + path1: Source path, glob pattern, or list of them. + path2: Destination path, or list of paths. + + Returns: + True for a string source that is neither a glob pattern nor ends + with a slash, copied to a string destination that does not end + with a slash. :meth:`copy_pairs` then needs + ``destination_is_dir``. """ - self._fs = fs + return ( + isinstance(path1, str) + and not S3PathPairing._is_glob(path1) + and not trailing_sep(path1) + and isinstance(path2, str) + and not trailing_sep(path2) + ) + @staticmethod def copy_pairs( - self, path1: str | list[str], path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, - isdir: Callable[[str], bool] | None = None, + sources: Sequence[str] = (), + destination_is_dir: bool | None = None, ) -> list[tuple[str, str]]: """Pair the sources of a copy with their destinations as fsspec's ``copy()`` does. Args: - path1: Source S3 path, glob pattern, or list of them. - path2: Destination path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. - isdir: Whether a destination path is a directory, by default the - filesystem's ``isdir``; ``get()`` passes the local - filesystem's. + path1: Source S3 path, glob pattern, or list of them, as given to + the copy. + path2: Destination path, or a list of paths: as many as the + sources, or, when ``path1`` is a list, its destinations. + sources: The expansion of ``path1``, without the directories that + :meth:`skips_directories` leaves out. Not used when both + ``path1`` and ``path2`` are lists. + destination_is_dir: Whether ``path2`` is a directory; needed when + :meth:`looks_up_destination` is true and there are sources. Returns: - The sources and their destinations. When both ``path1`` and - ``path2`` are lists, they are paired as given. Empty if ``path1`` - is a string that matches only directories without ``recursive``. + The sources and their destinations, which keep the form of + ``path2``. When both ``path1`` and ``path2`` are lists, they are + paired as given, up to the end of the shorter list, as in + fsspec's ``copy()``. Empty if there are no sources. Raises: - ValueError: If ``maxdepth`` is less than 1. - FileNotFoundError: If the expansion of ``path1`` matches nothing: - a glob pattern without matches, or a missing path with - ``recursive``. + ValueError: If ``destination_is_dir`` is needed and None. """ - if isinstance(path1, list) and isinstance(path2, list): + if not S3PathPairing.expands(path1, path2): return list(zip(path1, path2, strict=False)) - fs = self._fs + if not sources: + return [] + if destination_is_dir is None and S3PathPairing.looks_up_destination(path1, path2): + raise ValueError("destination_is_dir is needed to pair the paths.") source_is_str = isinstance(path1, str) - paths1 = fs.expand_path(path1, recursive=recursive, maxdepth=maxdepth) - if source_is_str and (not recursive or maxdepth is not None): - # Non-recursive glob does not copy directories. - paths1 = [p for p in paths1 if not (trailing_sep(p) or fs.isdir(p))] - if not paths1: - return [] - glob = isinstance(path1, str) and has_magic(path1) and not S3Path.has_version_id(path1) - # The destination is looked up only when it decides the mapping. + glob = isinstance(path1, str) and S3PathPairing._is_glob(path1) exists = source_is_str and ( - (glob and len(paths1) == 1) + (glob and len(sources) == 1) or ( not glob and not trailing_sep(path1) and isinstance(path2, str) - and (trailing_sep(path2) or (isdir or fs.isdir)(path2)) + and (trailing_sep(path2) or bool(destination_is_dir)) ) ) - names = [S3Path.split_version_id(p)[0] for p in paths1] - paths2 = other_paths(names, path2, exists=exists, flatten=not source_is_str) - return list(zip(paths1, paths2, strict=True)) + names = [S3Path.split_version_id(p)[0] for p in sources] + destinations = other_paths(names, path2, exists=exists, flatten=not source_is_str) + return list(zip(sources, destinations, strict=True)) + + @staticmethod + def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: + """Return the sources of a move whose conflicts depend on an object at their key. + + A source with another source below it may be a directory. If no object + exists at its key, it is not copied, so it writes no destination and + is left out of the conflict checks of :meth:`move_pairs`. A path with + a version always names an object. + + Args: + pairs: The sources and destinations of the move, as + :meth:`copy_pairs` pairs them. + Returns: + The sources, by their :attr:`S3Path.target`, that have another + source below them and a destination that conflicts, in the order + of the pairs; empty if nothing needs to be looked up. + """ + moved, sources, directories, counts = S3PathPairing._moves(pairs) + return list( + dict.fromkeys( + source + for versioned, source, dest in moved + if (counts[dest] > 1 or dest in sources) and source in directories and not versioned + ) + ) + + @staticmethod def move_pairs( - self, - path1: str | list[str], - path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, + pairs: Sequence[tuple[str, str]], missing: Collection[str] | None = None ) -> list[tuple[str, str]]: - """Pair the sources of a move with their destinations as :meth:`copy_pairs` does. + """Check the pairs of a move and leave out the sources that stay in place. Args: - path1: Source S3 path, glob pattern, or list of paths. - path2: Destination S3 path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. + pairs: The sources and destinations of the move, as + :meth:`copy_pairs` pairs them. + missing: The :meth:`conflict_candidates` without an object at their + key; needed when there are candidates. Returns: - The sources and their destinations, except the sources whose - destination is the source itself or, for a ``null`` version, the - key of the source. + The pairs, except those whose destination is the source itself or, + for a ``null`` version, the key of the source. Raises: ValueError: If two sources have the same destination, or a destination is another source, including one left in place, except for a directory with no object at its key, which is not - copied. Also if ``maxdepth`` is less than 1. - FileNotFoundError: If the expansion of ``path1`` matches nothing, - as for :meth:`copy_pairs`. + copied. Also if ``missing`` is needed and None. """ - # The paths are moved as given, and compared by what they name. - named = [ - (p1, p2, str(S3Path.parse(p1).target), str(S3Path.parse(p2).target)) - for p1, p2 in self.copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) - ] - pairs = [(p1, p2) for p1, p2, source, dest in named if source != dest] - moved = [(p1, source, dest) for p1, _, source, dest in named if source != dest] - # The sources left in place count too; a copy onto one overwrites it. - sources = {source for _, _, source, _ in named} - counts = Counter(dest for _, _, dest in moved) - # A source with another source below it may be a directory. - directories: set[str] = set() - for source in sources: - parent = source.rpartition("/")[0] - while parent and parent not in directories: - directories.add(parent) - parent = parent.rpartition("/")[0] - # A directory without an object at its key is not copied, so it - # writes no destination and is left out of the checks. A path with a - # version always names an object. + moved, sources, directories, counts = S3PathPairing._moves(pairs) + if missing is None and S3PathPairing.conflict_candidates(pairs): + raise ValueError("missing is needed to check the pairs.") + missing = missing or () writers = [ (source, dest) - for p1, source, dest in moved + for versioned, source, dest in moved if not ( (counts[dest] > 1 or dest in sources) and source in directories - and not S3Path.parse(p1).version_id - and self._fs._head_object(source) is None + and not versioned + and source in missing ) ] counts = Counter(dest for _, dest in writers) @@ -173,26 +211,25 @@ def move_pairs( raise ValueError("Cannot move several paths to the same destination.") if dest in sources: raise ValueError("Cannot move a path onto another path that is moved.") - return pairs + return [ + (p1, p2) + for p1, p2 in pairs + if str(S3Path.parse(p1).target) != str(S3Path.parse(p2).target) + ] - def delete_paths( - self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None - ) -> list[str]: - """Expand the paths that ``rm()`` deletes. + @staticmethod + def delete_paths(path: str | list[str]) -> tuple[list[str], list[str]]: + """Split the paths that ``rm()`` deletes into those deleted as given and those expanded. Args: path: S3 path or list of paths. - recursive: Whether to include all objects below the paths. - maxdepth: Maximum depth to expand when ``recursive`` is True. Returns: - The paths with a version ID as given, followed by the expansion - of the other paths by the filesystem's ``expand_path``. + The paths with a version ID, which are deleted as given, and the + other paths, which the filesystem expands. Raises: - ValueError: If a path is a bucket, or ``maxdepth`` is less than 1. - FileNotFoundError: If the expansion of the paths without a version - ID matches nothing, as for :meth:`copy_pairs`. + ValueError: If a path is a bucket. """ paths = [path] if isinstance(path, str) else list(path) versioned_paths, unversioned_paths = [], [] @@ -205,11 +242,53 @@ def delete_paths( versioned_paths.append(p) else: unversioned_paths.append(p) + return versioned_paths, unversioned_paths + + @staticmethod + def _is_glob(path: str) -> bool: + """Return whether a path is a glob pattern; a version ID query is not one. + + Args: + path: The path. + + Returns: + True if the path has glob characters and no version ID query. + """ + return has_magic(path) and not S3Path.has_version_id(path) + + @staticmethod + def _moves( + pairs: Sequence[tuple[str, str]], + ) -> tuple[list[tuple[bool, str, str]], set[str], set[str], Counter[str]]: + """Compare the paths of a move by what they name. + + Args: + pairs: The sources and destinations of the move. - if unversioned_paths: - # Versioned paths are deleted as given, without the lookup that - # expand_path makes for them with recursive. - unversioned_paths = self._fs.expand_path( - unversioned_paths, recursive=recursive, maxdepth=maxdepth + Returns: + The moved sources (whether the source has a version, its target + and the target of its destination), the targets of all sources, + including those left in place, the parents of the sources, and + the number of moved sources per destination target. + """ + named = [] + for p1, p2 in pairs: + source_path = S3Path.parse(p1) + named.append( + ( + bool(source_path.version_id), + str(source_path.target), + str(S3Path.parse(p2).target), + ) ) - return versioned_paths + unversioned_paths + moved = [(versioned, source, dest) for versioned, source, dest in named if source != dest] + # The sources left in place count too; a copy onto one overwrites it. + sources = {source for _, source, _ in named} + # A source with another source below it may be a directory. + directories: set[str] = set() + for source in sources: + parent = source.rpartition("/")[0] + while parent and parent not in directories: + directories.add(parent) + parent = parent.rpartition("/")[0] + return moved, sources, directories, Counter(dest for _, _, dest in moved) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index c50b1d7f..34b50ca6 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -16,6 +16,7 @@ import urllib.parse import urllib.request import uuid +import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain @@ -41,7 +42,6 @@ from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path -from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -175,7 +175,6 @@ def _make_fs(): fs._core = S3Core(client, retry_config=RetryConfig()) # The requests of the core and of the filesystem go to one mock. fs._call = fs._core.call = mock.MagicMock() - fs._pairing = S3PathPairing(fs) fs.max_workers = 4 fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE fs.allow_bucket_creation = False @@ -1810,6 +1809,57 @@ def test_get_version_outside_destination(self, tmp_path, rpath): fs.get(rpath, f"{tmp_path}/d/") assert sorted(p.name for p in tmp_path.rglob("*")) == ["d"] + @pytest.mark.parametrize( + ("path2", "lookups", "expected"), + [ + # The source is checked for a directory; the destination decides + # the pairing, so it is looked up too. + ("s3://bucket/d", ["bucket/b?versionId=v1", "s3://bucket/d"], "s3://bucket/d/b"), + # A trailing slash decides it without a lookup. + ("s3://bucket/d/", ["bucket/b?versionId=v1"], "s3://bucket/d/b"), + ], + ) + def test_copy_pairs_destination_lookup(self, path2, lookups, expected): + fs = self._make_fs() + self._serve_keys(fs, {"b"}) + # Only the destination is a directory. + fs.isdir = mock.MagicMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) + + pairs = fs._copy_pairs("s3://bucket/b?versionId=v1", path2) + + assert pairs == [("bucket/b?versionId=v1", expected)] + assert [c.args[0] for c in fs.isdir.call_args_list] == lookups + + def test_move_pairs_looks_up_only_conflict_candidates(self): + # Only a source that may be a directory and whose destination + # conflicts is looked up with HeadObject. + fs = self._make_fs() + self._serve_keys(fs, {"d/x", "e/y", "f"}) + fs._head_object = mock.MagicMock(return_value=None) + + pairs = fs._move_pairs( + ["s3://bucket/d", "s3://bucket/d/x", "s3://bucket/e/y", "s3://bucket/f"], + ["s3://bucket/e", "s3://bucket/e", "s3://bucket/out", "s3://bucket/g"], + ) + + assert len(pairs) == 4 + fs._head_object.assert_called_once_with("bucket/d") + + def test_freed_by_reference_counting(self): + # The filesystem holds no reference cycle, so a filesystem that is + # not cached, such as the internal one of a cursor (GH-978), is freed + # as soon as it is unused. + fs = S3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + ref = weakref.ref(fs) + gc.disable() + try: + del fs + assert ref() is None + finally: + gc.enable() + def test_mv_nothing_within_maxdepth(self): # Only directories within maxdepth: nothing is moved, as with copy(). fs = self._make_fs() diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 9663602c..fa43bb95 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -869,12 +869,13 @@ async def test_rm_maxdepth(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) - sync_fs.find = mock.MagicMock(return_value=["bucket/dir/a"]) - sync_fs.exists = mock.MagicMock(return_value=True) + # The paths are expanded with the coroutines of fsspec's _expand_path(). + fs._find = mock.AsyncMock(return_value=["bucket/dir/a"]) + fs._exists = mock.AsyncMock(return_value=True) # batch_size is part of fsspec's async _rm() signature. await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) - sync_fs.find.assert_called_once_with("bucket/dir", maxdepth=1, withdirs=True, detail=False) + fs._find.assert_awaited_once_with("bucket/dir", maxdepth=1, withdirs=True) (call,) = sync_fs._call.call_args_list assert call.kwargs["Delete"]["Objects"] == [{"Key": "dir"}, {"Key": "dir/a"}] @@ -1105,6 +1106,40 @@ def upload_part_copy(**kw): == "requester" ) + @pytest.mark.parametrize( + ("path2", "lookups", "expected"), + [ + ("s3://bucket/d", ["bucket/b?versionId=v1", "s3://bucket/d"], "s3://bucket/d/b"), + ("s3://bucket/d/", ["bucket/b?versionId=v1"], "s3://bucket/d/b"), + ], + ) + @pytest.mark.asyncio + async def test_copy_pairs_destination_lookup(self, path2, lookups, expected): + # See TestS3FileSystem.test_copy_pairs_destination_lookup; the aio + # filesystem expands and looks up the paths with its own coroutines. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + # Only the destination is a directory. + fs._isdir = mock.AsyncMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) + + pairs = await fs._copy_pairs("s3://bucket/b?versionId=v1", path2) + + assert pairs == [("bucket/b?versionId=v1", expected)] + assert [c.args[0] for c in fs._isdir.call_args_list] == lookups + + @pytest.mark.asyncio + async def test_move_pairs_looks_up_only_conflict_candidates(self): + # See TestS3FileSystem.test_move_pairs_looks_up_only_conflict_candidates. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs._sync_fs._head_object = mock.MagicMock(return_value=None) + + pairs = await fs._move_pairs( + ["s3://bucket/d", "s3://bucket/d/x", "s3://bucket/e/y", "s3://bucket/f"], + ["s3://bucket/e", "s3://bucket/e", "s3://bucket/out", "s3://bucket/g"], + ) + + assert len(pairs) == 4 + fs._sync_fs._head_object.assert_called_once_with("bucket/d") + def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance # cache, so skip_instance_cache=True instances shared it. diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 4a16fee6..3323999f 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -5,113 +5,173 @@ # # SPDX-License-Identifier: MIT -from unittest import mock - import pytest -from pyathena.filesystem.s3_async import AioS3FileSystem from pyathena.filesystem.s3_path_pairing import S3PathPairing -from tests.pyathena.filesystem import test_s3 -def _make_fs(keys): - # An S3FileSystem whose requests are answered from the keys of "bucket". - fs = test_s3.TestS3FileSystem._make_fs() - test_s3.TestS3FileSystem._serve_keys(fs, keys) - return fs +class TestS3PathPairing: + @pytest.mark.parametrize( + ("path1", "path2", "expected"), + [ + ("s3://bucket/a", "s3://bucket/b", True), + (["s3://bucket/a"], "s3://bucket/b", True), + ("s3://bucket/a", ["s3://bucket/b"], True), + # Lists are paired as given. + (["s3://bucket/a"], ["s3://bucket/b"], False), + ], + ) + def test_expands(self, path1, path2, expected): + assert S3PathPairing.expands(path1, path2) is expected + @pytest.mark.parametrize( + ("path1", "recursive", "maxdepth", "expected"), + [ + ("s3://bucket/d", False, None, True), + ("s3://bucket/d", True, None, False), + ("s3://bucket/d", True, 1, True), + (["s3://bucket/d"], False, None, False), + ], + ) + def test_skips_directories(self, path1, recursive, maxdepth, expected): + assert S3PathPairing.skips_directories(path1, recursive, maxdepth) is expected -class TestS3PathPairing: - def test_filesystems_share_one_pairing(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - assert isinstance(fs.pairing, S3PathPairing) - assert fs.pairing is fs._sync_fs.pairing - assert fs._sync_fs.pairing is fs._sync_fs.pairing + @pytest.mark.parametrize( + ("path1", "path2", "expected"), + [ + ("s3://bucket/a", "s3://bucket/b", True), + # A version ID query is not a glob pattern. + ("s3://bucket/a?versionId=v1", "s3://bucket/b", True), + # A trailing slash decides the pairing without a lookup. + ("s3://bucket/a/", "s3://bucket/b", False), + ("s3://bucket/a", "s3://bucket/b/", False), + # So does a glob pattern. + ("s3://bucket/a*", "s3://bucket/b", False), + (["s3://bucket/a"], "s3://bucket/b", False), + ("s3://bucket/a", ["s3://bucket/b"], False), + ], + ) + def test_looks_up_destination(self, path1, path2, expected): + assert S3PathPairing.looks_up_destination(path1, path2) is expected @pytest.mark.parametrize( - ("path1", "path2", "kwargs", "expected"), + ("path1", "path2", "sources", "destination_is_dir", "expected"), [ - # The directory and the objects below it, as fsspec's copy() pairs - # them. + # The directory and the objects below it, as fsspec's copy() + # pairs them. ( "s3://bucket/d", "s3://bucket/out", - {"recursive": True}, + ["bucket/d", "bucket/d/a", "bucket/d/b"], + False, [ ("bucket/d", "s3://bucket/out"), ("bucket/d/a", "s3://bucket/out/a"), ("bucket/d/b", "s3://bucket/out/b"), ], ), - # A version is copied to a destination named after its key. + # A version is copied into a directory under its key, + ( + "s3://bucket/b?versionId=v1", + "s3://bucket/d", + ["bucket/b?versionId=v1"], + True, + [("bucket/b?versionId=v1", "s3://bucket/d/b")], + ), + # or to the destination itself. + ( + "s3://bucket/b?versionId=v1", + "s3://bucket/c", + ["bucket/b?versionId=v1"], + False, + [("bucket/b?versionId=v1", "s3://bucket/c")], + ), + # A trailing slash needs no lookup. ( "s3://bucket/b?versionId=v1", "s3://bucket/d/", - {}, + ["bucket/b?versionId=v1"], + None, [("bucket/b?versionId=v1", "s3://bucket/d/b")], ), - # Without recursive, a directory is not copied. - ("s3://bucket/d", "s3://bucket/out", {}, []), - # Lists are paired as given. + # Lists are paired as given, up to the end of the shorter one. ( - ["s3://bucket/b", "s3://bucket/d/a"], + ["s3://bucket/a", "s3://bucket/b", "s3://bucket/c"], ["s3://bucket/x", "s3://bucket/y"], - {}, - [("s3://bucket/b", "s3://bucket/x"), ("s3://bucket/d/a", "s3://bucket/y")], + (), + None, + [("s3://bucket/a", "s3://bucket/x"), ("s3://bucket/b", "s3://bucket/y")], ), + # Nothing to copy. + ("s3://bucket/d", "s3://bucket/out", [], None, []), ], ) - def test_copy_pairs(self, path1, path2, kwargs, expected): - fs = _make_fs({"d/a", "d/b", "b"}) - assert fs.pairing.copy_pairs(path1, path2, **kwargs) == expected - - def test_copy_pairs_isdir(self): - # The given isdir, such as the local filesystem's for get(), decides - # whether the destination is a directory. - fs = _make_fs({"b"}) - isdir = mock.MagicMock(return_value=True) + def test_copy_pairs(self, path1, path2, sources, destination_is_dir, expected): + assert S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) == expected - pairs = fs.pairing.copy_pairs("s3://bucket/b?versionId=v1", "/tmp/out", isdir=isdir) - - assert pairs == [("bucket/b?versionId=v1", "/tmp/out/b")] - isdir.assert_called_once_with("/tmp/out") + def test_copy_pairs_needs_destination_lookup(self): + with pytest.raises(ValueError, match="destination_is_dir"): + S3PathPairing.copy_pairs("s3://bucket/a", "s3://bucket/b", ["bucket/a"]) def test_move_pairs(self): # The "null" version of a key moved onto the key stays in place. - fs = _make_fs({"b", "d/a"}) - - pairs = fs.pairing.move_pairs( - ["s3://bucket/b?versionId=null", "s3://bucket/d/a"], - ["s3://bucket/b", "s3://bucket/z"], - ) + pairs = [("s3://bucket/b?versionId=null", "s3://bucket/b"), ("bucket/d/a", "s3://bucket/z")] - assert pairs == [("s3://bucket/d/a", "s3://bucket/z")] + assert S3PathPairing.conflict_candidates(pairs) == [] + assert S3PathPairing.move_pairs(pairs) == [("bucket/d/a", "s3://bucket/z")] @pytest.mark.parametrize( - ("path2", "match"), + ("pairs", "match"), [ - (["s3://bucket/o", "s3://bucket/o"], "same destination"), - (["s3://bucket/z", "s3://bucket/n"], "another path that is moved"), + ( + [("s3://bucket/b", "s3://bucket/o"), ("s3://bucket/z", "s3://bucket/o")], + "same destination", + ), + ( + [("s3://bucket/b", "s3://bucket/z"), ("s3://bucket/z", "s3://bucket/n")], + "another path that is moved", + ), ], ) - def test_move_pairs_conflicts(self, path2, match): - fs = _make_fs({"b", "z"}) + def test_move_pairs_conflicts(self, pairs, match): with pytest.raises(ValueError, match=match): - fs.pairing.move_pairs(["s3://bucket/b", "s3://bucket/z"], path2) + S3PathPairing.move_pairs(pairs) + + def test_move_pairs_directory_without_object(self): + # A source with another source below it may be a directory; one + # without an object at its key is not copied, so its destination + # does not conflict. + pairs = [ + ("s3://bucket/d", "s3://bucket/e"), + ("s3://bucket/d/x", "s3://bucket/e"), + ("s3://bucket/e/y", "s3://bucket/out"), + ] + + assert S3PathPairing.conflict_candidates(pairs) == ["bucket/d"] + with pytest.raises(ValueError, match="missing"): + S3PathPairing.move_pairs(pairs) + with pytest.raises(ValueError, match="same destination"): + S3PathPairing.move_pairs(pairs, missing=set()) + assert S3PathPairing.move_pairs(pairs, missing={"bucket/d"}) == pairs + + def test_move_pairs_version_names_an_object(self): + # A version is never taken for a directory. + pairs = [ + ("s3://bucket/d?versionId=null", "s3://bucket/out"), + ("s3://bucket/d/x", "s3://bucket/x"), + ("s3://bucket/a", "s3://bucket/out"), + ] + + assert S3PathPairing.conflict_candidates(pairs) == [] + with pytest.raises(ValueError, match="same destination"): + S3PathPairing.move_pairs(pairs) def test_delete_paths(self): - # The versions are deleted as given, before the expanded paths. - fs = _make_fs({"d/a", "d/b"}) - - paths = fs.pairing.delete_paths( - ["s3://bucket/d", "s3://bucket/b?versionId=v1"], recursive=True - ) - - assert paths == ["s3://bucket/b?versionId=v1", "bucket/d", "bucket/d/a", "bucket/d/b"] + assert S3PathPairing.delete_paths( + ["s3://bucket/d", "s3://bucket/b?versionId=v1", "s3://bucket/c"] + ) == (["s3://bucket/b?versionId=v1"], ["s3://bucket/d", "s3://bucket/c"]) @pytest.mark.parametrize("path", ["s3://bucket", "s3://bucket/", "s3://bucket//"]) def test_delete_paths_bucket(self, path): - fs = _make_fs(set()) with pytest.raises(ValueError, match="Cannot delete the bucket"): - fs.pairing.delete_paths(path) - fs._call.assert_not_called() + S3PathPairing.delete_paths(path) From 05f37dbec72e5ce69ec6390cbd4274a543b88764 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:13:28 +0900 Subject: [PATCH 03/10] Require the sources, normalize missing, parse each move path once, and run the aio lookups concurrently Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 7 +- pyathena/filesystem/s3_async.py | 23 +++- pyathena/filesystem/s3_path_pairing.py | 113 +++++++++++------- .../filesystem/test_s3_path_pairing.py | 7 ++ 4 files changed, 96 insertions(+), 54 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 59ff6b5e..4f5184a9 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1514,7 +1514,8 @@ def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kw As fsspec's ``copy()``, except that a source with a version ID copies that version of the object to a destination named after its key, as - :meth:`S3PathPairing.copy_pairs` pairs them. + :meth:`~pyathena.filesystem.s3_path_pairing.S3PathPairing.copy_pairs` + pairs them. Args: path1: Source S3 path, glob pattern, or list of them. @@ -1544,8 +1545,8 @@ def get( As fsspec's ``get()``, except that a source with a version ID downloads that version of the object to a local path named after its - key, as :meth:`S3PathPairing.copy_pairs` pairs them. Those destinations are checked to - lie under ``lpath``. + key, as :meth:`~pyathena.filesystem.s3_path_pairing.S3PathPairing.copy_pairs` + pairs them. Those destinations are checked to lie under ``lpath``. Args: rpath: Source S3 path, glob pattern, or list of them. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 00224aee..fa3ad97e 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -431,10 +431,13 @@ async def _move_pairs( ValueError: If the move has conflicting paths. """ pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) - missing = set() - for source in S3PathPairing.conflict_candidates(pairs): - if await asyncio.to_thread(self._sync_fs._head_object, source) is None: - missing.add(source) + candidates = S3PathPairing.conflict_candidates(pairs) + objects = await asyncio.gather( + *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates] + ) + missing = { + source for source, object_ in zip(candidates, objects, strict=True) if not object_ + } return S3PathPairing.move_pairs(pairs, missing=missing) async def _copy_pairs( @@ -463,12 +466,20 @@ async def _copy_pairs( return S3PathPairing.copy_pairs(path1, path2) sources = await self._expand_path(path1, recursive=recursive, maxdepth=maxdepth) if S3PathPairing.skips_directories(path1, recursive, maxdepth): - sources = [p for p in sources if not (trailing_sep(p) or await self._isdir(p))] + # A path with a trailing slash is a directory without a lookup. + files = [p for p in sources if not trailing_sep(p)] + directories = await asyncio.gather(*[self._isdir(p) for p in files]) + sources = [p for p, is_dir in zip(files, directories, strict=True) if not is_dir] destination_is_dir = None if sources and S3PathPairing.looks_up_destination(path1, path2): # A string, as looks_up_destination() checks. destination = cast(str, path2) - destination_is_dir = isdir(destination) if isdir else await self._isdir(destination) + destination_is_dir = ( + # A local isdir, which can block, runs in a thread. + await asyncio.to_thread(isdir, destination) + if isdir + else await self._isdir(destination) + ) return S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 6ed46c0b..24a32a6b 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -26,7 +26,7 @@ class S3PathPairing: pairs them, except that a path with a version ID names that version of an object: it is not a glob pattern, and its destination is named after its key without the version. A move compares the paths by what they name - (see :attr:`S3Path.target`). + (see :attr:`~pyathena.filesystem.s3_path.S3Path.target`). The rules are pure functions of the paths and of the lookups that they need, which the caller makes: the filesystems expand the paths and look @@ -96,7 +96,7 @@ def looks_up_destination(path1: str | list[str], path2: str | list[str]) -> bool def copy_pairs( path1: str | list[str], path2: str | list[str], - sources: Sequence[str] = (), + sources: Sequence[str] | None = None, destination_is_dir: bool | None = None, ) -> list[tuple[str, str]]: """Pair the sources of a copy with their destinations as fsspec's ``copy()`` does. @@ -107,8 +107,9 @@ def copy_pairs( path2: Destination path, or a list of paths: as many as the sources, or, when ``path1`` is a list, its destinations. sources: The expansion of ``path1``, without the directories that - :meth:`skips_directories` leaves out. Not used when both - ``path1`` and ``path2`` are lists. + :meth:`skips_directories` leaves out; needed unless both + ``path1`` and ``path2`` are lists, which are paired without + it. destination_is_dir: Whether ``path2`` is a directory; needed when :meth:`looks_up_destination` is true and there are sources. @@ -119,10 +120,13 @@ def copy_pairs( fsspec's ``copy()``. Empty if there are no sources. Raises: - ValueError: If ``destination_is_dir`` is needed and None. + ValueError: If ``sources`` or ``destination_is_dir`` is needed and + None. """ if not S3PathPairing.expands(path1, path2): return list(zip(path1, path2, strict=False)) + if sources is None: + raise ValueError("sources is needed to pair the paths.") if not sources: return [] if destination_is_dir is None and S3PathPairing.looks_up_destination(path1, path2): @@ -156,18 +160,12 @@ def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: :meth:`copy_pairs` pairs them. Returns: - The sources, by their :attr:`S3Path.target`, that have another - source below them and a destination that conflicts, in the order - of the pairs; empty if nothing needs to be looked up. + The sources, in ``bucket/key`` form (their + :attr:`~pyathena.filesystem.s3_path.S3Path.target`), that have + another source below them and a destination that conflicts, in the + order of the pairs; empty if nothing needs to be looked up. """ - moved, sources, directories, counts = S3PathPairing._moves(pairs) - return list( - dict.fromkeys( - source - for versioned, source, dest in moved - if (counts[dest] > 1 or dest in sources) and source in directories and not versioned - ) - ) + return S3PathPairing._candidates(*S3PathPairing._moves(pairs)) @staticmethod def move_pairs( @@ -179,7 +177,8 @@ def move_pairs( pairs: The sources and destinations of the move, as :meth:`copy_pairs` pairs them. missing: The :meth:`conflict_candidates` without an object at their - key; needed when there are candidates. + key, in any form that names them; needed when there are + candidates. Returns: The pairs, except those whose destination is the source itself or, @@ -191,31 +190,23 @@ def move_pairs( except for a directory with no object at its key, which is not copied. Also if ``missing`` is needed and None. """ - moved, sources, directories, counts = S3PathPairing._moves(pairs) - if missing is None and S3PathPairing.conflict_candidates(pairs): + named, sources, directories, counts = S3PathPairing._moves(pairs) + candidates = S3PathPairing._candidates(named, sources, directories, counts) + if missing is None and candidates: raise ValueError("missing is needed to check the pairs.") - missing = missing or () - writers = [ - (source, dest) - for versioned, source, dest in moved - if not ( - (counts[dest] > 1 or dest in sources) - and source in directories - and not versioned - and source in missing - ) - ] - counts = Counter(dest for _, dest in writers) - for _, dest in writers: - if counts[dest] > 1: + # A directory without an object at its key writes no destination. + skipped = {str(S3Path.parse(path).target) for path in missing or ()} + writers = Counter( + dest + for _, _, _, source, dest in named + if source != dest and not (source in candidates and source in skipped) + ) + for dest in writers: + if writers[dest] > 1: raise ValueError("Cannot move several paths to the same destination.") if dest in sources: raise ValueError("Cannot move a path onto another path that is moved.") - return [ - (p1, p2) - for p1, p2 in pairs - if str(S3Path.parse(p1).target) != str(S3Path.parse(p2).target) - ] + return [(p1, p2) for p1, p2, _, source, dest in named if source != dest] @staticmethod def delete_paths(path: str | list[str]) -> tuple[list[str], list[str]]: @@ -259,16 +250,16 @@ def _is_glob(path: str) -> bool: @staticmethod def _moves( pairs: Sequence[tuple[str, str]], - ) -> tuple[list[tuple[bool, str, str]], set[str], set[str], Counter[str]]: + ) -> tuple[list[tuple[str, str, bool, str, str]], set[str], set[str], Counter[str]]: """Compare the paths of a move by what they name. Args: pairs: The sources and destinations of the move. Returns: - The moved sources (whether the source has a version, its target - and the target of its destination), the targets of all sources, - including those left in place, the parents of the sources, and + Each pair with whether its source has a version and the targets + of its source and destination; the targets of all sources, + including those left in place; the parents of the sources; and the number of moved sources per destination target. """ named = [] @@ -276,14 +267,15 @@ def _moves( source_path = S3Path.parse(p1) named.append( ( + p1, + p2, bool(source_path.version_id), str(source_path.target), str(S3Path.parse(p2).target), ) ) - moved = [(versioned, source, dest) for versioned, source, dest in named if source != dest] # The sources left in place count too; a copy onto one overwrites it. - sources = {source for _, source, _ in named} + sources = {source for _, _, _, source, _ in named} # A source with another source below it may be a directory. directories: set[str] = set() for source in sources: @@ -291,4 +283,35 @@ def _moves( while parent and parent not in directories: directories.add(parent) parent = parent.rpartition("/")[0] - return moved, sources, directories, Counter(dest for _, _, dest in moved) + counts = Counter(dest for _, _, _, source, dest in named if source != dest) + return named, sources, directories, counts + + @staticmethod + def _candidates( + named: list[tuple[str, str, bool, str, str]], + sources: set[str], + directories: set[str], + counts: Counter[str], + ) -> list[str]: + """Select the conflict candidates of a move compared by :meth:`_moves`. + + Args: + named: The pairs with their versions and targets. + sources: The targets of all sources. + directories: The parents of the sources. + counts: The number of moved sources per destination target. + + Returns: + The moved, unversioned sources that are a parent of another + source and whose destination conflicts, without duplicates. + """ + return list( + dict.fromkeys( + source + for _, _, versioned, source, dest in named + if source != dest + and (counts[dest] > 1 or dest in sources) + and source in directories + and not versioned + ) + ) diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 3323999f..90973ae0 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -113,6 +113,11 @@ def test_copy_pairs_needs_destination_lookup(self): with pytest.raises(ValueError, match="destination_is_dir"): S3PathPairing.copy_pairs("s3://bucket/a", "s3://bucket/b", ["bucket/a"]) + def test_copy_pairs_needs_sources(self): + # Not passing the expansion is not the same as expanding to nothing. + with pytest.raises(ValueError, match="sources"): + S3PathPairing.copy_pairs("s3://bucket/a/", "s3://bucket/b/") + def test_move_pairs(self): # The "null" version of a key moved onto the key stays in place. pairs = [("s3://bucket/b?versionId=null", "s3://bucket/b"), ("bucket/d/a", "s3://bucket/z")] @@ -153,6 +158,8 @@ def test_move_pairs_directory_without_object(self): with pytest.raises(ValueError, match="same destination"): S3PathPairing.move_pairs(pairs, missing=set()) assert S3PathPairing.move_pairs(pairs, missing={"bucket/d"}) == pairs + # The missing sources can be given in any form that names them. + assert S3PathPairing.move_pairs(pairs, missing={"s3://bucket/d"}) == pairs def test_move_pairs_version_names_an_object(self): # A version is never taken for a directory. From 4794aa0b95dbbed432b537e8638096d9edb13228 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:18:09 +0900 Subject: [PATCH 04/10] Simplify the move analysis, look up the aio sources in one thread, and keep the planning off the event loop Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 15 +++--- pyathena/filesystem/s3.py | 3 +- pyathena/filesystem/s3_async.py | 16 ++++-- pyathena/filesystem/s3_path.py | 4 +- pyathena/filesystem/s3_path_pairing.py | 57 ++++++++-------------- tests/pyathena/filesystem/test_s3.py | 4 +- tests/pyathena/filesystem/test_s3_async.py | 7 ++- 7 files changed, 48 insertions(+), 58 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 96e1a190..1e2ae2b6 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -350,13 +350,14 @@ filesystems' `cp_file()`, `copy()` and `mv()` run these plans. ## Path pairing -`S3PathPairing` holds the rules by which `copy()`, `get()`, `mv()` and `rm()` pair -their paths, as fsspec pairs them, except that a path with a version ID names that -version, and its destination is named after its key. The rules are pure functions: -the filesystems expand the paths, make the lookups that the rules ask for -(`skips_directories()`, `looks_up_destination()` and `conflict_candidates()`), and -pass the results to `copy_pairs()`, `move_pairs()` and `delete_paths()`. A rule that -needs a lookup that is not passed raises `ValueError`. +`S3PathPairing` holds the rules by which `mv()` and `rm()` pair and expand their +paths, and by which `copy()` and `get()` pair them when a source has a version ID +(fsspec pairs the others). The pairing is fsspec's, except that a path with a version +ID names that version, and its destination is named after its key. The rules are pure +functions: the filesystems ask `expands()`, `skips_directories()`, +`looks_up_destination()` and `conflict_candidates()` what to look up, expand and look +up the paths, and pass the results to `copy_pairs()`, `move_pairs()` and +`delete_paths()`. A rule that needs a lookup that is not passed raises `ValueError`. ```python from pyathena.filesystem.s3_path_pairing import S3PathPairing diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 4f5184a9..f34d310f 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1121,8 +1121,7 @@ def _delete_paths( """ versioned_paths, unversioned_paths = S3PathPairing.delete_paths(path) if not unversioned_paths: - # Versioned paths are deleted as given, without the lookup that - # expand_path makes for them with recursive. + # expand_path raises FileNotFoundError for no paths. return versioned_paths return versioned_paths + self.expand_path( unversioned_paths, recursive=recursive, maxdepth=maxdepth diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index fa3ad97e..de76eedf 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -347,6 +347,7 @@ async def _delete_paths( """ versioned_paths, unversioned_paths = S3PathPairing.delete_paths(path) if not unversioned_paths: + # _expand_path raises FileNotFoundError for no paths. return versioned_paths return versioned_paths + await self._expand_path( unversioned_paths, recursive=recursive, maxdepth=maxdepth @@ -431,14 +432,14 @@ async def _move_pairs( ValueError: If the move has conflicting paths. """ pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) - candidates = S3PathPairing.conflict_candidates(pairs) + candidates = await asyncio.to_thread(S3PathPairing.conflict_candidates, pairs) objects = await asyncio.gather( *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates] ) missing = { source for source, object_ in zip(candidates, objects, strict=True) if not object_ } - return S3PathPairing.move_pairs(pairs, missing=missing) + return await asyncio.to_thread(S3PathPairing.move_pairs, pairs, missing=missing) async def _copy_pairs( self, @@ -467,9 +468,11 @@ async def _copy_pairs( sources = await self._expand_path(path1, recursive=recursive, maxdepth=maxdepth) if S3PathPairing.skips_directories(path1, recursive, maxdepth): # A path with a trailing slash is a directory without a lookup. + # The paths are looked up in one thread, mostly from the cache. files = [p for p in sources if not trailing_sep(p)] - directories = await asyncio.gather(*[self._isdir(p) for p in files]) - sources = [p for p, is_dir in zip(files, directories, strict=True) if not is_dir] + sources = await asyncio.to_thread( + lambda: [p for p in files if not self._sync_fs.isdir(p)] + ) destination_is_dir = None if sources and S3PathPairing.looks_up_destination(path1, path2): # A string, as looks_up_destination() checks. @@ -480,7 +483,10 @@ async def _copy_pairs( if isdir else await self._isdir(destination) ) - return S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) + # The pairing of many paths takes long enough to block the event loop. + return await asyncio.to_thread( + S3PathPairing.copy_pairs, path1, path2, sources, destination_is_dir + ) def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. diff --git a/pyathena/filesystem/s3_path.py b/pyathena/filesystem/s3_path.py index 1620bbca..1f2a4f38 100644 --- a/pyathena/filesystem/s3_path.py +++ b/pyathena/filesystem/s3_path.py @@ -135,9 +135,9 @@ def uri(self) -> str: @property def target(self) -> S3Path: - """The object that a write to the path replaces, for comparing paths. + """What the path names when paths are compared: a ``null`` version names its key. - A ``null`` version names the key itself: in a bucket without + A ``null`` version is taken to name the key itself: in a bucket without versioning, or with versioning suspended, a write to the key replaces its ``null`` version. With versioning enabled, a write adds a new version instead, but the ``null`` version is still taken to name the diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 24a32a6b..a1c0ff04 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -29,9 +29,11 @@ class S3PathPairing: (see :attr:`~pyathena.filesystem.s3_path.S3Path.target`). The rules are pure functions of the paths and of the lookups that they - need, which the caller makes: the filesystems expand the paths and look - them up with their own requests and cache, and pass the results in. A - lookup that a rule needs and that is not passed raises ``ValueError``. + need, which the caller makes: the filesystems ask :meth:`expands`, + :meth:`skips_directories`, :meth:`looks_up_destination` and + :meth:`conflict_candidates` what to look up, expand and look up the paths + with their own requests and cache, and pass the results in. A lookup that + a rule needs and that is not passed raises ``ValueError``. Example: >>> sources = fs.expand_path("s3://bucket/dir", recursive=True) @@ -129,7 +131,8 @@ def copy_pairs( raise ValueError("sources is needed to pair the paths.") if not sources: return [] - if destination_is_dir is None and S3PathPairing.looks_up_destination(path1, path2): + looks_up = S3PathPairing.looks_up_destination(path1, path2) + if destination_is_dir is None and looks_up: raise ValueError("destination_is_dir is needed to pair the paths.") source_is_str = isinstance(path1, str) glob = isinstance(path1, str) and S3PathPairing._is_glob(path1) @@ -139,8 +142,9 @@ def copy_pairs( not glob and not trailing_sep(path1) and isinstance(path2, str) - and (trailing_sep(path2) or bool(destination_is_dir)) + and trailing_sep(path2) ) + or (looks_up and bool(destination_is_dir)) ) names = [S3Path.split_version_id(p)[0] for p in sources] destinations = other_paths(names, path2, exists=exists, flatten=not source_is_str) @@ -165,7 +169,7 @@ def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: another source below them and a destination that conflicts, in the order of the pairs; empty if nothing needs to be looked up. """ - return S3PathPairing._candidates(*S3PathPairing._moves(pairs)) + return S3PathPairing._moves(pairs)[2] @staticmethod def move_pairs( @@ -190,16 +194,15 @@ def move_pairs( except for a directory with no object at its key, which is not copied. Also if ``missing`` is needed and None. """ - named, sources, directories, counts = S3PathPairing._moves(pairs) - candidates = S3PathPairing._candidates(named, sources, directories, counts) + named, sources, candidates = S3PathPairing._moves(pairs) if missing is None and candidates: raise ValueError("missing is needed to check the pairs.") # A directory without an object at its key writes no destination. - skipped = {str(S3Path.parse(path).target) for path in missing or ()} + skipped = set(candidates).intersection( + str(S3Path.parse(path).target) for path in missing or () + ) writers = Counter( - dest - for _, _, _, source, dest in named - if source != dest and not (source in candidates and source in skipped) + dest for _, _, _, source, dest in named if source != dest and source not in skipped ) for dest in writers: if writers[dest] > 1: @@ -250,7 +253,7 @@ def _is_glob(path: str) -> bool: @staticmethod def _moves( pairs: Sequence[tuple[str, str]], - ) -> tuple[list[tuple[str, str, bool, str, str]], set[str], set[str], Counter[str]]: + ) -> tuple[list[tuple[str, str, bool, str, str]], set[str], list[str]]: """Compare the paths of a move by what they name. Args: @@ -259,8 +262,8 @@ def _moves( Returns: Each pair with whether its source has a version and the targets of its source and destination; the targets of all sources, - including those left in place; the parents of the sources; and - the number of moved sources per destination target. + including those left in place; and the conflict candidates (see + :meth:`conflict_candidates`). """ named = [] for p1, p2 in pairs: @@ -284,28 +287,7 @@ def _moves( directories.add(parent) parent = parent.rpartition("/")[0] counts = Counter(dest for _, _, _, source, dest in named if source != dest) - return named, sources, directories, counts - - @staticmethod - def _candidates( - named: list[tuple[str, str, bool, str, str]], - sources: set[str], - directories: set[str], - counts: Counter[str], - ) -> list[str]: - """Select the conflict candidates of a move compared by :meth:`_moves`. - - Args: - named: The pairs with their versions and targets. - sources: The targets of all sources. - directories: The parents of the sources. - counts: The number of moved sources per destination target. - - Returns: - The moved, unversioned sources that are a parent of another - source and whose destination conflicts, without duplicates. - """ - return list( + candidates = list( dict.fromkeys( source for _, _, versioned, source, dest in named @@ -315,3 +297,4 @@ def _candidates( and not versioned ) ) + return named, sources, candidates diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 34b50ca6..5987b5ce 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1849,9 +1849,7 @@ def test_freed_by_reference_counting(self): # The filesystem holds no reference cycle, so a filesystem that is # not cached, such as the internal one of a cursor (GH-978), is freed # as soon as it is unused. - fs = S3FileSystem( - key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True - ) + fs = self._stubbed_fs() ref = weakref.ref(fs) gc.disable() try: diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index fa43bb95..f9164aff 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1118,13 +1118,16 @@ async def test_copy_pairs_destination_lookup(self, path2, lookups, expected): # See TestS3FileSystem.test_copy_pairs_destination_lookup; the aio # filesystem expands and looks up the paths with its own coroutines. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - # Only the destination is a directory. + # Only the destination is a directory. The sources are looked up in + # one thread with the sync isdir, the destination with _isdir. + fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._isdir = mock.AsyncMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) pairs = await fs._copy_pairs("s3://bucket/b?versionId=v1", path2) assert pairs == [("bucket/b?versionId=v1", expected)] - assert [c.args[0] for c in fs._isdir.call_args_list] == lookups + assert [c.args[0] for c in fs._sync_fs.isdir.call_args_list] == lookups[:1] + assert [c.args[0] for c in fs._isdir.call_args_list] == lookups[1:] @pytest.mark.asyncio async def test_move_pairs_looks_up_only_conflict_candidates(self): From 594d94d752a20bd44c3e0c7d7fda0aeaa76d48e8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:19:01 +0900 Subject: [PATCH 05/10] Say exactly when copy() and get() pair through S3PathPairing Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 1e2ae2b6..44e67426 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -351,8 +351,8 @@ filesystems' `cp_file()`, `copy()` and `mv()` run these plans. ## Path pairing `S3PathPairing` holds the rules by which `mv()` and `rm()` pair and expand their -paths, and by which `copy()` and `get()` pair them when a source has a version ID -(fsspec pairs the others). The pairing is fsspec's, except that a path with a version +paths, and by which `copy()` and `get()` pair them when a source has a version ID and +the destination is one path (fsspec pairs the others). The pairing is fsspec's, except that a path with a version ID names that version, and its destination is named after its key. The rules are pure functions: the filesystems ask `expands()`, `skips_directories()`, `looks_up_destination()` and `conflict_candidates()` what to look up, expand and look From 26b399571ede24035dfdc01148c3dc9b5e3a373d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:28:54 +0900 Subject: [PATCH 06/10] Keep versioned writers in the move checks, raise the first aio lookup error, and plan off the event loop Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 4 ++-- pyathena/filesystem/s3_async.py | 20 ++++++++++++------ pyathena/filesystem/s3_path_pairing.py | 6 +++++- tests/pyathena/filesystem/test_s3_async.py | 21 +++++++++++++++++++ .../filesystem/test_s3_path_pairing.py | 14 +++++++++++++ 5 files changed, 56 insertions(+), 9 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 44e67426..46740151 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -350,8 +350,8 @@ filesystems' `cp_file()`, `copy()` and `mv()` run these plans. ## Path pairing -`S3PathPairing` holds the rules by which `mv()` and `rm()` pair and expand their -paths, and by which `copy()` and `get()` pair them when a source has a version ID and +`S3PathPairing` holds the rules by which `mv()` pairs its paths and `rm()` expands +them, and by which `copy()` and `get()` pair them when a source has a version ID and the destination is one path (fsspec pairs the others). The pairing is fsspec's, except that a path with a version ID names that version, and its destination is named after its key. The rules are pure functions: the filesystems ask `expands()`, `skips_directories()`, diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index de76eedf..ec3296e6 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -345,7 +345,9 @@ async def _delete_paths( Raises: ValueError: If a path is a bucket. """ - versioned_paths, unversioned_paths = S3PathPairing.delete_paths(path) + versioned_paths, unversioned_paths = await asyncio.to_thread( + S3PathPairing.delete_paths, path + ) if not unversioned_paths: # _expand_path raises FileNotFoundError for no paths. return versioned_paths @@ -434,11 +436,17 @@ async def _move_pairs( pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) candidates = await asyncio.to_thread(S3PathPairing.conflict_candidates, pairs) objects = await asyncio.gather( - *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates] + *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates], + return_exceptions=True, ) - missing = { - source for source, object_ in zip(candidates, objects, strict=True) if not object_ - } + missing = set() + # Every lookup finishes; the error of the first candidate is raised, + # as when the lookups run in order. + for source, object_ in zip(candidates, objects, strict=True): + if isinstance(object_, BaseException): + raise object_ + if object_ is None: + missing.add(source) return await asyncio.to_thread(S3PathPairing.move_pairs, pairs, missing=missing) async def _copy_pairs( @@ -464,7 +472,7 @@ async def _copy_pairs( The sources and their destinations. """ if not S3PathPairing.expands(path1, path2): - return S3PathPairing.copy_pairs(path1, path2) + return await asyncio.to_thread(S3PathPairing.copy_pairs, path1, path2) sources = await self._expand_path(path1, recursive=recursive, maxdepth=maxdepth) if S3PathPairing.skips_directories(path1, recursive, maxdepth): # A path with a trailing slash is a directory without a lookup. diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index a1c0ff04..970ef169 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -201,8 +201,12 @@ def move_pairs( skipped = set(candidates).intersection( str(S3Path.parse(path).target) for path in missing or () ) + # A path with a version always names an object, so it writes its + # destination even when the key has no current object. writers = Counter( - dest for _, _, _, source, dest in named if source != dest and source not in skipped + dest + for _, _, versioned, source, dest in named + if source != dest and (versioned or source not in skipped) ) for dest in writers: if writers[dest] > 1: diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f9164aff..ad71cc85 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1143,6 +1143,27 @@ async def test_move_pairs_looks_up_only_conflict_candidates(self): assert len(pairs) == 4 fs._sync_fs._head_object.assert_called_once_with("bucket/d") + @pytest.mark.asyncio + async def test_move_pairs_raises_the_first_lookup_error(self): + # The conflict lookups run concurrently, but the error of the first + # candidate is raised, as when they run in order. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + + def head_object(path): + if path == "bucket/d1": + time.sleep(0.1) + raise PermissionError(path) + raise TimeoutError(path) + + fs._sync_fs._head_object = mock.MagicMock(side_effect=head_object) + + with pytest.raises(PermissionError, match="bucket/d1"): + await fs._move_pairs( + ["s3://bucket/d1", "s3://bucket/d1/x", "s3://bucket/d2", "s3://bucket/d2/x"], + ["s3://bucket/o", "s3://bucket/o", "s3://bucket/p", "s3://bucket/p"], + ) + assert fs._sync_fs._head_object.call_count == 2 + def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance # cache, so skip_instance_cache=True instances shared it. diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 90973ae0..05e5e314 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -173,6 +173,20 @@ def test_move_pairs_version_names_an_object(self): with pytest.raises(ValueError, match="same destination"): S3PathPairing.move_pairs(pairs) + def test_move_pairs_version_of_missing_directory_key(self): + # The "null" version of a key without a current object still names an + # object, so it writes its destination although the key is missing. + pairs = [ + ("src/d", "dst/out"), + ("src/d?versionId=null", "dst/out"), + ("src/a", "dst/out"), + ("src/d/x", "dst/x"), + ] + + assert S3PathPairing.conflict_candidates(pairs) == ["src/d"] + with pytest.raises(ValueError, match="same destination"): + S3PathPairing.move_pairs(pairs, missing={"src/d"}) + def test_delete_paths(self): assert S3PathPairing.delete_paths( ["s3://bucket/d", "s3://bucket/b?versionId=v1", "s3://bucket/c"] From 9cb261ef687c7de79b54a14032041262275b03cf Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:44:37 +0900 Subject: [PATCH 07/10] Make S3PathPairing a frozen value of the paths it pairs S3PathPairing now holds path1, path2, recursive and maxdepth, so the questions (expands, skips_directories, looks_up_destination) are properties and copy_pairs(), conflict_candidates() and move_pairs() are methods of one pairing instead of static functions that took the same paths again. delete_paths() stays a classmethod, as rm() has one path. The adapters build one pairing per operation and pass it to _copy_pairs(). The pairing results are unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 27 ++-- pyathena/filesystem/s3.py | 45 +++--- pyathena/filesystem/s3_async.py | 47 +++--- pyathena/filesystem/s3_path_pairing.py | 145 ++++++++---------- tests/pyathena/filesystem/test_s3.py | 3 +- tests/pyathena/filesystem/test_s3_async.py | 3 +- .../filesystem/test_s3_path_pairing.py | 44 +++--- 7 files changed, 152 insertions(+), 162 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 01b07c7c..431d2620 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -373,23 +373,24 @@ except Exception: ## Path pairing -`S3PathPairing` holds the rules by which `mv()` pairs its paths and `rm()` expands -them, and by which `copy()` and `get()` pair them when a source has a version ID and -the destination is one path (fsspec pairs the others). The pairing is fsspec's, -except that a path with a version ID names that version, and its destination is named -after its key. The rules are pure functions: the filesystems ask `expands()`, -`skips_directories()`, `looks_up_destination()` and `conflict_candidates()` what to -look up, expand and look up the paths, and pass the results to `copy_pairs()`, -`move_pairs()` and `delete_paths()`. A rule that needs a lookup that is not passed -raises `ValueError`. +`S3PathPairing` pairs the paths of one `copy()`, `get()` or `mv()`: `mv()` always, and +`copy()` and `get()` when a source has a version ID and the destination is one path +(fsspec pairs the others). Its `delete_paths()` splits the paths of an `rm()` into +those deleted as given and those expanded. The pairing is fsspec's, except that a +path with a version ID names that version, and its destination is named after its key. + +A pairing is a pure value of `path1`, `path2`, `recursive` and `maxdepth`. The +filesystems read from it what to look up (`expands`, `skips_directories`, +`looks_up_destination` and `conflict_candidates()`), expand and look up the paths, and +pass the results to `copy_pairs()` and `move_pairs()`. A rule that needs a lookup that +is not passed raises `ValueError`. ```python from pyathena.filesystem.s3_path_pairing import S3PathPairing -sources = fs.expand_path("s3://YOUR_S3_BUCKET/src/", recursive=True) -for source, destination in S3PathPairing.copy_pairs( - "s3://YOUR_S3_BUCKET/src/", "s3://YOUR_S3_BUCKET/dst/", sources -): +pairing = S3PathPairing("s3://YOUR_S3_BUCKET/src/", "s3://YOUR_S3_BUCKET/dst/", recursive=True) +sources = fs.expand_path(pairing.path1, recursive=pairing.recursive) +for source, destination in pairing.copy_pairs(sources): print(source, "->", destination) ``` diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 1161939d..3e2f8a01 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1463,50 +1463,44 @@ def _move_pairs( Raises: ValueError: If the move has conflicting paths. """ - pairs = self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = self._copy_pairs(pairing) missing = { source - for source in S3PathPairing.conflict_candidates(pairs) + for source in pairing.conflict_candidates(pairs) if self._head_object(source) is None } - return S3PathPairing.move_pairs(pairs, missing=missing) + return pairing.move_pairs(pairs, missing=missing) def _copy_pairs( - self, - path1: str | list[str], - path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, - isdir: Callable[[str], bool] | None = None, + self, pairing: S3PathPairing, isdir: Callable[[str], bool] | None = None ) -> list[tuple[str, str]]: """Expand and pair the paths of a copy (see :meth:`S3PathPairing.copy_pairs`). The destination is looked up only when it decides the pairing. Args: - path1: Source S3 path, glob pattern, or list of them. - path2: Destination path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. + pairing: The paths of the copy. isdir: Whether the destination is a directory, by default ``self.isdir``; ``get()`` passes the local filesystem's. Returns: The sources and their destinations. """ - if not S3PathPairing.expands(path1, path2): - return S3PathPairing.copy_pairs(path1, path2) - sources = self.expand_path(path1, recursive=recursive, maxdepth=maxdepth) - if S3PathPairing.skips_directories(path1, recursive, maxdepth): + if not pairing.expands: + return pairing.copy_pairs() + sources = self.expand_path( + pairing.path1, recursive=pairing.recursive, maxdepth=pairing.maxdepth + ) + if pairing.skips_directories: sources = [p for p in sources if not (trailing_sep(p) or self.isdir(p))] destination_is_dir = ( - # A string, as looks_up_destination() checks. - (isdir or self.isdir)(cast(str, path2)) - if sources and S3PathPairing.looks_up_destination(path1, path2) + # A string, as looks_up_destination checks. + (isdir or self.isdir)(cast(str, pairing.path2)) + if sources and pairing.looks_up_destination else None ) - return S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) + return pairing.copy_pairs(sources, destination_is_dir) def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kwargs) -> None: """Copy files within S3. @@ -1529,7 +1523,9 @@ def copy(self, path1, path2, recursive=False, maxdepth=None, on_error=None, **kw """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = self._copy_pairs( + S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + ) if not pairs: return path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] @@ -1564,7 +1560,8 @@ def get( if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): root = make_path_posix(lpath) pairs = self._copy_pairs( - rpath, root, recursive=recursive, maxdepth=maxdepth, isdir=LocalFileSystem().isdir + 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) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 37fe614f..38ba2870 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -433,8 +433,9 @@ async def _move_pairs( Raises: ValueError: If the move has conflicting paths. """ - pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) - candidates = await asyncio.to_thread(S3PathPairing.conflict_candidates, pairs) + pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = await self._copy_pairs(pairing) + candidates = await asyncio.to_thread(pairing.conflict_candidates, pairs) objects = await asyncio.gather( *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates], return_exceptions=True, @@ -447,34 +448,27 @@ async def _move_pairs( raise object_ if object_ is None: missing.add(source) - return await asyncio.to_thread(S3PathPairing.move_pairs, pairs, missing=missing) + return await asyncio.to_thread(pairing.move_pairs, pairs, missing=missing) async def _copy_pairs( - self, - path1: str | list[str], - path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, - isdir: Callable[[str], bool] | None = None, + self, pairing: S3PathPairing, isdir: Callable[[str], bool] | None = None ) -> list[tuple[str, str]]: """Expand and pair the paths of a copy, as ``S3FileSystem._copy_pairs`` does. Args: - path1: Source S3 path, glob pattern, or list of them. - path2: Destination path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. + pairing: The paths of the copy. isdir: Whether the destination is a directory, by default ``self._isdir``; ``_get()`` passes the local filesystem's. Returns: The sources and their destinations. """ - if not S3PathPairing.expands(path1, path2): - return await asyncio.to_thread(S3PathPairing.copy_pairs, path1, path2) - sources = await self._expand_path(path1, recursive=recursive, maxdepth=maxdepth) - if S3PathPairing.skips_directories(path1, recursive, maxdepth): + if not pairing.expands: + return await asyncio.to_thread(pairing.copy_pairs) + sources = await self._expand_path( + pairing.path1, recursive=pairing.recursive, maxdepth=pairing.maxdepth + ) + if pairing.skips_directories: # A path with a trailing slash is a directory without a lookup. # The paths are looked up in one thread, mostly from the cache. files = [p for p in sources if not trailing_sep(p)] @@ -482,9 +476,9 @@ async def _copy_pairs( lambda: [p for p in files if not self._sync_fs.isdir(p)] ) destination_is_dir = None - if sources and S3PathPairing.looks_up_destination(path1, path2): - # A string, as looks_up_destination() checks. - destination = cast(str, path2) + if sources and pairing.looks_up_destination: + # A string, as looks_up_destination checks. + destination = cast(str, pairing.path2) destination_is_dir = ( # A local isdir, which can block, runs in a thread. await asyncio.to_thread(isdir, destination) @@ -492,9 +486,7 @@ async def _copy_pairs( else await self._isdir(destination) ) # The pairing of many paths takes long enough to block the event loop. - return await asyncio.to_thread( - S3PathPairing.copy_pairs, path1, path2, sources, destination_is_dir - ) + return await asyncio.to_thread(pairing.copy_pairs, sources, destination_is_dir) def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. @@ -539,7 +531,9 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = await self._copy_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = await self._copy_pairs( + S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + ) if not pairs: return path1, path2 = [p1 for p1, _ in pairs], [p2 for _, p2 in pairs] @@ -578,7 +572,8 @@ async def _get( 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._copy_pairs( - rpath, root, recursive=recursive, maxdepth=maxdepth, isdir=LocalFileSystem().isdir + 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) diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 970ef169..7c4bef8c 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -11,6 +11,7 @@ from collections import Counter from collections.abc import Collection, Sequence +from dataclasses import dataclass from glob import has_magic from fsspec.implementations.local import trailing_sep @@ -19,123 +20,111 @@ from pyathena.filesystem.s3_path import S3Path +@dataclass(frozen=True) class S3PathPairing: - """The rules that pair the paths of ``copy()``, ``get()``, ``mv()`` and ``rm()``. + """The pairing of the paths of one ``copy()``, ``get()`` or ``mv()``. The sources are paired with their destinations as fsspec's ``copy()`` pairs them, except that a path with a version ID names that version of an object: it is not a glob pattern, and its destination is named after its key without the version. A move compares the paths by what they name (see :attr:`~pyathena.filesystem.s3_path.S3Path.target`). + :meth:`delete_paths` splits the paths of an ``rm()``. - The rules are pure functions of the paths and of the lookups that they - need, which the caller makes: the filesystems ask :meth:`expands`, - :meth:`skips_directories`, :meth:`looks_up_destination` and - :meth:`conflict_candidates` what to look up, expand and look up the paths - with their own requests and cache, and pass the results in. A lookup that - a rule needs and that is not passed raises ``ValueError``. + The pairing is a pure value that holds no filesystem. The caller makes + the lookups that it asks for (:attr:`expands`, :attr:`skips_directories`, + :attr:`looks_up_destination` and :meth:`conflict_candidates`) with its + own requests and cache, and passes the results in. A lookup that a rule + needs and that is not passed raises ``ValueError``. + + Attributes: + path1: The source path, glob pattern, or list of them, as given to + the copy. + path2: The destination path, or list of paths. + recursive: Whether the copy includes the contents of directories. + maxdepth: The maximum depth of the expansion. Example: - >>> sources = fs.expand_path("s3://bucket/dir", recursive=True) - >>> S3PathPairing.copy_pairs("s3://bucket/dir", "s3://bucket/copy/", sources) + >>> pairing = S3PathPairing("s3://bucket/dir/", "s3://bucket/copy/", recursive=True) + >>> sources = fs.expand_path(pairing.path1, recursive=True) + >>> pairing.copy_pairs(sources) """ - @staticmethod - def expands(path1: str | list[str], path2: str | list[str]) -> bool: - """Return whether a copy expands its sources. + path1: str | list[str] + path2: str | list[str] + recursive: bool = False + maxdepth: int | None = None - Args: - path1: Source path, glob pattern, or list of them. - path2: Destination path, or list of paths. + @property + def expands(self) -> bool: + """Whether the sources are expanded. - Returns: - False if both are lists, which are paired as given, without - expansion or lookups. + False if both paths are lists, which are paired as given, without + expansion or lookups. """ - return not (isinstance(path1, list) and isinstance(path2, list)) + return not (isinstance(self.path1, list) and isinstance(self.path2, list)) - @staticmethod - def skips_directories(path1: str | list[str], recursive: bool, maxdepth: int | None) -> bool: - """Return whether a copy leaves out the directories among its expanded sources. + @property + def skips_directories(self) -> bool: + """Whether the directories among the expanded sources are left out. - A directory is a source that ends with a slash or that the filesystem - reports as a directory. - - Args: - path1: Source path, glob pattern, or list of them. - recursive: Whether the copy includes the contents of directories. - maxdepth: Maximum depth of the expansion. - - Returns: - True for a string source copied without ``recursive``, or with a - ``maxdepth``. + True for a string source copied without ``recursive``, or with a + ``maxdepth``. A directory is a source that ends with a slash or that + the filesystem reports as a directory. """ - return isinstance(path1, str) and (not recursive or maxdepth is not None) + return isinstance(self.path1, str) and (not self.recursive or self.maxdepth is not None) - @staticmethod - def looks_up_destination(path1: str | list[str], path2: str | list[str]) -> bool: - """Return whether the pairing of a copy depends on its destination being a directory. - - Args: - path1: Source path, glob pattern, or list of them. - path2: Destination path, or list of paths. + @property + def looks_up_destination(self) -> bool: + """Whether the pairing depends on the destination being a directory. - Returns: - True for a string source that is neither a glob pattern nor ends - with a slash, copied to a string destination that does not end - with a slash. :meth:`copy_pairs` then needs - ``destination_is_dir``. + True for a string source that is neither a glob pattern nor ends with + a slash, copied to a string destination that does not end with a + slash. :meth:`copy_pairs` then needs ``destination_is_dir``. """ return ( - isinstance(path1, str) - and not S3PathPairing._is_glob(path1) - and not trailing_sep(path1) - and isinstance(path2, str) - and not trailing_sep(path2) + isinstance(self.path1, str) + and not self._is_glob(self.path1) + and not trailing_sep(self.path1) + and isinstance(self.path2, str) + and not trailing_sep(self.path2) ) - @staticmethod def copy_pairs( - path1: str | list[str], - path2: str | list[str], + self, sources: Sequence[str] | None = None, destination_is_dir: bool | None = None, ) -> list[tuple[str, str]]: - """Pair the sources of a copy with their destinations as fsspec's ``copy()`` does. + """Pair the sources with their destinations as fsspec's ``copy()`` does. Args: - path1: Source S3 path, glob pattern, or list of them, as given to - the copy. - path2: Destination path, or a list of paths: as many as the - sources, or, when ``path1`` is a list, its destinations. sources: The expansion of ``path1``, without the directories that - :meth:`skips_directories` leaves out; needed unless both - ``path1`` and ``path2`` are lists, which are paired without - it. + :attr:`skips_directories` leaves out; needed when + :attr:`expands`. destination_is_dir: Whether ``path2`` is a directory; needed when - :meth:`looks_up_destination` is true and there are sources. + :attr:`looks_up_destination` is true and there are sources. Returns: The sources and their destinations, which keep the form of - ``path2``. When both ``path1`` and ``path2`` are lists, they are - paired as given, up to the end of the shorter list, as in - fsspec's ``copy()``. Empty if there are no sources. + ``path2``. When both paths are lists, they are paired as given, up + to the end of the shorter list, as in fsspec's ``copy()``. Empty + if there are no sources. Raises: ValueError: If ``sources`` or ``destination_is_dir`` is needed and None. """ - if not S3PathPairing.expands(path1, path2): + path1, path2 = self.path1, self.path2 + if not self.expands: return list(zip(path1, path2, strict=False)) if sources is None: raise ValueError("sources is needed to pair the paths.") if not sources: return [] - looks_up = S3PathPairing.looks_up_destination(path1, path2) - if destination_is_dir is None and looks_up: + if destination_is_dir is None and self.looks_up_destination: raise ValueError("destination_is_dir is needed to pair the paths.") source_is_str = isinstance(path1, str) - glob = isinstance(path1, str) and S3PathPairing._is_glob(path1) + glob = isinstance(path1, str) and self._is_glob(path1) exists = source_is_str and ( (glob and len(sources) == 1) or ( @@ -144,14 +133,13 @@ def copy_pairs( and isinstance(path2, str) and trailing_sep(path2) ) - or (looks_up and bool(destination_is_dir)) + or (self.looks_up_destination and bool(destination_is_dir)) ) names = [S3Path.split_version_id(p)[0] for p in sources] destinations = other_paths(names, path2, exists=exists, flatten=not source_is_str) return list(zip(sources, destinations, strict=True)) - @staticmethod - def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: + def conflict_candidates(self, pairs: Sequence[tuple[str, str]]) -> list[str]: """Return the sources of a move whose conflicts depend on an object at their key. A source with another source below it may be a directory. If no object @@ -161,7 +149,7 @@ def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: Args: pairs: The sources and destinations of the move, as - :meth:`copy_pairs` pairs them. + :meth:`copy_pairs` of this pairing returns them. Returns: The sources, in ``bucket/key`` form (their @@ -171,15 +159,14 @@ def conflict_candidates(pairs: Sequence[tuple[str, str]]) -> list[str]: """ return S3PathPairing._moves(pairs)[2] - @staticmethod def move_pairs( - pairs: Sequence[tuple[str, str]], missing: Collection[str] | None = None + self, pairs: Sequence[tuple[str, str]], missing: Collection[str] | None = None ) -> list[tuple[str, str]]: """Check the pairs of a move and leave out the sources that stay in place. Args: pairs: The sources and destinations of the move, as - :meth:`copy_pairs` pairs them. + :meth:`copy_pairs` of this pairing returns them. missing: The :meth:`conflict_candidates` without an object at their key, in any form that names them; needed when there are candidates. @@ -215,8 +202,8 @@ def move_pairs( raise ValueError("Cannot move a path onto another path that is moved.") return [(p1, p2) for p1, p2, _, source, dest in named if source != dest] - @staticmethod - def delete_paths(path: str | list[str]) -> tuple[list[str], list[str]]: + @classmethod + def delete_paths(cls, path: str | list[str]) -> tuple[list[str], list[str]]: """Split the paths that ``rm()`` deletes into those deleted as given and those expanded. Args: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 6531e3da..6f99e53a 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -44,6 +44,7 @@ from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path +from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -1832,7 +1833,7 @@ def test_copy_pairs_destination_lookup(self, path2, lookups, expected): # Only the destination is a directory. fs.isdir = mock.MagicMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) - pairs = fs._copy_pairs("s3://bucket/b?versionId=v1", path2) + pairs = fs._copy_pairs(S3PathPairing("s3://bucket/b?versionId=v1", path2)) assert pairs == [("bucket/b?versionId=v1", expected)] assert [c.args[0] for c in fs.isdir.call_args_list] == lookups diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 0fc8af6a..d53c6416 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -31,6 +31,7 @@ S3StorageClass, ) from pyathena.filesystem.s3_path import S3Path +from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -1135,7 +1136,7 @@ async def test_copy_pairs_destination_lookup(self, path2, lookups, expected): fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._isdir = mock.AsyncMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) - pairs = await fs._copy_pairs("s3://bucket/b?versionId=v1", path2) + pairs = await fs._copy_pairs(S3PathPairing("s3://bucket/b?versionId=v1", path2)) assert pairs == [("bucket/b?versionId=v1", expected)] assert [c.args[0] for c in fs._sync_fs.isdir.call_args_list] == lookups[:1] diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 05e5e314..81829928 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -10,6 +10,11 @@ from pyathena.filesystem.s3_path_pairing import S3PathPairing +def _pairing(pairs): + # The pairing of a move of the paths as lists. + return S3PathPairing([p1 for p1, _ in pairs], [p2 for _, p2 in pairs]) + + class TestS3PathPairing: @pytest.mark.parametrize( ("path1", "path2", "expected"), @@ -22,7 +27,7 @@ class TestS3PathPairing: ], ) def test_expands(self, path1, path2, expected): - assert S3PathPairing.expands(path1, path2) is expected + assert S3PathPairing(path1, path2).expands is expected @pytest.mark.parametrize( ("path1", "recursive", "maxdepth", "expected"), @@ -34,7 +39,10 @@ def test_expands(self, path1, path2, expected): ], ) def test_skips_directories(self, path1, recursive, maxdepth, expected): - assert S3PathPairing.skips_directories(path1, recursive, maxdepth) is expected + assert ( + S3PathPairing(path1, "s3://bucket/out", recursive, maxdepth).skips_directories + is expected + ) @pytest.mark.parametrize( ("path1", "path2", "expected"), @@ -52,7 +60,7 @@ def test_skips_directories(self, path1, recursive, maxdepth, expected): ], ) def test_looks_up_destination(self, path1, path2, expected): - assert S3PathPairing.looks_up_destination(path1, path2) is expected + assert S3PathPairing(path1, path2).looks_up_destination is expected @pytest.mark.parametrize( ("path1", "path2", "sources", "destination_is_dir", "expected"), @@ -107,23 +115,23 @@ def test_looks_up_destination(self, path1, path2, expected): ], ) def test_copy_pairs(self, path1, path2, sources, destination_is_dir, expected): - assert S3PathPairing.copy_pairs(path1, path2, sources, destination_is_dir) == expected + assert S3PathPairing(path1, path2).copy_pairs(sources, destination_is_dir) == expected def test_copy_pairs_needs_destination_lookup(self): with pytest.raises(ValueError, match="destination_is_dir"): - S3PathPairing.copy_pairs("s3://bucket/a", "s3://bucket/b", ["bucket/a"]) + S3PathPairing("s3://bucket/a", "s3://bucket/b").copy_pairs(["bucket/a"]) def test_copy_pairs_needs_sources(self): # Not passing the expansion is not the same as expanding to nothing. with pytest.raises(ValueError, match="sources"): - S3PathPairing.copy_pairs("s3://bucket/a/", "s3://bucket/b/") + S3PathPairing("s3://bucket/a/", "s3://bucket/b/").copy_pairs() def test_move_pairs(self): # The "null" version of a key moved onto the key stays in place. pairs = [("s3://bucket/b?versionId=null", "s3://bucket/b"), ("bucket/d/a", "s3://bucket/z")] - assert S3PathPairing.conflict_candidates(pairs) == [] - assert S3PathPairing.move_pairs(pairs) == [("bucket/d/a", "s3://bucket/z")] + assert _pairing(pairs).conflict_candidates(pairs) == [] + assert _pairing(pairs).move_pairs(pairs) == [("bucket/d/a", "s3://bucket/z")] @pytest.mark.parametrize( ("pairs", "match"), @@ -140,7 +148,7 @@ def test_move_pairs(self): ) def test_move_pairs_conflicts(self, pairs, match): with pytest.raises(ValueError, match=match): - S3PathPairing.move_pairs(pairs) + _pairing(pairs).move_pairs(pairs) def test_move_pairs_directory_without_object(self): # A source with another source below it may be a directory; one @@ -152,14 +160,14 @@ def test_move_pairs_directory_without_object(self): ("s3://bucket/e/y", "s3://bucket/out"), ] - assert S3PathPairing.conflict_candidates(pairs) == ["bucket/d"] + assert _pairing(pairs).conflict_candidates(pairs) == ["bucket/d"] with pytest.raises(ValueError, match="missing"): - S3PathPairing.move_pairs(pairs) + _pairing(pairs).move_pairs(pairs) with pytest.raises(ValueError, match="same destination"): - S3PathPairing.move_pairs(pairs, missing=set()) - assert S3PathPairing.move_pairs(pairs, missing={"bucket/d"}) == pairs + _pairing(pairs).move_pairs(pairs, missing=set()) + assert _pairing(pairs).move_pairs(pairs, missing={"bucket/d"}) == pairs # The missing sources can be given in any form that names them. - assert S3PathPairing.move_pairs(pairs, missing={"s3://bucket/d"}) == pairs + assert _pairing(pairs).move_pairs(pairs, missing={"s3://bucket/d"}) == pairs def test_move_pairs_version_names_an_object(self): # A version is never taken for a directory. @@ -169,9 +177,9 @@ def test_move_pairs_version_names_an_object(self): ("s3://bucket/a", "s3://bucket/out"), ] - assert S3PathPairing.conflict_candidates(pairs) == [] + assert _pairing(pairs).conflict_candidates(pairs) == [] with pytest.raises(ValueError, match="same destination"): - S3PathPairing.move_pairs(pairs) + _pairing(pairs).move_pairs(pairs) def test_move_pairs_version_of_missing_directory_key(self): # The "null" version of a key without a current object still names an @@ -183,9 +191,9 @@ def test_move_pairs_version_of_missing_directory_key(self): ("src/d/x", "dst/x"), ] - assert S3PathPairing.conflict_candidates(pairs) == ["src/d"] + assert _pairing(pairs).conflict_candidates(pairs) == ["src/d"] with pytest.raises(ValueError, match="same destination"): - S3PathPairing.move_pairs(pairs, missing={"src/d"}) + _pairing(pairs).move_pairs(pairs, missing={"src/d"}) def test_delete_paths(self): assert S3PathPairing.delete_paths( From 97aea57d328d22ad3a15662c6fdc0ebb58d6683f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:50:35 +0900 Subject: [PATCH 08/10] Reject a string as missing, look up the aio conflicts in order in one thread, and use one isdir in aio Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 12 ++++--- pyathena/filesystem/s3_async.py | 31 +++++++------------ pyathena/filesystem/s3_path_pairing.py | 16 +++++++--- tests/pyathena/filesystem/test_s3_async.py | 16 +++++----- .../filesystem/test_s3_path_pairing.py | 5 +++ 5 files changed, 41 insertions(+), 39 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 431d2620..181b09b3 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -379,11 +379,13 @@ except Exception: those deleted as given and those expanded. The pairing is fsspec's, except that a path with a version ID names that version, and its destination is named after its key. -A pairing is a pure value of `path1`, `path2`, `recursive` and `maxdepth`. The -filesystems read from it what to look up (`expands`, `skips_directories`, -`looks_up_destination` and `conflict_candidates()`), expand and look up the paths, and -pass the results to `copy_pairs()` and `move_pairs()`. A rule that needs a lookup that -is not passed raises `ValueError`. +A pairing is a frozen dataclass of `path1`, `path2`, `recursive` and `maxdepth`, and +holds no filesystem. The filesystems read from it what to look up (`expands`, +`skips_directories`, `looks_up_destination` and `conflict_candidates()`), expand the +sources, leave out the directories when `skips_directories` is true, look up the +paths, and pass the results to `copy_pairs()` and `move_pairs()`. A `sources`, +`destination_is_dir` or `missing` that a rule needs and that is not passed raises +`ValueError`. ```python from pyathena.filesystem.s3_path_pairing import S3PathPairing diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 38ba2870..c19702a7 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -436,18 +436,11 @@ async def _move_pairs( pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) pairs = await self._copy_pairs(pairing) candidates = await asyncio.to_thread(pairing.conflict_candidates, pairs) - objects = await asyncio.gather( - *[asyncio.to_thread(self._sync_fs._head_object, source) for source in candidates], - return_exceptions=True, + # Looked up in order in one thread, as S3FileSystem does, so that the + # first error stops the lookups. + missing = await asyncio.to_thread( + lambda: {s for s in candidates if self._sync_fs._head_object(s) is None} ) - missing = set() - # Every lookup finishes; the error of the first candidate is raised, - # as when the lookups run in order. - for source, object_ in zip(candidates, objects, strict=True): - if isinstance(object_, BaseException): - raise object_ - if object_ is None: - missing.add(source) return await asyncio.to_thread(pairing.move_pairs, pairs, missing=missing) async def _copy_pairs( @@ -457,8 +450,9 @@ async def _copy_pairs( Args: pairing: The paths of the copy. - isdir: Whether the destination is a directory, by default - ``self._isdir``; ``_get()`` passes the local filesystem's. + isdir: Whether the destination is a directory, by default the + wrapped ``S3FileSystem.isdir``; ``_get()`` passes the local + filesystem's. Returns: The sources and their destinations. @@ -477,13 +471,10 @@ async def _copy_pairs( ) destination_is_dir = None if sources and pairing.looks_up_destination: - # A string, as looks_up_destination checks. - destination = cast(str, pairing.path2) - destination_is_dir = ( - # A local isdir, which can block, runs in a thread. - await asyncio.to_thread(isdir, destination) - if isdir - else await self._isdir(destination) + # A string, as looks_up_destination checks; looked up in a thread + # as the sources are. + destination_is_dir = await asyncio.to_thread( + isdir or self._sync_fs.isdir, cast(str, pairing.path2) ) # The pairing of many paths takes long enough to block the event loop. return await asyncio.to_thread(pairing.copy_pairs, sources, destination_is_dir) diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 7c4bef8c..e2f61433 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -31,11 +31,14 @@ class S3PathPairing: (see :attr:`~pyathena.filesystem.s3_path.S3Path.target`). :meth:`delete_paths` splits the paths of an ``rm()``. - The pairing is a pure value that holds no filesystem. The caller makes - the lookups that it asks for (:attr:`expands`, :attr:`skips_directories`, - :attr:`looks_up_destination` and :meth:`conflict_candidates`) with its - own requests and cache, and passes the results in. A lookup that a rule - needs and that is not passed raises ``ValueError``. + The pairing is a frozen dataclass of the paths as given, and holds no + filesystem. The caller makes the lookups that it asks for + (:attr:`expands`, :attr:`skips_directories`, :attr:`looks_up_destination` + and :meth:`conflict_candidates`) with its own requests and cache: it + expands the sources, leaves out the directories when + :attr:`skips_directories` says so, and passes the results in. A + ``sources``, ``destination_is_dir`` or ``missing`` that a rule needs and + that is not passed raises ``ValueError``. Attributes: path1: The source path, glob pattern, or list of them, as given to @@ -180,7 +183,10 @@ def move_pairs( destination is another source, including one left in place, except for a directory with no object at its key, which is not copied. Also if ``missing`` is needed and None. + TypeError: If ``missing`` is a string instead of a collection. """ + if isinstance(missing, str): + raise TypeError("missing is a collection of paths, not a path.") named, sources, candidates = S3PathPairing._moves(pairs) if missing is None and candidates: raise ValueError("missing is needed to check the pairs.") diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index d53c6416..b436d63d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1131,16 +1131,14 @@ async def test_copy_pairs_destination_lookup(self, path2, lookups, expected): # See TestS3FileSystem.test_copy_pairs_destination_lookup; the aio # filesystem expands and looks up the paths with its own coroutines. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - # Only the destination is a directory. The sources are looked up in - # one thread with the sync isdir, the destination with _isdir. - fs._sync_fs.isdir = mock.MagicMock(return_value=False) - fs._isdir = mock.AsyncMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) + # Only the destination is a directory; the paths are looked up in a + # thread with the cached isdir of the wrapped filesystem. + fs._sync_fs.isdir = mock.MagicMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) pairs = await fs._copy_pairs(S3PathPairing("s3://bucket/b?versionId=v1", path2)) assert pairs == [("bucket/b?versionId=v1", expected)] - assert [c.args[0] for c in fs._sync_fs.isdir.call_args_list] == lookups[:1] - assert [c.args[0] for c in fs._isdir.call_args_list] == lookups[1:] + assert [c.args[0] for c in fs._sync_fs.isdir.call_args_list] == lookups @pytest.mark.asyncio async def test_move_pairs_looks_up_only_conflict_candidates(self): @@ -1158,8 +1156,8 @@ async def test_move_pairs_looks_up_only_conflict_candidates(self): @pytest.mark.asyncio async def test_move_pairs_raises_the_first_lookup_error(self): - # The conflict lookups run concurrently, but the error of the first - # candidate is raised, as when they run in order. + # The conflict lookups run in order, as in S3FileSystem, so the first + # error is raised and stops them. fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) def head_object(path): @@ -1175,7 +1173,7 @@ def head_object(path): ["s3://bucket/d1", "s3://bucket/d1/x", "s3://bucket/d2", "s3://bucket/d2/x"], ["s3://bucket/o", "s3://bucket/o", "s3://bucket/p", "s3://bucket/p"], ) - assert fs._sync_fs._head_object.call_count == 2 + assert fs._sync_fs._head_object.call_count == 1 def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index 81829928..ccdfb445 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -169,6 +169,11 @@ def test_move_pairs_directory_without_object(self): # The missing sources can be given in any form that names them. assert _pairing(pairs).move_pairs(pairs, missing={"s3://bucket/d"}) == pairs + def test_move_pairs_missing_is_a_collection(self): + pairs = [("s3://bucket/d", "s3://bucket/e"), ("s3://bucket/d/x", "s3://bucket/e")] + with pytest.raises(TypeError, match="collection"): + _pairing(pairs).move_pairs(pairs, missing="bucket/d") + def test_move_pairs_version_names_an_object(self): # A version is never taken for a directory. pairs = [ From 30e48052650506d2758755e6527ff688caddac32 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 18:59:11 +0900 Subject: [PATCH 09/10] Run the sync pairing orchestration in a thread from aio, as master did The aio expansion with fsspec's async glob lists with a prefix, which is not cached, so the directory filter of a non-recursive move sent a HeadObject per matched path. AioS3FileSystem now calls the sync _copy_pairs, _move_pairs and _delete_paths in one thread, which sends the same requests as master and keeps the expansion off the event loop; the aio orchestration and its tests are removed. Also write the exists rule of copy_pairs() in the shape of fsspec's copy(), and call _moves() through self. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 114 +++------------------ pyathena/filesystem/s3_path_pairing.py | 17 ++- tests/pyathena/filesystem/test_s3_async.py | 64 +----------- 3 files changed, 22 insertions(+), 173 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index c19702a7..59256218 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -13,14 +13,13 @@ import logging import mimetypes import os -from collections.abc import Callable from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast from fsspec.asyn import AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK from fsspec.core import get_compression -from fsspec.implementations.local import LocalFileSystem, make_path_posix, trailing_sep +from fsspec.implementations.local import LocalFileSystem, make_path_posix from fsspec.utils import check_contained from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem @@ -325,35 +324,10 @@ async def _rm( ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = await self._delete_paths(path, recursive=recursive, maxdepth=maxdepth) - await self._delete_objects(paths, **kwargs) - - async def _delete_paths( - self, path: str | list[str], recursive: bool = False, maxdepth: int | None = None - ) -> list[str]: - """Expand the paths that ``_rm()`` deletes, as ``S3FileSystem._delete_paths`` does. - - Args: - path: S3 path or list of paths. - recursive: Whether to include all objects below the paths. - maxdepth: Maximum depth to expand when ``recursive`` is True. - - Returns: - The paths with a version ID as given, followed by the expansion - of the other paths by ``_expand_path``. - - Raises: - ValueError: If a path is a bucket. - """ - versioned_paths, unversioned_paths = await asyncio.to_thread( - S3PathPairing.delete_paths, path - ) - if not unversioned_paths: - # _expand_path raises FileNotFoundError for no paths. - return versioned_paths - return versioned_paths + await self._expand_path( - unversioned_paths, recursive=recursive, maxdepth=maxdepth + paths = await asyncio.to_thread( + self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth ) + await self._delete_objects(paths, **kwargs) async def _delete_objects(self, paths: list[str], **kwargs) -> None: """Delete objects with DeleteObjects requests run with ``asyncio.gather``. @@ -398,7 +372,9 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N """ if path1 == path2: return - pairs = await self._move_pairs(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = await asyncio.to_thread( + self._sync_fs._move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth + ) # Every copy finishes before a failure is raised, as in fsspec's # _copy(), and nothing is deleted after a failure. results = await asyncio.gather( @@ -411,74 +387,6 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N [p1 for (p1, _), copied in zip(pairs, results, strict=True) if copied] ) - async def _move_pairs( - self, - path1: str | list[str], - path2: str | list[str], - recursive: bool = False, - maxdepth: int | None = None, - ) -> list[tuple[str, str]]: - """Pair and check the paths of a move, as ``S3FileSystem._move_pairs`` does. - - Args: - path1: Source S3 path, glob pattern, or list of paths. - path2: Destination S3 path, or list of paths when ``path1`` is a - list. - recursive: Whether to include the contents of the directories. - maxdepth: Maximum depth of the expansion. - - Returns: - The sources and destinations that are moved. - - Raises: - ValueError: If the move has conflicting paths. - """ - pairing = S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) - pairs = await self._copy_pairs(pairing) - candidates = await asyncio.to_thread(pairing.conflict_candidates, pairs) - # Looked up in order in one thread, as S3FileSystem does, so that the - # first error stops the lookups. - missing = await asyncio.to_thread( - lambda: {s for s in candidates if self._sync_fs._head_object(s) is None} - ) - return await asyncio.to_thread(pairing.move_pairs, pairs, missing=missing) - - async def _copy_pairs( - self, pairing: S3PathPairing, isdir: Callable[[str], bool] | None = None - ) -> list[tuple[str, str]]: - """Expand and pair the paths of a copy, as ``S3FileSystem._copy_pairs`` does. - - Args: - pairing: The paths of the copy. - isdir: Whether the destination is a directory, by default the - wrapped ``S3FileSystem.isdir``; ``_get()`` passes the local - filesystem's. - - Returns: - The sources and their destinations. - """ - if not pairing.expands: - return await asyncio.to_thread(pairing.copy_pairs) - sources = await self._expand_path( - pairing.path1, recursive=pairing.recursive, maxdepth=pairing.maxdepth - ) - if pairing.skips_directories: - # A path with a trailing slash is a directory without a lookup. - # The paths are looked up in one thread, mostly from the cache. - files = [p for p in sources if not trailing_sep(p)] - sources = await asyncio.to_thread( - lambda: [p for p in files if not self._sync_fs.isdir(p)] - ) - destination_is_dir = None - if sources and pairing.looks_up_destination: - # A string, as looks_up_destination checks; looked up in a thread - # as the sources are. - destination_is_dir = await asyncio.to_thread( - isdir or self._sync_fs.isdir, cast(str, pairing.path2) - ) - # The pairing of many paths takes long enough to block the event loop. - return await asyncio.to_thread(pairing.copy_pairs, sources, destination_is_dir) - def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. @@ -522,8 +430,9 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = await self._copy_pairs( - S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth) + pairs = await asyncio.to_thread( + self._sync_fs._copy_pairs, + S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth), ) if not pairs: return @@ -562,7 +471,8 @@ async def _get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) 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._copy_pairs( + pairs = await asyncio.to_thread( + self._sync_fs._copy_pairs, S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), isdir=LocalFileSystem().isdir, ) diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index e2f61433..52c36d2d 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -128,15 +128,12 @@ def copy_pairs( raise ValueError("destination_is_dir is needed to pair the paths.") source_is_str = isinstance(path1, str) glob = isinstance(path1, str) and self._is_glob(path1) + # As in fsspec's copy(); destination_is_dir is None only when the + # destination is a list or ends with a slash, or the source is a glob + # pattern or ends with a slash, which decide without it. + dest_is_dir = isinstance(path2, str) and (trailing_sep(path2) or bool(destination_is_dir)) exists = source_is_str and ( - (glob and len(sources) == 1) - or ( - not glob - and not trailing_sep(path1) - and isinstance(path2, str) - and trailing_sep(path2) - ) - or (self.looks_up_destination and bool(destination_is_dir)) + (glob and len(sources) == 1) or (not glob and dest_is_dir and not trailing_sep(path1)) ) names = [S3Path.split_version_id(p)[0] for p in sources] destinations = other_paths(names, path2, exists=exists, flatten=not source_is_str) @@ -160,7 +157,7 @@ def conflict_candidates(self, pairs: Sequence[tuple[str, str]]) -> list[str]: another source below them and a destination that conflicts, in the order of the pairs; empty if nothing needs to be looked up. """ - return S3PathPairing._moves(pairs)[2] + return self._moves(pairs)[2] def move_pairs( self, pairs: Sequence[tuple[str, str]], missing: Collection[str] | None = None @@ -187,7 +184,7 @@ def move_pairs( """ if isinstance(missing, str): raise TypeError("missing is a collection of paths, not a path.") - named, sources, candidates = S3PathPairing._moves(pairs) + named, sources, candidates = self._moves(pairs) if missing is None and candidates: raise ValueError("missing is needed to check the pairs.") # A directory without an object at its key writes no destination. diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b436d63d..19f442c5 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -31,7 +31,6 @@ S3StorageClass, ) from pyathena.filesystem.s3_path import S3Path -from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -880,13 +879,12 @@ async def test_rm_maxdepth(self): fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) - # The paths are expanded with the coroutines of fsspec's _expand_path(). - fs._find = mock.AsyncMock(return_value=["bucket/dir/a"]) - fs._exists = mock.AsyncMock(return_value=True) + sync_fs.find = mock.MagicMock(return_value=["bucket/dir/a"]) + sync_fs.exists = mock.MagicMock(return_value=True) # batch_size is part of fsspec's async _rm() signature. await fs._rm("s3://bucket/dir", recursive=True, maxdepth=1, batch_size=10) - fs._find.assert_awaited_once_with("bucket/dir", maxdepth=1, withdirs=True) + sync_fs.find.assert_called_once_with("bucket/dir", maxdepth=1, withdirs=True, detail=False) (call,) = sync_fs._call.call_args_list assert call.kwargs["Delete"]["Objects"] == [{"Key": "dir"}, {"Key": "dir/a"}] @@ -1119,62 +1117,6 @@ def upload_part_copy(**kw): == "requester" ) - @pytest.mark.parametrize( - ("path2", "lookups", "expected"), - [ - ("s3://bucket/d", ["bucket/b?versionId=v1", "s3://bucket/d"], "s3://bucket/d/b"), - ("s3://bucket/d/", ["bucket/b?versionId=v1"], "s3://bucket/d/b"), - ], - ) - @pytest.mark.asyncio - async def test_copy_pairs_destination_lookup(self, path2, lookups, expected): - # See TestS3FileSystem.test_copy_pairs_destination_lookup; the aio - # filesystem expands and looks up the paths with its own coroutines. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - # Only the destination is a directory; the paths are looked up in a - # thread with the cached isdir of the wrapped filesystem. - fs._sync_fs.isdir = mock.MagicMock(side_effect=lambda p: p.rstrip("/").endswith("/d")) - - pairs = await fs._copy_pairs(S3PathPairing("s3://bucket/b?versionId=v1", path2)) - - assert pairs == [("bucket/b?versionId=v1", expected)] - assert [c.args[0] for c in fs._sync_fs.isdir.call_args_list] == lookups - - @pytest.mark.asyncio - async def test_move_pairs_looks_up_only_conflict_candidates(self): - # See TestS3FileSystem.test_move_pairs_looks_up_only_conflict_candidates. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - fs._sync_fs._head_object = mock.MagicMock(return_value=None) - - pairs = await fs._move_pairs( - ["s3://bucket/d", "s3://bucket/d/x", "s3://bucket/e/y", "s3://bucket/f"], - ["s3://bucket/e", "s3://bucket/e", "s3://bucket/out", "s3://bucket/g"], - ) - - assert len(pairs) == 4 - fs._sync_fs._head_object.assert_called_once_with("bucket/d") - - @pytest.mark.asyncio - async def test_move_pairs_raises_the_first_lookup_error(self): - # The conflict lookups run in order, as in S3FileSystem, so the first - # error is raised and stops them. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) - - def head_object(path): - if path == "bucket/d1": - time.sleep(0.1) - raise PermissionError(path) - raise TimeoutError(path) - - fs._sync_fs._head_object = mock.MagicMock(side_effect=head_object) - - with pytest.raises(PermissionError, match="bucket/d1"): - await fs._move_pairs( - ["s3://bucket/d1", "s3://bucket/d1/x", "s3://bucket/d2", "s3://bucket/d2/x"], - ["s3://bucket/o", "s3://bucket/o", "s3://bucket/p", "s3://bucket/p"], - ) - assert fs._sync_fs._head_object.call_count == 1 - def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance # cache, so skip_instance_cache=True instances shared it. From 1a73556cdfd0d10e144d9e53e0972b005841293a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 4 Oct 2026 19:09:05 +0900 Subject: [PATCH 10/10] Keep one conflict candidate per pair, so that moves send the same HEAD requests as master Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_path_pairing.py | 28 +++++++++---------- .../filesystem/test_s3_path_pairing.py | 17 +++++++++++ 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/pyathena/filesystem/s3_path_pairing.py b/pyathena/filesystem/s3_path_pairing.py index 52c36d2d..d4233416 100644 --- a/pyathena/filesystem/s3_path_pairing.py +++ b/pyathena/filesystem/s3_path_pairing.py @@ -128,9 +128,8 @@ def copy_pairs( raise ValueError("destination_is_dir is needed to pair the paths.") source_is_str = isinstance(path1, str) glob = isinstance(path1, str) and self._is_glob(path1) - # As in fsspec's copy(); destination_is_dir is None only when the - # destination is a list or ends with a slash, or the source is a glob - # pattern or ends with a slash, which decide without it. + # As in fsspec's copy(); destination_is_dir is None only when + # looks_up_destination is false, where the other terms decide. dest_is_dir = isinstance(path2, str) and (trailing_sep(path2) or bool(destination_is_dir)) exists = source_is_str and ( (glob and len(sources) == 1) or (not glob and dest_is_dir and not trailing_sep(path1)) @@ -154,8 +153,9 @@ def conflict_candidates(self, pairs: Sequence[tuple[str, str]]) -> list[str]: Returns: The sources, in ``bucket/key`` form (their :attr:`~pyathena.filesystem.s3_path.S3Path.target`), that have - another source below them and a destination that conflicts, in the - order of the pairs; empty if nothing needs to be looked up. + another source below them and a destination that conflicts, one + per pair in the order of the pairs; empty if nothing needs to be + looked up. """ return self._moves(pairs)[2] @@ -281,14 +281,12 @@ def _moves( directories.add(parent) parent = parent.rpartition("/")[0] counts = Counter(dest for _, _, _, source, dest in named if source != dest) - candidates = list( - dict.fromkeys( - source - for _, _, versioned, source, dest in named - if source != dest - and (counts[dest] > 1 or dest in sources) - and source in directories - and not versioned - ) - ) + candidates = [ + source + for _, _, versioned, source, dest in named + if source != dest + and (counts[dest] > 1 or dest in sources) + and source in directories + and not versioned + ] return named, sources, candidates diff --git a/tests/pyathena/filesystem/test_s3_path_pairing.py b/tests/pyathena/filesystem/test_s3_path_pairing.py index ccdfb445..f57fbacf 100644 --- a/tests/pyathena/filesystem/test_s3_path_pairing.py +++ b/tests/pyathena/filesystem/test_s3_path_pairing.py @@ -102,6 +102,14 @@ def test_looks_up_destination(self, path1, path2, expected): None, [("bucket/b?versionId=v1", "s3://bucket/d/b")], ), + # A list of sources is copied into the destination by name. + ( + ["s3://bucket/b?versionId=v1"], + "s3://bucket/d", + ["bucket/b?versionId=v1"], + None, + [("bucket/b?versionId=v1", "s3://bucket/d/b")], + ), # Lists are paired as given, up to the end of the shorter one. ( ["s3://bucket/a", "s3://bucket/b", "s3://bucket/c"], @@ -174,6 +182,15 @@ def test_move_pairs_missing_is_a_collection(self): with pytest.raises(TypeError, match="collection"): _pairing(pairs).move_pairs(pairs, missing="bucket/d") + def test_conflict_candidates_one_per_pair(self): + # A source given twice is looked up twice, as each pair is checked. + pairs = [ + ("s3://bucket/d", "s3://bucket/out"), + ("s3://bucket/d", "s3://bucket/out"), + ("s3://bucket/d/x", "s3://bucket/x"), + ] + assert _pairing(pairs).conflict_candidates(pairs) == ["bucket/d", "bucket/d"] + def test_move_pairs_version_names_an_object(self): # A version is never taken for a directory. pairs = [